[tree] Improve logic to refresh branch addresses of TChain friends

vepadulano · vepadulano · commit b0868f3c3f34 · 2025-07-29T21:55:59.000+02:00
Part of the logic of TChain::LoadTree deals with updating the status of friend
trees. First, the current entry is loaded on all friends. In case the friend
tree is a chain that needs to switch to a new tree, the branch addresses of the
friend are also updated.

This logic only considered the case of top-level friends of a TChain. Notably,
it did not consider the case of a main TChain where one (or more) of the trees
in the chain have themselves some friends, but not the main chain itself:

TChain c1
--&gt; TTree t1 &lt;friendswith&gt; TChain c2

This commit condenses the logic that updates branch addresses of friends in a
separate method of TChain, namely RefreshFriendAddresses. The new method is
called in TChain::LoadTree when either the chain itself or the current tree
have friends.

Part of the new logic relies on an addition to TFriendElement. In
particular, when the current tree of the chain is loading a new entry
and this results in loading the entry of one of its friends which may
have to switch to a new tree, the friend calls MarkUpdated to let the
tree know it was updated. In case the tree belongs to the chain, it
can't know a priori which branches are requested by the user (an
information stored in the fStatus data member of TChain). Thus,
TFriendElement now has two bits that represent the updated status. The
first is the preexisting kUpdated, the second is the new
kUpdatedForChain. A call to MarkUpdated sets both bits. In
TTree::LoadTree, only the preexisting bit is checked and reset. The new
bit is instead checked and reset only in TChain::RefreshFriendAddresses.

