RootTree.h

Go to the documentation of this file.

#ifndef IOPool_Input_RootTree_h
#define IOPool_Input_RootTree_h

/*----------------------------------------------------------------------

RootTree.h // used by ROOT input sources

----------------------------------------------------------------------*/

#include "DataFormats/Provenance/interface/BranchDescription.h"
#include "DataFormats/Provenance/interface/IndexIntoFile.h"
#include "DataFormats/Provenance/interface/BranchKey.h"
#include "DataFormats/Provenance/interface/ProvenanceFwd.h"
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Utilities/interface/InputType.h"

#include "Rtypes.h"
#include "TBranch.h"

#include <memory>
#include <string>
#include <vector>
#include <unordered_set>
#include <unordered_map>

class TBranch;
class TClass;
class TTree;
class TTreeCache;

namespace edm {
  class BranchKey;
  class RootDelayedReader;
  class InputFile;
  class RootTree;

  class StreamContext;
  class ModuleCallingContext;
  
  namespace signalslot {
    template <typename T> class Signal;
  }


  namespace roottree {
    unsigned int const defaultCacheSize = 20U * 1024 * 1024;
    unsigned int const defaultNonEventCacheSize = 1U * 1024 * 1024;
    unsigned int const defaultLearningEntries = 20U;
    unsigned int const defaultNonEventLearningEntries = 1U;
    typedef IndexIntoFile::EntryNumber_t EntryNumber;
    struct BranchInfo {
      BranchInfo(BranchDescription const& prod) :
        branchDescription_(prod),
        productBranch_(nullptr),
        provenanceBranch_(nullptr),
        classCache_(nullptr),
        offsetToWrapperBase_(0) {}
      BranchDescription const branchDescription_;
      TBranch* productBranch_;
      TBranch* provenanceBranch_; // For backward compatibility
      mutable TClass* classCache_;
      mutable Int_t offsetToWrapperBase_;
    };
    
    class BranchMap {
      enum {
        kKeys,
        kInfos,
      };
    public:
      void reserve(size_t iSize) {
        map_.reserve(iSize);
      }
      void insert(edm::BranchID const& iKey, BranchInfo const& iInfo) {
        map_.emplace(iKey.id(),iInfo);
      }
      BranchInfo const* find(BranchID const& iKey) const {
        auto itFound = map_.find(iKey.id());
        if(itFound == map_.end()) {return nullptr;}
        return &itFound->second;
      }
      BranchInfo* find(BranchID const& iKey) {
        auto itFound = map_.find(iKey.id());
        if(itFound == map_.end()) {return nullptr;}
        return &itFound->second;
      }

    private:
      std::unordered_map<unsigned int,BranchInfo> map_;
    };

    Int_t getEntry(TBranch* branch, EntryNumber entryNumber);
    Int_t getEntry(TTree* tree, EntryNumber entryNumber);
    std::unique_ptr<TTreeCache> trainCache(TTree* tree, InputFile& file, unsigned int cacheSize, char const* branchNames);
  }

  class RootTree {
  public:
    typedef roottree::BranchMap BranchMap;
    typedef roottree::EntryNumber EntryNumber;
    RootTree(std::shared_ptr<InputFile> filePtr,
             BranchType const& branchType,
             unsigned int nIndexes,
             unsigned int maxVirtualSize,
             unsigned int cacheSize,
             unsigned int learningEntries,
             bool enablePrefetching,
             InputType inputType);
    ~RootTree();

    RootTree(RootTree const&) = delete; // Disallow copying and moving
    RootTree& operator=(RootTree const&) = delete; // Disallow copying and moving

    bool isValid() const;
    void numberOfBranchesToAdd(size_t iSize) { branches_.reserve(iSize);}
    void addBranch(BranchDescription const& prod,
                   std::string const& oldBranchName);
    void dropBranch(std::string const& oldBranchName);
    void getEntry(TBranch *branch, EntryNumber entry) const;
    void setPresence(BranchDescription& prod,
                   std::string const& oldBranchName);

