Merge pull request #23664 from Dr15Jones/useBranchIDinDelayedReader

Use BranchID in DelayedReader

FWCore/Framework/interface/DelayedReader.h

@@ -15,7 +15,7 @@ uses input sources to retrieve EDProducts from external storage.
 
 namespace edm {
 
-  class BranchKey;
+  class BranchID;
   class EDProductGetter;
   class ModuleCallingContext;
   class SharedResourcesAcquirer;
@@ -28,7 +28,7 @@ namespace edm {
   class DelayedReader {
   public:
     virtual ~DelayedReader();
-    std::unique_ptr<WrapperBase> getProduct(BranchKey const& k,
+    std::unique_ptr<WrapperBase> getProduct(BranchID const& k,
                                             EDProductGetter const* ep,
                                             ModuleCallingContext const* mcc = nullptr);
 
@@ -45,7 +45,7 @@ namespace edm {
 
 
   private:
-    virtual std::unique_ptr<WrapperBase> getProduct_(BranchKey const& k, EDProductGetter const* ep) = 0;
+    virtual std::unique_ptr<WrapperBase> getProduct_(BranchID const& k, EDProductGetter const* ep) = 0;
     virtual void mergeReaders_(DelayedReader*) = 0;
     virtual void reset_() = 0;
     virtual std::pair<SharedResourcesAcquirer*, std::recursive_mutex*> sharedResources_() const;

FWCore/Framework/src/DelayedReader.cc

@@ -16,7 +16,7 @@ namespace edm {
   DelayedReader::~DelayedReader() {}
 
   std::unique_ptr<WrapperBase>
-  DelayedReader::getProduct(BranchKey const& k,
+  DelayedReader::getProduct(BranchID const& k,
                             EDProductGetter const* ep,
                             ModuleCallingContext const* mcc) {
 

FWCore/Framework/src/ProductResolvers.cc

@@ -140,7 +140,7 @@ namespace edm {
         }
         if ( not productResolved()) {
           //another thread could have beaten us here
-          putProduct( reader->getProduct(BranchKey(branchDescription()), &principal, mcc));
+          putProduct( reader->getProduct(branchDescription().branchID(), &principal, mcc));
         }
       }
     });
@@ -158,8 +158,7 @@ namespace edm {
 
       //Can't use resolveProductImpl since it first checks to see
       // if the product was already retrieved and then returns if it is
-      BranchKey const bk = BranchKey(branchDescription());
-      std::unique_ptr<WrapperBase> edp(reader->getProduct(bk, &principal));
+      std::unique_ptr<WrapperBase> edp(reader->getProduct(branchDescription().branchID(), &principal));
 
       if(edp.get() != nullptr) {
         putOrMergeProduct(std::move(edp));
@@ -194,7 +193,7 @@ namespace edm {
               }
               if ( not productResolved()) {
                 //another thread could have finished this while we were waiting
-                putProduct( reader->getProduct(BranchKey(branchDescription()), &principal, mcc));
+                putProduct( reader->getProduct(branchDescription().branchID(), &principal, mcc));
               }
             }
           });

FWCore/TFWLiteSelector/src/TFWLiteSelectorBasic.cc

@@ -57,18 +57,19 @@
 #include <memory>
 #include <string>
 #include <vector>
+#include <unordered_map>
 
 namespace edm {
   namespace root {
     class FWLiteDelayedReader : public DelayedReader {
      public:
-      FWLiteDelayedReader() : entry_(-1), eventTree_(nullptr), reg_() {}
+      FWLiteDelayedReader() : entry_(-1), eventTree_(nullptr) {}
       void setEntry(Long64_t iEntry) { entry_ = iEntry; }
       void setTree(TTree* iTree) {eventTree_ = iTree;}
-      void set(std::shared_ptr<ProductRegistry const> iReg) { reg_ = iReg;}
+      void set(std::shared_ptr<std::unordered_map<unsigned int, BranchDescription const*>> iMap) { bidToDesc_ = std::move(iMap);}
      private:
-      std::unique_ptr<WrapperBase> getTheProduct(BranchKey const& k) const;
-      std::unique_ptr<WrapperBase> getProduct_(BranchKey const& k, EDProductGetter const* ep) override;
+      std::unique_ptr<WrapperBase> getTheProduct(BranchID const& k) const;
+      std::unique_ptr<WrapperBase> getProduct_(BranchID const& k, EDProductGetter const* ep) override;
       virtual std::unique_ptr<EventEntryDescription> getProvenance_(BranchKey const&) const {
         return std::unique_ptr<EventEntryDescription>();
       }
@@ -84,21 +85,21 @@ namespace edm {
 
       Long64_t entry_;
       TTree* eventTree_;
-      std::shared_ptr<ProductRegistry const>(reg_);
+      std::shared_ptr<std::unordered_map<unsigned int, BranchDescription const*>> bidToDesc_;
     };
 
     std::unique_ptr<WrapperBase>
-    FWLiteDelayedReader::getProduct_(BranchKey const& k, EDProductGetter const* /*ep*/) {
+    FWLiteDelayedReader::getProduct_(BranchID const& k, EDProductGetter const* /*ep*/) {
       return getTheProduct(k);
     }
 
     std::unique_ptr<WrapperBase>
-    FWLiteDelayedReader::getTheProduct(BranchKey const& k) const {
-      ProductRegistry::ProductList::const_iterator itFind= reg_->productList().find(k);
-      if(itFind == reg_->productList().end()) {
+    FWLiteDelayedReader::getTheProduct(BranchID const& k) const {
+      auto itFind= bidToDesc_->find(k.id());
+      if(itFind == bidToDesc_->end()) {
         throw Exception(errors::ProductNotFound) << "could not find entry for product " << k;
       }
-      BranchDescription const& bDesc = itFind->second;
+      BranchDescription const& bDesc = *(itFind->second);
 
       TBranch* branch= eventTree_->GetBranch(bDesc.branchName().c_str());
       if(nullptr == branch) {
@@ -140,6 +141,7 @@ namespace edm {
       TFWLiteSelectorMembers() :
       tree_(nullptr),
       reg_(new ProductRegistry()),
+      bidToDesc_(std::make_shared<std::unordered_map<unsigned int, BranchDescription const*>>()),
       phreg_(new ProcessHistoryRegistry()),
       branchIDListHelper_(new BranchIDListHelper()),
       // Note that thinned collections are not supported yet, the next
@@ -154,7 +156,7 @@ namespace edm {
       prov_(),
       pointerToBranchBuffer_(),
       provRetriever_(new edm::ProductProvenanceRetriever(0)) {
-        reader_->set(reg());
+        reader_->set(get_underlying_safe(bidToDesc_));
       }
       void setTree(TTree* iTree) {
         tree_ = iTree;
@@ -172,7 +174,7 @@ namespace edm {
 
       edm::propagate_const<TTree*> tree_;
       edm::propagate_const<std::shared_ptr<ProductRegistry>> reg_;
-      edm::propagate_const<std::shared_ptr<ProcessHistoryRegistry>> phreg_;
+      edm::propagate_const<std::shared_ptr<std::unordered_map<unsigned int, BranchDescription const*>>> bidToDesc_;      edm::propagate_const<std::shared_ptr<ProcessHistoryRegistry>> phreg_;
       edm::propagate_const<std::shared_ptr<BranchIDListHelper>> branchIDListHelper_;
       edm::propagate_const<std::shared_ptr<ThinnedAssociationsHelper>> thinnedAssociationsHelper_;
       ProcessHistory processNames_;
@@ -498,6 +500,10 @@ TFWLiteSelectorBasic::setupNewFile(TFile& iFile) {
   }
   m_->branchIDListHelper_->updateFromInput(*branchIDListsPtr);
   m_->reg_->setFrozen();
+  m_->bidToDesc_->clear();
+  for(auto const&p: m_->reg_->productList()) {
+    m_->bidToDesc_->emplace(p.second.branchID().id(),&p.second);
+  }
   m_->ep_ = std::make_shared<edm::EventPrincipal>(m_->reg(), m_->branchIDListHelper(), m_->thinnedAssociationsHelper(), m_->pc_, nullptr);
   everythingOK_ = true;
 }

IOPool/Input/src/RootDelayedReader.cc

@@ -46,20 +46,19 @@ namespace edm {
   }
 
   std::unique_ptr<WrapperBase>
-  RootDelayedReader::getProduct_(BranchKey const& k, EDProductGetter const* ep) {
+  RootDelayedReader::getProduct_(BranchID const& k, EDProductGetter const* ep) {
     if (lastException_) {
       std::rethrow_exception(lastException_);
     }
-    iterator iter = branchIter(k);
-    if (!found(iter)) {
+    auto branchInfo = getBranchInfo(k);
+    if (not branchInfo) {
       if (nextReader_) {
         return nextReader_->getProduct(k, ep);
       } else {
         return std::unique_ptr<WrapperBase>();
       }
     }
-    roottree::BranchInfo const& branchInfo = getBranchInfo(iter);
-    TBranch* br = branchInfo.productBranch_;
+    TBranch* br = branchInfo->productBranch_;
     if (br == nullptr) {
       if (nextReader_) {
         return nextReader_->getProduct(k, ep);
@@ -72,14 +71,14 @@ namespace edm {
     //make code exception safe
     std::shared_ptr<void> refCoreStreamerGuard(nullptr,[](void*){    setRefCoreStreamer(false);
       ;});
-    TClass* cp = branchInfo.classCache_;
+    TClass* cp = branchInfo->classCache_;
     if(nullptr == cp) {
-      branchInfo.classCache_ = TClass::GetClass(branchInfo.branchDescription_.wrappedName().c_str());
-      cp = branchInfo.classCache_;
-      branchInfo.offsetToWrapperBase_ = cp->GetBaseClassOffset(wrapperBaseTClass_);
+      branchInfo->classCache_ = TClass::GetClass(branchInfo->branchDescription_.wrappedName().c_str());
+      cp = branchInfo->classCache_;
+      branchInfo->offsetToWrapperBase_ = cp->GetBaseClassOffset(wrapperBaseTClass_);
     }
     void* p = cp->New();
-    std::unique_ptr<WrapperBase> edp = getWrapperBasePtr(p, branchInfo.offsetToWrapperBase_); 
+    std::unique_ptr<WrapperBase> edp = getWrapperBasePtr(p, branchInfo->offsetToWrapperBase_);
     br->SetAddress(&p);
     try{
       //Run and Lumi only have 1 entry number, which is index 0

IOPool/Input/src/RootDelayedReader.h

@@ -7,7 +7,7 @@ RootDelayedReader.h // used by ROOT input sources
 
 ----------------------------------------------------------------------*/
 
-#include "DataFormats/Provenance/interface/BranchKey.h"
+#include "DataFormats/Provenance/interface/BranchID.h"
 #include "FWCore/Framework/interface/DelayedReader.h"
 #include "FWCore/Utilities/interface/InputType.h"
 #include "FWCore/Utilities/interface/propagate_const.h"
@@ -33,7 +33,6 @@ namespace edm {
   public:
     typedef roottree::BranchInfo BranchInfo;
     typedef roottree::BranchMap BranchMap;
-    typedef roottree::BranchMap::const_iterator iterator;
     typedef roottree::EntryNumber EntryNumber;
     RootDelayedReader(
       RootTree const& tree,
@@ -59,15 +58,13 @@ namespace edm {
     }
 
   private:
-    std::unique_ptr<WrapperBase> getProduct_(BranchKey const& k, EDProductGetter const* ep) override;
+    std::unique_ptr<WrapperBase> getProduct_(BranchID const& k, EDProductGetter const* ep) override;
     void mergeReaders_(DelayedReader* other) override {nextReader_ = other;}
     void reset_() override {nextReader_ = nullptr;}
     std::pair<SharedResourcesAcquirer*, std::recursive_mutex*> sharedResources_() const override;
 
     BranchMap const& branches() const {return tree_.branches();}
-    iterator branchIter(BranchKey const& k) const {return branches().find(k);}
-    bool found(iterator const& iter) const {return iter != branches().end();}
-    BranchInfo const& getBranchInfo(iterator const& iter) const {return iter->second; }
+    BranchInfo const* getBranchInfo(BranchID const& k) const {return branches().find(k);}
     // NOTE: filePtr_ appears to be unused, but is needed to prevent
     // the file containing the branch from being reclaimed.
     RootTree const& tree_;

IOPool/Input/src/RootFile.cc

@@ -492,9 +492,23 @@ namespace edm {
 
     // Set up information from the product registry.
     ProductRegistry::ProductList const& prodList = productRegistry()->productList();
+
+    {
+      std::array<size_t,NumBranchTypes> nBranches;
+      nBranches.fill(0);
+      for(auto const& product : prodList) {
+        ++nBranches[product.second.branchType()];
+      }
+
+      int i = 0;
+      for(auto t: treePointers_) {
+        t->numberOfBranchesToAdd(nBranches[i]);
+        ++i;
+      }
+    }
     for(auto const& product : prodList) {
       BranchDescription const& prod = product.second;
-      treePointers_[prod.branchType()]->addBranch(product.first, prod,
+      treePointers_[prod.branchType()]->addBranch(prod,
                                                   newBranchToOldBranch(prod.branchName()));
     }
 

IOPool/Input/src/RootTree.cc

@@ -48,7 +48,7 @@ namespace edm {
     entryNumber_(-1),
     entryNumberForIndex_(new std::vector<EntryNumber>(nIndexes, IndexIntoFile::invalidEntry)),
     branchNames_(),
-    branches_(new BranchMap),
+    branches_{},
     trainNow_(false),
     switchOverEntry_(-1),
     rawTriggerSwitchOverEntry_(-1),
@@ -134,13 +134,12 @@ namespace edm {
   }
 
   void
-  RootTree::addBranch(BranchKey const& key,
-                      BranchDescription const& prod,
+  RootTree::addBranch(BranchDescription const& prod,
                       std::string const& oldBranchName) {
       assert(isValid());
       //use the translated branch name
       TBranch* branch = tree_->GetBranch(oldBranchName.c_str());
-      roottree::BranchInfo info = roottree::BranchInfo(BranchDescription(prod));
+      roottree::BranchInfo info = roottree::BranchInfo(prod);
       info.productBranch_ = nullptr;
       if (prod.present()) {
         info.productBranch_ = branch;
@@ -149,7 +148,7 @@ namespace edm {
       }
       TTree* provTree = (metaTree_ != nullptr ? metaTree_ : tree_);
       info.provenanceBranch_ = provTree->GetBranch(oldBranchName.c_str());
-      branches_->insert(std::make_pair(key, info));
+      branches_.insert(prod.branchID(), info);
   }
 
   void
@@ -176,7 +175,7 @@ namespace edm {
   }
 
   roottree::BranchMap const&
-  RootTree::branches() const {return *branches_;}
+  RootTree::branches() const {return branches_;}
 
   void
   RootTree::setCacheSize(unsigned int cacheSize) {

IOPool/Input/src/RootTree.h

@@ -9,18 +9,19 @@ 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 <map>
 #include <memory>
 #include <string>
 #include <vector>
 #include <unordered_set>
+#include <unordered_map>
 
 class TBranch;
 class TClass;
@@ -60,7 +61,34 @@ namespace edm {
       mutable TClass* classCache_;
       mutable Int_t offsetToWrapperBase_;
     };
-    typedef std::map<BranchKey const, BranchInfo> BranchMap;
+
+    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);
@@ -84,8 +112,8 @@ namespace edm {
     RootTree& operator=(RootTree const&) = delete; // Disallow copying and moving
 
     bool isValid() const;
-    void addBranch(BranchKey const& key,
-                   BranchDescription const& prod,
+    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;
@@ -193,7 +221,7 @@ namespace edm {
     EntryNumber entryNumber_;
     std::unique_ptr<std::vector<EntryNumber> > entryNumberForIndex_;
     std::vector<std::string> branchNames_;
-    std::shared_ptr<BranchMap> branches_;
+    BranchMap branches_;
     bool trainNow_;
     EntryNumber switchOverEntry_;
     mutable EntryNumber rawTriggerSwitchOverEntry_;