A new test is added with a main TChain that has both a top-level TChain
friend and where its only tree is also connected to another friend
TChain.
diff --git a/tree/tree/inc/TChain.h b/tree/tree/inc/TChain.h
@@ -49,6 +49,7 @@ class TChain : public TTree {
    TChain& operator=(const TChain&); // not implemented
    void
    ParseTreeFilename(const char *name, TString &filename, TString &treename, TString &query, TString &suffix) const;
+   Long64_t RefreshFriendAddresses();
 
 protected:
    void InvalidateCurrentTree();
diff --git a/tree/tree/inc/TFriendElement.h b/tree/tree/inc/TFriendElement.h
@@ -47,7 +47,8 @@ class TFriendElement : public TNamed {
 public:
    enum EStatusBits {
       kFromChain = BIT(9),  // Indicates a TChain inserted this element in one of its content TTree
-      kUpdated   = BIT(10)  // Indicates that the chain 'fTree' went through a LoadTree
+      kUpdated   = BIT(10),  // Indicates that the chain 'fTree' went through a LoadTree
+      kUpdatedForChain = BIT(11), // Same as kUpdated, but detected only by the chain itself and not the current tree
    };
    TFriendElement();
    TFriendElement(TTree *tree, const char *treename, const char *filename);
@@ -65,8 +66,10 @@ class TFriendElement : public TNamed {
    void        ls(Option_t *option="") const override;
            void        Reset() { fTree = nullptr; fFile = nullptr; }
            bool        IsUpdated() const { return TestBit(kUpdated); }
+           bool        IsUpdatedForChain() const { return TestBit(kUpdatedForChain); }
            void        ResetUpdated() { ResetBit(kUpdated); }
-           void        MarkUpdated() { SetBit(kUpdated); }
+           void        ResetUpdatedForChain() { ResetBit(kUpdatedForChain); }
+           void        MarkUpdated() { SetBit(kUpdated); SetBit(kUpdatedForChain); }
    void        RecursiveRemove(TObject *obj) override;
 
 
diff --git a/tree/tree/src/TChain.cxx b/tree/tree/src/TChain.cxx
@@ -1231,6 +1231,88 @@ Int_t TChain::LoadBaskets(Long64_t /*maxmemory*/)
    return 0;
 }
 
+////////////////////////////////////////////////////////////////////////////////
+/// Refresh branch/leaf addresses of friend trees
+///
+/// The method acts only on the current tree in the chain (fTree), but it may
+/// be called in two different scenarios: when there are friends of the chain
+/// or when there are friends of fTree itself.
+Long64_t TChain::RefreshFriendAddresses()
+{
+   assert(fTree != nullptr);
+
+   bool needUpdate = false;
+   if (auto *innerFriendList = fTree->GetListOfFriends()) {
+      // If the current tree has friends, check if they were mark for update
+      // when switching to the following tree, detect it so that we later we
+      // actually refresh the addresses of the friends.
+      for (auto *frEl : ROOT::Detail::TRangeStaticCast<TFriendElement>(*innerFriendList)) {
+         if (frEl->IsUpdated()) {
+            needUpdate = true;
+            frEl->ResetUpdated();
+         }
+         if (frEl->IsUpdatedForChain()) {
+            needUpdate = true;
+            frEl->ResetUpdatedForChain();
+         }
+      }
+   }
+
+   if (!needUpdate)
+      return 0;
+
+   // Update the branch/leaf addresses and the list of leaves in all
+   // TTreeFormula of the TTreePlayer (if any).
+   for (auto *chainEl : ROOT::Detail::TRangeStaticCast<TChainElement>(*fStatus)) {
+      // Set the branch status of all the chain elements, which may include also
+      // branches that are available in friends. Only set the branch status
+      // if it has a value provided by the user
+      Int_t status = chainEl->GetStatus();
+      if (status != -1)
+         fTree->SetBranchStatus(chainEl->GetName(), status);
+
+      // Set the branch addresses for the newly opened file.
+      void *addr = chainEl->GetBaddress();
+      if (!addr)
+         continue;
+
+      TBranch *br = fTree->GetBranch(chainEl->GetName());
+      TBranch **pp = chainEl->GetBranchPtr();
+      if (pp) {
+         // FIXME: What if br is zero here?
+         *pp = br;
+      }
+      if (!br)
+         continue;
+
+      if (!chainEl->GetCheckedType()) {
+         Int_t res = CheckBranchAddressType(br, TClass::GetClass(chainEl->GetBaddressClassName()),
+                                            (EDataType)chainEl->GetBaddressType(), chainEl->GetBaddressIsPtr());
+         if ((res & kNeedEnableDecomposedObj) && !br->GetMakeClass()) {
+            br->SetMakeClass(true);
+         }
+         chainEl->SetDecomposedObj(br->GetMakeClass());
+         chainEl->SetCheckedType(true);
+      }
+      // FIXME: We may have to tell the branch it should
+      //        not be an owner of the object pointed at.
+      br->SetAddress(addr);
+      if (TestBit(kAutoDelete)) {
+         br->SetAutoDelete(true);
+      }
+   }
+   if (fPlayer) {
+      fPlayer->UpdateFormulaLeaves();
+   }
+   // Notify user if requested.
+   if (fNotify) {
+      if (!fNotify->Notify())
+         return -6;
+   }
+
+   return 0;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 /// Find the tree which contains entry, and set it as the current tree.
 ///
@@ -1294,93 +1376,34 @@ Long64_t TChain::LoadTree(Long64_t entry)
 
    // If entry belongs to the current tree return entry.
    if (fTree && treenum == fTreeNumber) {
-      // First set the entry the tree on its owns friends
-      // (the friends of the chain will be updated in the
-      // next loop).
+      // First load entry on the current tree, this will set the cursor also
+      // on its friend trees. Their branch addresses cannot be updated yet,
+      // as the required branch names are only available via `fStatus`, i.e.
+      // only the chain knows about them. This is taken care of in the following
+      // RefreshFriendAddresses call
       fTree->LoadTree(treeReadEntry);
 
-      if (fFriends) {
-         // The current tree has not changed but some of its friends might.
-         //
-         TIter next(fFriends);
+      if (fFriends || fTree->GetListOfFriends()) {
          TFriendLock lock(this, kLoadTree);
-         TFriendElement* fe = nullptr;
-         while ((fe = (TFriendElement*) next())) {
-            TTree* at = fe->GetTree();
-            // If the tree is a
-            // direct friend of the chain, it should be scanned
-            // used the chain entry number and NOT the tree entry
-            // number (treeReadEntry) hence we do:
-            at->LoadTreeFriend(entry, this);
-         }
-         bool needUpdate = false;
-         if (fTree->GetListOfFriends()) {
-            for(auto fetree : ROOT::Detail::TRangeStaticCast<TFriendElement>(*fTree->GetListOfFriends())) {
-               if (fetree->IsUpdated()) {
-                  needUpdate = true;
-                  fetree->ResetUpdated();
-               }
+         if (fFriends) {
+            // Make sure we load friends of the chain aligned to the current global entry number
+            for (auto *frEl : ROOT::Detail::TRangeStaticCast<TFriendElement>(*fFriends)) {
+               auto *frTree = frEl->GetTree();
+               frTree->LoadTreeFriend(entry, this);
             }
          }
-         if (needUpdate) {
-            // Update the branch/leaf addresses and
-            // the list of leaves in all TTreeFormula of the TTreePlayer (if any).
-
-            // Set the branch statuses for the newly opened file.
-            TChainElement *frelement;
-            TIter fnext(fStatus);
-            while ((frelement = (TChainElement*) fnext())) {
-               Int_t status = frelement->GetStatus();
-               // Only set the branch status if it has a value provided
-               // by the user
-               if (status != -1)
-                  fTree->SetBranchStatus(frelement->GetName(), status);
-            }
 
-            // Set the branch addresses for the newly opened file.
-            fnext.Reset();
-            while ((frelement = (TChainElement*) fnext())) {
-               void* addr = frelement->GetBaddress();
-               if (addr) {
-                  TBranch* br = fTree->GetBranch(frelement->GetName());
-                  TBranch** pp = frelement->GetBranchPtr();
-                  if (pp) {
-                     // FIXME: What if br is zero here?
-                     *pp = br;
-                  }
-                  if (br) {
-                     if (!frelement->GetCheckedType()) {
-                        Int_t res = CheckBranchAddressType(br, TClass::GetClass(frelement->GetBaddressClassName()),
-                                                         (EDataType) frelement->GetBaddressType(), frelement->GetBaddressIsPtr());
-                        if ((res & kNeedEnableDecomposedObj) && !br->GetMakeClass()) {
-                           br->SetMakeClass(true);
-                        }
-                        frelement->SetDecomposedObj(br->GetMakeClass());
-                        frelement->SetCheckedType(true);
-                     }
-                     // FIXME: We may have to tell the branch it should
-                     //        not be an owner of the object pointed at.
-                     br->SetAddress(addr);
-                     if (TestBit(kAutoDelete)) {
-                        br->SetAutoDelete(true);
-                     }
-                  }
-               }
-            }
-            if (fPlayer) {
-               fPlayer->UpdateFormulaLeaves();
-            }
-            // Notify user if requested.
-            if (fNotify) {
-               if(!fNotify->Notify()) return -6;
-            }
-         }
+         // Now refresh branch addresses of friend trees. This acts on the current tree of the chain, whether the
+         // friends are friends of the chain or friends of the tree itself.
+         if (auto refreshFriendAddressesRet = RefreshFriendAddresses(); refreshFriendAddressesRet == -6)
+            return refreshFriendAddressesRet;
       }
+
       return treeReadEntry;
    }
 
    if (fExternalFriends) {
-      for(auto external_fe : ROOT::Detail::TRangeStaticCast<TFriendElement>(*fExternalFriends)) {
+      for (auto external_fe : ROOT::Detail::TRangeStaticCast<TFriendElement>(*fExternalFriends)) {
          external_fe->MarkUpdated();
       }
    }
diff --git a/tree/treeplayer/test/CMakeLists.txt b/tree/treeplayer/test/CMakeLists.txt
@@ -17,6 +17,8 @@ endif()
 
 ROOT_ADD_GTEST(treeplayer_gh16804 gh16804.cxx LIBRARIES TreePlayer)
 
+ROOT_ADD_GTEST(treeplayer_gh16805 gh16805.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)
@@ -39,3 +41,5 @@ if(imt)
       ROOT_ADD_GTEST(treeprocessormt_remotefiles treeprocs/treeprocessormt_remotefiles.cxx LIBRARIES TreePlayer)
    endif()
 endif()
+
+
diff --git a/tree/treeplayer/test/gh16805.cxx b/tree/treeplayer/test/gh16805.cxx