    bool next() {return ++entryNumber_ < entries_;}
    bool nextWithCache();
    bool previous() {return --entryNumber_ >= 0;}
    bool current() const {return entryNumber_ < entries_ && entryNumber_ >= 0;}
    bool current(EntryNumber entry) const {return entry < entries_ && entry >= 0;}
    void rewind() {entryNumber_ = 0;}
    void close();
    bool skipEntries(unsigned int& offset);
    EntryNumber const& entryNumber() const {return entryNumber_;}
    EntryNumber const& entryNumberForIndex(unsigned int index) const;
    EntryNumber const& entries() const {return entries_;}
    void setEntryNumber(EntryNumber theEntryNumber);
    void insertEntryForIndex(unsigned int index);
    std::vector<std::string> const& branchNames() const {return branchNames_;}
    DelayedReader* rootDelayedReader() const;
    DelayedReader* resetAndGetRootDelayedReader() const;
    template <typename T>
    void fillAux(T*& pAux) {
      auxBranch_->SetAddress(&pAux);
      getEntry(auxBranch_, entryNumber_);
    }
    template <typename T>
    void fillBranchEntryMeta(TBranch* branch, T*& pbuf) {
      if (metaTree_ != nullptr) {
        // Metadata was in separate tree.  Not cached.
        branch->SetAddress(&pbuf);
        roottree::getEntry(branch, entryNumber_);
      } else {
        fillBranchEntry<T>(branch, pbuf);
      }
    }

    template <typename T>
    void fillBranchEntry(TBranch* branch, T*& pbuf) {
      branch->SetAddress(&pbuf);
      getEntry(branch, entryNumber_);
    }

    template <typename T>
    void fillBranchEntryMeta(TBranch* branch, EntryNumber entryNumber, T*& pbuf) {
      if (metaTree_ != nullptr) {
        // Metadata was in separate tree.  Not cached.
        branch->SetAddress(&pbuf);
        roottree::getEntry(branch, entryNumber);
      } else {
        fillBranchEntry<T>(branch, entryNumber, pbuf);
      }
    }
    
    template <typename T>
    void fillBranchEntry(TBranch* branch, EntryNumber entryNumber, T*& pbuf) {
      branch->SetAddress(&pbuf);
      getEntry(branch, entryNumber);
    }
    
    TTree const* tree() const {return tree_;}
    TTree* tree() {return tree_;}
    TTree const* metaTree() const {return metaTree_;}
    BranchMap const& branches() const;

    //For backwards compatibility
    TBranch* branchEntryInfoBranch() const {return branchEntryInfoBranch_;}

    inline TTreeCache* checkTriggerCache(TBranch* branch, EntryNumber entryNumber) const;
    TTreeCache* checkTriggerCacheImpl(TBranch* branch, EntryNumber entryNumber) const;
    inline TTreeCache* selectCache(TBranch* branch, EntryNumber entryNumber) const;
    void trainCache(char const* branchNames);
    void resetTraining() {trainNow_ = true;}

    BranchType branchType() const {return branchType_;}
    
    void setSignals(signalslot::Signal<void(StreamContext const&, ModuleCallingContext const&)> const* preEventReadSource,
                    signalslot::Signal<void(StreamContext const&, ModuleCallingContext const&)> const* postEventReadSource);

  private:
    void setCacheSize(unsigned int cacheSize);
    void setTreeMaxVirtualSize(int treeMaxVirtualSize);
    void startTraining();
    void stopTraining();

    std::shared_ptr<InputFile> filePtr_;
// We use bare pointers for pointers to some ROOT entities.
// Root owns them and uses bare pointers internally.
// Therefore,using smart pointers here will do no good.
    TTree* tree_;
    TTree* metaTree_;
    BranchType branchType_;
    TBranch* auxBranch_;
// We use a smart pointer to own the TTreeCache.
// Unfortunately, ROOT owns it when attached to a TFile, but not after it is detached.
// So, we make sure to it is detached before closing the TFile so there is no double delete.
    std::shared_ptr<TTreeCache> treeCache_;
    std::shared_ptr<TTreeCache> rawTreeCache_;
    mutable std::shared_ptr<TTreeCache> triggerTreeCache_;
    mutable std::shared_ptr<TTreeCache> rawTriggerTreeCache_;
    mutable std::unordered_set<TBranch*> trainedSet_;
    mutable std::unordered_set<TBranch*> triggerSet_;
    EntryNumber entries_;
    EntryNumber entryNumber_;
    std::unique_ptr<std::vector<EntryNumber> > entryNumberForIndex_;
    std::vector<std::string> branchNames_;
    BranchMap branches_;
    bool trainNow_;
    EntryNumber switchOverEntry_;
    mutable EntryNumber rawTriggerSwitchOverEntry_;
    mutable bool performedSwitchOver_;
    unsigned int learningEntries_;
    unsigned int cacheSize_;
    unsigned long treeAutoFlush_;
// Enable asynchronous I/O in ROOT (done in a separate thread).  Only takes
// effect on the primary treeCache_; all other caches have this explicitly disabled.
    bool enablePrefetching_;
    bool enableTriggerCache_;
    std::unique_ptr<RootDelayedReader> rootDelayedReader_;

    TBranch* branchEntryInfoBranch_; //backwards compatibility
    // below for backward compatibility
    TTree* infoTree_; // backward compatibility
    TBranch* statusBranch_; // backward compatibility
  };
}
#endif