Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-12-16 10:31:41

0001 // @(#)root/tree:$Id$
0002 // Author: Axel Naumann, 2010-08-02
0003 // Author: Vincenzo Eduardo Padulano CERN 09/2024
0004 
0005 /*************************************************************************
0006  * Copyright (C) 1995-2024, Rene Brun and Fons Rademakers.               *
0007  * All rights reserved.                                                  *
0008  *                                                                       *
0009  * For the licensing terms see $ROOTSYS/LICENSE.                         *
0010  * For the list of contributors see $ROOTSYS/README/CREDITS.             *
0011  *************************************************************************/
0012 
0013 #ifndef ROOT_TTreeReader
0014 #define ROOT_TTreeReader
0015 
0016 ////////////////////////////////////////////////////////////////////////////
0017 //                                                                        //
0018 // TTreeReader                                                            //
0019 //                                                                        //
0020 // A simple interface for reading trees or chains.                        //
0021 //                                                                        //
0022 //                                                                        //
0023 ////////////////////////////////////////////////////////////////////////////
0024 
0025 #include "TTree.h"
0026 #include "TTreeReaderUtils.h"
0027 #include "TNotifyLink.h"
0028 
0029 #include <deque>
0030 #include <iterator>
0031 #include <unordered_map>
0032 #include <set>
0033 #include <string>
0034 
0035 class TDictionary;
0036 class TDirectory;
0037 class TFileCollection;
0038 
0039 namespace ROOT {
0040 namespace Internal {
0041 class TBranchProxyDirector;
0042 class TFriendProxy;
0043 } // namespace Internal
0044 } // namespace ROOT
0045 
0046 class TTreeReader : public TObject {
0047 public:
0048    ///\class TTreeReader::Iterator_t
0049    /// Iterate through the entries of a TTree.
0050    ///
0051    /// This iterator drives the associated TTreeReader; its
0052    /// dereferencing (and actually even the iteration) will
0053    /// set the entry number represented by this iterator.
0054    /// It does not really represent a data element; it simply
0055    /// returns the entry number (or -1 once the end of the tree
0056    /// is reached).
0057    class Iterator_t {
0058    private:
0059       Long64_t fEntry;      ///< Entry number of the tree referenced by this iterator; -1 is invalid.
0060       TTreeReader *fReader; ///< The reader we select the entries on.
0061 
0062       /// Whether the iterator points to a valid entry.
0063       bool IsValid() const { return fEntry >= 0; }
0064 
0065    public:
0066       using iterator_category = std::input_iterator_tag;
0067       using value_type = const Long64_t;
0068       using difference_type = Long64_t;
0069       using pointer = const Long64_t *;
0070       using const_pointer = const Long64_t *;
0071       using reference = const Long64_t &;
0072 
0073       /// Default-initialize the iterator as "past the end".
0074       Iterator_t() : fEntry(-1), fReader(nullptr) {}
0075 
0076       /// Initialize the iterator with the reader it steers and a
0077       /// tree entry number; -1 is invalid.
0078       Iterator_t(TTreeReader &reader, Long64_t entry) : fEntry(entry), fReader(&reader) {}
0079 
0080       /// Compare two iterators for equality.
0081       bool operator==(const Iterator_t &lhs) const
0082       {
0083          // From C++14: value initialized (past-end) it compare equal.
0084          if (!IsValid() && !lhs.IsValid())
0085             return true;
0086          // The iterators refer to different readers
0087          if (fReader != lhs.fReader)
0088             return false;
0089          // #16249: range based loop and the tree has zero entries
0090          // as well as analogous cases.
0091          // Getting the number of events can have a cost, for example in
0092          // case of chains of remote files accessible with high latency.
0093          // However, it is reasonable to assume that if iterators are
0094          // being compared is because an iteration is taking place,
0095          // therefore such cost has to be paid anyway, it's just
0096          // anticipated.
0097          if (fReader->GetTree()->GetEntriesFast() == 0 && fEntry == 0 && !lhs.IsValid()) {
0098             return true;
0099          }
0100          if (lhs.fReader->GetTree()->GetEntriesFast() == 0 && lhs.fEntry == 0 && !IsValid()) {
0101             return true;
0102          }
0103          return fEntry == lhs.fEntry;
0104       }
0105 
0106       /// Compare two iterators for inequality.
0107       bool operator!=(const Iterator_t &lhs) const { return !(*this == lhs); }
0108 
0109       /// Increment the iterator (postfix i++).
0110       Iterator_t operator++(int)
0111       {
0112          Iterator_t ret = *this;
0113          this->operator++();
0114          return ret;
0115       }
0116 
0117       /// Increment the iterator (prefix ++i).
0118       Iterator_t &operator++()
0119       {
0120          if (IsValid()) {
0121             ++fEntry;
0122             // Force validity check of new fEntry.
0123             this->operator*();
0124             // Don't set the old entry: op* will if needed, and
0125             // in most cases it just adds a lot of spinning back
0126             // and forth: in most cases the sequence is ++i; *i.
0127          }
0128          return *this;
0129       }
0130 
0131       /// Set the entry number in the reader and return it.
0132       const Long64_t &operator*()
0133       {
0134          if (IsValid()) {
0135             // If we cannot access that entry, mark the iterator invalid.
0136             if (fReader->SetEntry(fEntry) != kEntryValid) {
0137                fEntry = -1;
0138             }
0139          }
0140          // There really is no data in this iterator; return the number.
0141          return fEntry;
0142       }
0143 
0144       const Long64_t &operator*() const { return **const_cast<Iterator_t *>(this); }
0145    };
0146 
0147    typedef Iterator_t iterator;
0148 
0149    enum EEntryStatus {
0150       kEntryValid = 0,                 ///< data read okay
0151       kEntryNotLoaded,                 ///< no entry has been loaded yet
0152       kEntryNoTree,                    ///< the tree does not exist
0153       kEntryNotFound,                  ///< the tree entry number does not exist
0154       kEntryChainSetupError,           ///< problem in accessing a chain element, e.g. file without the tree
0155       kEntryChainFileError,            ///< problem in opening a chain's file
0156       kEntryDictionaryError,           ///< problem reading dictionary info from tree
0157       kEntryBeyondEnd,                 ///< last entry loop has reached its end
0158       kEntryBadReader,                 ///< One of the readers was not successfully initialized.
0159       kIndexedFriendNoMatch,           ///< A friend with TTreeIndex doesn't have an entry for this index
0160       kMissingBranchWhenSwitchingTree, ///< A branch was not found when switching to the next TTree in the chain
0161       kEntryUnknownError               ///< LoadTree return less than -6, likely a 'newer' error code.
0162    };
0163 
0164    enum ELoadTreeStatus {
0165       kNoTree = 0,           ///< default state, no TTree is connected (formerly 'Zombie' state)
0166       kLoadTreeNone,         ///< Notify has not been called yet.
0167       kInternalLoadTree,     ///< Notify/LoadTree was last called from SetEntryBase
0168       kExternalLoadTree,     ///< User code called LoadTree directly.
0169       kMissingBranchFromTree ///< Missing expected branch when loading new tree
0170    };
0171 
0172    static constexpr const char *const fgEntryStatusText[kEntryUnknownError + 1] = {
0173       "valid entry",
0174       "the tree does not exist",
0175       "the tree entry number does not exist",
0176       "cannot access chain element",
0177       "problem in opening a chain's file",
0178       "problem reading dictionary info from tree",
0179       "last entry loop has reached its end",
0180       "one of the readers was not successfully initialized",
0181       "A friend with TTreeIndex doesn't have an entry for this index",
0182       "A branch was not found when switching to the next TTree in the chain",
0183       "LoadTree return less than -6, likely a 'newer' error code"};
0184 
0185    TTreeReader();
0186 
0187    TTreeReader(TTree *tree, TEntryList *entryList = nullptr, bool warnAboutLongerFriends = true,
0188                const std::set<std::string> &suppressErrorsForMissingBranches = {});
0189    TTreeReader(const char *keyname, TDirectory *dir, TEntryList *entryList = nullptr);
0190    TTreeReader(const char *keyname, TEntryList *entryList = nullptr) : TTreeReader(keyname, nullptr, entryList) {}
0191 
0192    ~TTreeReader() override;
0193 
0194    void SetTree(TTree *tree, TEntryList *entryList = nullptr);
0195    void SetTree(const char *keyname, TEntryList *entryList = nullptr) { SetTree(keyname, nullptr, entryList); }
0196    void SetTree(const char *keyname, TDirectory *dir, TEntryList *entryList = nullptr);
0197 
0198    bool IsChain() const { return TestBit(kBitIsChain); }
0199 
0200    bool IsInvalid() const { return fLoadTreeStatus == kNoTree; }
0201 
0202    TTree *GetTree() const { return fTree; }
0203    TEntryList *GetEntryList() const { return fEntryList; }
0204 
0205    ///\{ \name Entry setters
0206 
0207    /// Move to the next entry (or index of the TEntryList if that is set).
0208    ///
0209    /// \return false if the previous entry was already the last entry. This allows
0210    ///   the function to be used in `while (reader.Next()) { ... }`
0211    bool Next() { return SetEntry(GetCurrentEntry() + 1) == kEntryValid; }
0212 
0213    /// Set the next entry (or index of the TEntryList if that is set).
0214    ///
0215    /// \param entry If not TEntryList is set, the entry is a global entry (i.e.
0216    /// not the entry number local to the chain's current tree).
0217    /// \returns the `entry`'s read status, i.e. whether the entry is available.
0218    EEntryStatus SetEntry(Long64_t entry) { return SetEntryBase(entry, false); }
0219 
0220    /// Set the next local tree entry. If a TEntryList is set, this function is
0221    /// equivalent to `SetEntry()`.
0222    ///
0223    /// \param entry Entry number of the TChain's current TTree. This is the
0224    /// entry number passed for instance by `TSelector::Process(entry)`, i.e.
0225    /// within `TSelector::Process()` always use `SetLocalEntry()` and not
0226    /// `SetEntry()`!
0227    /// \return the `entry`'s read status, i.e. whether the entry is available.
0228    EEntryStatus SetLocalEntry(Long64_t entry) { return SetEntryBase(entry, true); }
0229 
0230    EEntryStatus SetEntriesRange(Long64_t beginEntry, Long64_t endEntry);
0231 
0232    ///  Get the begin and end entry numbers
0233    ///
0234    /// \return a pair contained the begin and end entry numbers.
0235    std::pair<Long64_t, Long64_t> GetEntriesRange() const { return std::make_pair(fBeginEntry, fEndEntry); }
0236 
0237    /// Restart a Next() loop from entry 0 (of TEntryList index 0 of fEntryList is set).
0238    void Restart();
0239 
0240    ///\}
0241 
0242    EEntryStatus GetEntryStatus() const { return fEntryStatus; }
0243 
0244    Long64_t GetEntries() const;
0245    Long64_t GetEntries(bool force);
0246 
0247    /// Returns the index of the current entry being read.
0248    ///
0249    /// If `IsChain()`, the returned index corresponds to the global entry number
0250    /// (i.e. not the entry number local to the chain's current tree).
0251    /// If `fEntryList`, the returned index corresponds to an index in the
0252    /// TEntryList; to translate to the TChain's / TTree's entry number pass it
0253    /// through `reader.GetEntryList()->GetEntry(reader.GetCurrentEntry())`.
0254    Long64_t GetCurrentEntry() const { return fEntry; }
0255 
0256    bool Notify() override;
0257 
0258    /// Return an iterator to the 0th TTree entry.
0259    Iterator_t begin() { return Iterator_t(*this, 0); }
0260    /// Return an iterator beyond the last TTree entry.
0261    Iterator_t end() { return Iterator_t(*this, -1); }
0262 
0263    void SetSuppressErrorsForMissingBranches(const std::set<std::string> &suppressErrorsForMissingBranches)
0264    {
0265       fSuppressErrorsForMissingBranches = suppressErrorsForMissingBranches;
0266    }
0267 
0268 protected:
0269    using NamedProxies_t = std::unordered_map<std::string, std::unique_ptr<ROOT::Internal::TNamedBranchProxy>>;
0270    void Initialize();
0271    ROOT::Internal::TNamedBranchProxy *FindProxy(const char *branchname) const
0272    {
0273       const auto proxyIt = fProxies.find(branchname);
0274       return fProxies.end() != proxyIt ? proxyIt->second.get() : nullptr;
0275    }
0276 
0277    void AddProxy(std::unique_ptr<ROOT::Internal::TNamedBranchProxy> p)
0278    {
0279       auto bpName = p->GetName();
0280 #ifndef NDEBUG
0281       if (fProxies.end() != fProxies.find(bpName)) {
0282          std::string err = "A proxy with key " + std::string(bpName) + " was already stored!";
0283          throw std::runtime_error(err);
0284       }
0285 #endif
0286 
0287       fProxies[bpName] = std::move(p);
0288    }
0289 
0290    ROOT::Internal::TFriendProxy &AddFriendProxy(std::size_t friendIdx);
0291 
0292    bool RegisterValueReader(ROOT::Internal::TTreeReaderValueBase *reader);
0293    void DeregisterValueReader(ROOT::Internal::TTreeReaderValueBase *reader);
0294 
0295    EEntryStatus SetEntryBase(Long64_t entry, bool local);
0296 
0297    bool SetProxies();
0298 
0299 private:
0300    std::string GetProxyKey(const char *branchname)
0301    {
0302       std::string key(branchname);
0303       // key += reinterpret_cast<std::uintptr_t>(fTree);
0304       return key;
0305    }
0306 
0307    enum EStatusBits {
0308       kBitIsChain = BIT(14), ///< our tree is a chain
0309       kBitHaveWarnedAboutEntryListAttachedToTTree =
0310          BIT(15),                                ///< the tree had a TEntryList and we have warned about that
0311       kBitSetEntryBaseCallingLoadTree = BIT(16), ///< SetEntryBase is in the process of calling TChain/TTree::%LoadTree.
0312       kBitIsExternalTree = BIT(17)               ///< we do not own the tree
0313    };
0314 
0315    TTree *fTree = nullptr;                      ///< tree that's read
0316    TEntryList *fEntryList = nullptr;            ///< entry list to be used
0317    EEntryStatus fEntryStatus = kEntryNotLoaded; ///< status of most recent read request
0318    ELoadTreeStatus fLoadTreeStatus = kNoTree;   ///< Indicator on how LoadTree was called 'last' time.
0319    /// TTree and TChain will notify this object upon LoadTree, leading to a call to TTreeReader::Notify().
0320    TNotifyLink<TTreeReader> fNotify;
0321    std::unique_ptr<ROOT::Internal::TBranchProxyDirector> fDirector{nullptr}; ///< proxying director
0322    /// Proxies to friend trees, created in TTreeReader[Value,Array]::CreateProxy
0323    std::vector<std::unique_ptr<ROOT::Internal::TFriendProxy>> fFriendProxies;
0324    std::deque<ROOT::Internal::TTreeReaderValueBase *> fValues; ///< readers that use our director
0325    NamedProxies_t fProxies;                                    ///< attached ROOT::TNamedBranchProxies; owned
0326 
0327    Long64_t fEntry = -1; ///< Current (non-local) entry of fTree or of fEntryList if set.
0328 
0329    /// The end of the entry loop. When set (i.e. >= 0), it provides a way
0330    /// to stop looping over the TTree when we reach a certain entry: Next()
0331    /// returns false when GetCurrentEntry() reaches fEndEntry.
0332    Long64_t fEndEntry = -1LL;
0333    Long64_t fBeginEntry = 0LL;                ///< This allows us to propagate the range to the TTreeCache
0334    bool fProxiesSet = false;                  ///< True if the proxies have been set, false otherwise
0335    bool fSetEntryBaseCallingLoadTree = false; ///< True if during the LoadTree execution triggered by SetEntryBase.
0336 
0337    // Flag to activate or deactivate warnings in case the friend trees have
0338    // more entries than the main one. In some cases we may want to deactivate
0339    // this behaviour, notably in multithreaded runs where we have to partition
0340    // the main tree but keep the entire friend trees in every thread to ensure
0341    // alignment.
0342    bool fWarnAboutLongerFriends{true};
0343    void WarnIfFriendsHaveMoreEntries();
0344 
0345    // List of branches for which we want to suppress the printed error about
0346    // missing branch when switching to a new tree
0347    std::set<std::string> fSuppressErrorsForMissingBranches{};
0348    std::set<std::string> fMissingProxies{};
0349 
0350    friend class ROOT::Internal::TTreeReaderValueBase;
0351    friend class ROOT::Internal::TTreeReaderArrayBase;
0352 
0353    ClassDefOverride(TTreeReader, 0); // A simple interface to read trees
0354 };
0355 
0356 #endif // defined TTreeReader