Fixes for TTree with complex friend hierarchies

tree/tree/src/TChain.cxx

@@ -64,6 +64,8 @@ the trees in the chain.
 #include "strlcpy.h"
 #include "snprintf.h"
 
+#include <string_view>
+#include "ROOT/StringUtils.hxx"
 
 ////////////////////////////////////////////////////////////////////////////////
 /// Default constructor.
@@ -802,31 +804,69 @@ Long64_t TChain::Draw(const char* varexp, const char* selection,
 ////////////////////////////////////////////////////////////////////////////////
 /// See TTree::GetReadEntry().
 
-TBranch* TChain::FindBranch(const char* branchname)
+TBranch *TChain::FindBranch(const char *branchname)
 {
-   if (fTree) {
-      return fTree->FindBranch(branchname);
-   }
-   LoadTree(0);
-   if (fTree) {
-      return fTree->FindBranch(branchname);
+   auto findBranchImpl = [this](const char *resolvedBranchName) -> TBranch * {
+      if (fTree) {
+         return fTree->FindBranch(resolvedBranchName);
+      }
+      LoadTree(0);
+      if (fTree) {
+         return fTree->FindBranch(resolvedBranchName);
+      }
+      return nullptr;
+   };
+
+   // This will allow the branchname to be preceded by the name of this chain.
+   // See similar code in TTree::FindBranch
+   std::string_view branchNameView{branchname};
+   std::string_view chainPrefix = GetName();
+
+   if (ROOT::StartsWith(branchNameView, chainPrefix)) {
+      branchNameView.remove_prefix(chainPrefix.length());
+      if (!branchNameView.empty() && branchNameView.front() == '.') {
+         branchNameView.remove_prefix(1);
+         // We're only removing characters from the beginning of the view so we
+         // don't need to worry about missing null-termination character
+         return findBranchImpl(branchNameView.data());
+      }
    }
-   return nullptr;
+
+   return findBranchImpl(branchname);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 /// See TTree::GetReadEntry().
 
 TLeaf* TChain::FindLeaf(const char* searchname)
 {
-   if (fTree) {
-      return fTree->FindLeaf(searchname);
-   }
-   LoadTree(0);
-   if (fTree) {
-      return fTree->FindLeaf(searchname);
+   auto findLeafImpl = [this](const char *resolvedBranchName) -> TLeaf * {
+      if (fTree) {
+         return fTree->FindLeaf(resolvedBranchName);
+      }
+      LoadTree(0);
+      if (fTree) {
+         return fTree->FindLeaf(resolvedBranchName);
+      }
+      return nullptr;
+   };
+
+   // This will allow the branchname to be preceded by the name of this chain.
+   // See similar code in TTree::FindLeaf
+   std::string_view branchNameView{searchname};
+   std::string_view chainPrefix = GetName();
+
+   if (ROOT::StartsWith(branchNameView, chainPrefix)) {
+      branchNameView.remove_prefix(chainPrefix.length());
+      if (!branchNameView.empty() && branchNameView.front() == '.') {
+         branchNameView.remove_prefix(1);
+         // We're only removing characters from the beginning of the view so we
+         // don't need to worry about missing null-termination character
+         return findLeafImpl(branchNameView.data());
+      }
    }
-   return nullptr;
+
+   return findLeafImpl(searchname);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -1311,6 +1351,16 @@ Long64_t TChain::RefreshFriendAddresses()
          br->SetAutoDelete(true);
       }
    }
+
+   // We cannot know a priori if the branch(es) of the friend TChain(s) that were just
+   // updated were supposed to be connected to one of the TChainElement of this chain
+   // or possibly to another TChainElement belonging to another chain that has befriended
+   // this chain (i.e., one of the "external friends"). Thus, we forward the notification
+   // that one or more friend trees were updated to the friends of this chain.
+   if (fExternalFriends)
+      for (auto external_fe : ROOT::Detail::TRangeStaticCast<TFriendElement>(*fExternalFriends))
+         external_fe->MarkUpdated();
+
    if (fPlayer) {
       fPlayer->UpdateFormulaLeaves();
    }

tree/tree/src/TTree.cxx

@@ -6111,7 +6111,7 @@ TTree* TTree::GetFriend(const char *friendname) const
 /// that you may not be able to take advantage of this feature.
 ///
 
-const char* TTree::GetFriendAlias(TTree* tree) const
+const char *TTree::GetFriendAlias(TTree *tree) const
 {
    if ((tree == this) || (tree == GetTree())) {
       return nullptr;
@@ -6122,30 +6122,60 @@ const char* TTree::GetFriendAlias(TTree* tree) const
    if (kGetFriendAlias & fFriendLockStatus) {
       return nullptr;
    }
-   if (!fFriends) {
+
+   // This is a TTree and it does not have any friends, we can return early
+   if (GetTree() == this && !fFriends)
       return nullptr;
-   }
-   TFriendLock lock(const_cast<TTree*>(this), kGetFriendAlias);
-   TIter nextf(fFriends);
-   TFriendElement* fe = nullptr;
-   while ((fe = (TFriendElement*) nextf())) {
-      TTree* t = fe->GetTree();
-      if (t == tree) {
-         return fe->GetName();
+
+   TFriendLock lock(const_cast<TTree *>(this), kGetFriendAlias);
+
+   auto lookForFriendNameInListOfFriends = [tree](const TList &friends) -> const char * {
+      for (auto *frEl : ROOT::Detail::TRangeStaticCast<TFriendElement>(friends)) {
+         auto *frElTree = frEl->GetTree();
+         // Simplest case: we found a friend which tree is the same as the input tree
+         if (frElTree == tree)
+            return frEl->GetName();
+         // Try again: the friend tree might be actually a TChain
+         if (frElTree && frElTree->GetTree() == tree)
+            return frEl->GetName();
       }
-      // Case of a chain:
-      if (t && t->GetTree() == tree) {
-         return fe->GetName();
+      return nullptr;
+   };
+
+   // First, look for the immediate friends of this tree
+   if (fFriends) {
+      const char *friendAlias = lookForFriendNameInListOfFriends(*fFriends);
+      if (friendAlias)
+         return friendAlias;
+   }
+
+   // Then, check if this is a TChain and the current tree has friends
+   // The non-redundant scenario here is that the currently-available
+   // inner TTree of this TChain has a list of friends which the TChain
+   // itself doesn't know anything about.
+   if (const auto *innerListOfFriends = GetTree()->GetListOfFriends();
+       innerListOfFriends && innerListOfFriends != fFriends) {
+      const char *friendAlias = lookForFriendNameInListOfFriends(*innerListOfFriends);
+      if (friendAlias)
+         return friendAlias;
+   }
+
+   // Recursively look into the list of friends of this tree
+   if (fFriends) {
+      for (auto *frEl : ROOT::Detail::TRangeStaticCast<TFriendElement>(*fFriends)) {
+         const char *friendAlias = frEl->GetTree()->GetFriendAlias(tree);
+         if (friendAlias)
+            return friendAlias;
       }
    }
-   // After looking at the first level,
-   // let's see if it is a friend of friends.
-   nextf.Reset();
-   fe = nullptr;
-   while ((fe = (TFriendElement*) nextf())) {
-      const char* res = fe->GetTree()->GetFriendAlias(tree);
-      if (res) {
-         return res;
+
+   // Recursively look into the list of friends of the inner tree
+   if (const auto *innerListOfFriends = GetTree()->GetListOfFriends();
+       innerListOfFriends && innerListOfFriends != fFriends) {
+      for (auto *frEl : ROOT::Detail::TRangeStaticCast<TFriendElement>(*innerListOfFriends)) {
+         const char *friendAlias = frEl->GetTree()->GetFriendAlias(tree);
+         if (friendAlias)
+            return friendAlias;
       }
    }
    return nullptr;
@@ -6612,6 +6642,14 @@ Long64_t TTree::LoadTree(Long64_t entry)
          if (fNotify) {
             if(!fNotify->Notify()) return -6;
          }
+         // We cannot know a priori if the branch(es) of the friend TChain(s) that were just
+         // updated were supposed to be connected to possibly a TChainElement of another chain
+         // that has befriended this TTree (i.e., one of the "external friends"). Thus, we
+         // forward the notification that one or more friend trees were updated to the friends
+         // of this TTree.
+         if (fExternalFriends)
+            for (auto external_fe : ROOT::Detail::TRangeStaticCast<TFriendElement>(*fExternalFriends))
+               external_fe->MarkUpdated();
       }
    }
 

tree/treeplayer/src/TTreePlayer.cxx

@@ -2291,7 +2291,7 @@ Long64_t TTreePlayer::Process(TSelector *selector,Option_t *option, Long64_t nen
       fSelectorUpdate = selector;
       UpdateFormulaLeaves();
 
-      for (entry=firstentry;entry<firstentry+nentries;entry++) {
+      for (entry = firstentry; entry - firstentry < nentries; entry++) {
          entryNumber = fTree->GetEntryNumber(entry);
          if (entryNumber < 0) break;
          if (timer && timer->ProcessEvents()) break;
@@ -2711,9 +2711,7 @@ Long64_t TTreePlayer::Scan(const char *varexp, const char *selection,
    fSelectedRows = 0;
    Int_t tnumber = -1;
    bool exitloop = false;
-   for (entry=firstentry;
-        entry<(firstentry+nentries) && !exitloop;
-        entry++) {
+   for (entry = firstentry; entry - firstentry < nentries && !exitloop; entry++) {
       entryNumber = fTree->GetEntryNumber(entry);
       if (entryNumber < 0) break;
       Long64_t localEntry = fTree->LoadTree(entryNumber);

tree/treeplayer/test/CMSDASClasses.hxx

@@ -0,0 +1,69 @@
+#ifndef CMS_DAS_CLASSES
+#define CMS_DAS_CLASSES
+
+#include <Rtypes.h>
+#include <vector>
+
+namespace DAS {
+struct Weight {
+   int v{1};
+   int i{0};
+   operator int() const { return v; }
+
+   // https://root.cern/manual/io_custom_classes/#restrictions-on-types-root-io-can-handle
+   Weight() = default;
+
+   // For vector initialization in test
+   Weight(int a, int b) : v(a), i(b) {}
+
+   // https://root.cern/manual/io_custom_classes/#the-classdef-macro
+   // NV avoids marking the methods in I/O as virtual (not needed for this struct)
+   ClassDefNV(Weight, 1);
+};
+
+// Also this is a new type and it gets stored to disk, needs a dictionary
+using Weights = std::vector<Weight>;
+
+// Stored, needs a dictionary
+struct AbstractEvent {
+   // https://root.cern/manual/io_custom_classes/#restrictions-on-types-root-io-can-handle
+   AbstractEvent() = default;
+
+   virtual ~AbstractEvent() = default;
+   // Destructor is defined, thus we need to implement rule of five
+   // https://isocpp.github.io/CppCoreGuidelines/CppCoreGuidelines#c21-if-you-define-or-delete-any-copy-move-or-destructor-function-define-or-delete-them-all
+   AbstractEvent(const AbstractEvent &) = default;
+   AbstractEvent &operator=(const AbstractEvent &) = default;
+   AbstractEvent(AbstractEvent &&) = default;
+   AbstractEvent &operator=(AbstractEvent &&) = default;
+
+   Weights weights;
+
+   inline int Weight() const { return weights.front(); }
+
+   // https://root.cern/manual/io_custom_classes/#the-classdef-macro
+   ClassDef(AbstractEvent, 1);
+};
+
+struct GenEvent : public AbstractEvent {
+   // Default constructor for ROOT I/O
+   GenEvent() = default;
+
+   // This is a derived class of a pure virtual class, we need to override all virtual functions
+   // https://isocpp.github.io/CppCoreGuidelines/CppCoreGuidelines#c128-virtual-functions-should-specify-exactly-one-of-virtual-override-or-final
+   ~GenEvent() override = default;
+   // Destructor is defined, thus we need to implement rule of five
+   // https://isocpp.github.io/CppCoreGuidelines/CppCoreGuidelines#c21-if-you-define-or-delete-any-copy-move-or-destructor-function-define-or-delete-them-all
+   GenEvent(const GenEvent &) = default;
+   GenEvent &operator=(const GenEvent &) = default;
+   GenEvent(GenEvent &&) = default;
+   GenEvent &operator=(GenEvent &&) = default;
+
+   // https://root.cern/manual/io_custom_classes/#the-classdef-macro
+   // Override because this is a derived class
+   ClassDefOverride(GenEvent, 1);
+};
+
+} // namespace DAS
+
+#endif // CMS_DAS_CLASSES

tree/treeplayer/test/CMSDASClassesLinkDef.hxx

@@ -0,0 +1,8 @@
+#ifdef __CLING__
+
+#pragma link C++ class DAS::Weight+;
+#pragma link C++ class std::vector<DAS::Weight>+;
+#pragma link C++ class DAS::AbstractEvent+;
+#pragma link C++ class DAS::GenEvent+;
+
+#endif

tree/treeplayer/test/CMakeLists.txt

@@ -19,6 +19,9 @@ ROOT_ADD_GTEST(treeplayer_gh16804 gh16804.cxx LIBRARIES TreePlayer)
 
 ROOT_ADD_GTEST(treeplayer_gh16805 gh16805.cxx LIBRARIES TreePlayer)
 
+ROOT_GENERATE_DICTIONARY(CMSDASClassesDict ${CMAKE_CURRENT_SOURCE_DIR}/CMSDASClasses.hxx LINKDEF CMSDASClassesLinkDef.hxx NO_CXXMODULE DEPENDENCIES Core Tree RIO)
+ROOT_ADD_GTEST(treeplayer_gh20033 gh20033.cxx CMSDASClassesDict.cxx LIBRARIES TreePlayer)
+
 ROOT_ADD_GTEST(treeplayer_leafs leafs.cxx LIBRARIES TreePlayer)
 add_custom_command(TARGET treeplayer_leafs POST_BUILD
                    COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_SOURCE_DIR}/data.h data.h)