From 45198148444583b7f3d1f381d7fb849bef82e675 Mon Sep 17 00:00:00 2001 From: Vincenzo Eduardo Padulano Date: Wed, 9 Jul 2025 18:36:05 +0200 Subject: [PATCH] [tree] Improve logic to refresh branch addresses of TChain friends 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 --> TTree t1 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. --- tree/tree/inc/TChain.h | 1 + tree/tree/inc/TFriendElement.h | 7 +- tree/tree/src/TChain.cxx | 174 +++++++++++++---------- tree/treeplayer/test/CMakeLists.txt | 4 + tree/treeplayer/test/gh16805.cxx | 206 ++++++++++++++++++++++++++++ 5 files changed, 314 insertions(+), 78 deletions(-) create mode 100644 tree/treeplayer/test/gh16805.cxx diff --git a/tree/tree/inc/TChain.h b/tree/tree/inc/TChain.h index 72d12df657a39..55b88add59f9e 100644 --- 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 index 0c17fa2937435..1c78c7b48c9d4 100644 --- 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 index b131210c40bce..2536a0710d4c6 100644 --- a/tree/tree/src/TChain.cxx +++ b/tree/tree/src/TChain.cxx @@ -1231,6 +1231,87 @@ 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(*innerFriendList)) { + if (frEl->IsUpdated()) { + needUpdate = true; + frEl->ResetUpdated(); + } else 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(*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 +1375,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(*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(*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(*fExternalFriends)) { + for (auto external_fe : ROOT::Detail::TRangeStaticCast(*fExternalFriends)) { external_fe->MarkUpdated(); } } diff --git a/tree/treeplayer/test/CMakeLists.txt b/tree/treeplayer/test/CMakeLists.txt index c8229fbb00d60..0d84cd4c802e2 100644 --- 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 new file mode 100644 index 0000000000000..00267daa30c3b --- /dev/null +++ b/tree/treeplayer/test/gh16805.cxx @@ -0,0 +1,206 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +struct GH16805Regression : public ::testing::Test { + + const static inline std::vector fStepZeroFileNames{"gh16805_regression_stepzero_0.root", + "gh16805_regression_stepzero_1.root"}; + const static inline std::vector> fStepZeroEntryRanges{{0, 10}, {10, 20}}; + constexpr static auto fStepZeroTreeName{"stepzero"}; + + const static inline std::vector fTopLevelFriendFileNames{"gh16805_regression_toplevelfriend_0.root", + "gh16805_regression_toplevelfriend_1.root", + "gh16805_regression_toplevelfriend_2.root"}; + const static inline std::vector> fTopLevelFriendEntryRanges{{200, 207}, {207, 214}, {214, 220}}; + constexpr static auto fTopLevelFriendTreeName{"topLevelFriend"}; + + constexpr static auto fStepOneFileName{"gh16805_regression_stepone.root"}; + constexpr static std::pair fStepOneEntryRange{100, 120}; + constexpr static auto fStepOneTreeName{"stepone"}; + + static void writeStepZero(const char *treeName, const char *fileName, int begin, int end) + { + auto file = std::make_unique(fileName, "recreate"); + auto tree = std::make_unique(treeName, treeName); + int stepZeroBr1, stepZeroBr2 = 0; + tree->Branch("stepZeroBr1", &stepZeroBr1); + tree->Branch("stepZeroBr2", &stepZeroBr2); + + for (stepZeroBr1 = begin, stepZeroBr2 = begin * 2; stepZeroBr1 < end; + ++stepZeroBr1, stepZeroBr2 = stepZeroBr1 * 2) + tree->Fill(); + file->Write(); + } + + static void writeStepOne(const char *treeName, const char *fileName, int begin, int end) + { + auto file = std::make_unique(fileName, "recreate"); + auto tree = std::make_unique(treeName, treeName); + int br1 = -1; + tree->Branch("stepOneBr1", &br1); + + for (br1 = begin; br1 < end; ++br1) + tree->Fill(); + + file->Write(); + } + + static void writeTopLevelFriend(const char *treeName, const char *fileName, int begin, int end) + { + auto file = std::make_unique(fileName, "recreate"); + auto tree = std::make_unique(treeName, treeName); + int friendBr1, friendBr2 = 0; + tree->Branch("friendBr1", &friendBr1); + tree->Branch("friendBr2", &friendBr2); + + for (friendBr1 = begin, friendBr2 = begin * 2; friendBr1 < end; ++friendBr1, friendBr2 = friendBr1 * 2) + tree->Fill(); + file->Write(); + } + + static void SetUpTestSuite() + { + for (decltype(fStepZeroFileNames.size()) i = 0; i < fStepZeroFileNames.size(); i++) { + writeStepZero(fStepZeroTreeName, fStepZeroFileNames[i].c_str(), fStepZeroEntryRanges[i].first, + fStepZeroEntryRanges[i].second); + } + + for (decltype(fTopLevelFriendFileNames.size()) i = 0; i < fTopLevelFriendFileNames.size(); i++) { + writeTopLevelFriend(fTopLevelFriendTreeName, fTopLevelFriendFileNames[i].c_str(), + fTopLevelFriendEntryRanges[i].first, fTopLevelFriendEntryRanges[i].second); + } + writeStepOne(fStepOneTreeName, fStepOneFileName, fStepOneEntryRange.first, fStepOneEntryRange.second); + } + + static void TearDownTestSuite() + { + for (const auto &f : fStepZeroFileNames) + std::remove(f.c_str()); + + for (const auto &f : fTopLevelFriendFileNames) + std::remove(f.c_str()); + + std::remove(fStepOneFileName); + } +}; + +TEST_F(GH16805Regression, Test) +{ + // Tests that a chain is able to correctly update branch addresses of its friends, + // whether the friends are attached to the chain itself or to its first tree. + + auto stepZeroChain = std::make_unique(fStepZeroTreeName); + for (const auto &f : fStepZeroFileNames) + stepZeroChain->Add(f.c_str()); + + auto stepOneChain = std::make_unique(fStepOneTreeName); + stepOneChain->Add(fStepOneFileName); + + auto topLevelFriendChain = std::make_unique(fTopLevelFriendTreeName); + for (const auto &f : fTopLevelFriendFileNames) + topLevelFriendChain->Add(f.c_str()); + stepOneChain->AddFriend(topLevelFriendChain.get()); + + // Load the first tree in the chain + stepOneChain->LoadTree(0); + auto stepOneFirstTree = stepOneChain->GetTree()->GetTree(); + ASSERT_NE(stepOneFirstTree, nullptr); + + stepOneFirstTree->AddFriend(stepZeroChain.get()); + + // Now iterate over the chain and check the contents of the branches + // inherited from the friends. + int stepZeroBr1 = -1, stepZeroBr2 = -1, stepOneBr1 = -1, friendBr1 = -1, friendBr2 = -1; + ASSERT_EQ(stepOneChain->SetBranchAddress("stepZeroBr1", &stepZeroBr1), 0); + ASSERT_EQ(stepOneChain->SetBranchAddress("stepZeroBr2", &stepZeroBr2), 0); + ASSERT_EQ(stepOneChain->SetBranchAddress("stepOneBr1", &stepOneBr1), 0); + ASSERT_EQ(stepOneChain->SetBranchAddress("friendBr1", &friendBr1), 0); + ASSERT_EQ(stepOneChain->SetBranchAddress("friendBr2", &friendBr2), 0); + + std::vector expectedStepZeroBr1(20); + std::vector expectedStepZeroBr2(20); + std::vector expectedStepOneBr1(20); + std::vector expectedTopLevelFriendBr1(20); + std::vector expectedTopLevelFriendBr2(20); + + std::iota(expectedStepZeroBr1.begin(), expectedStepZeroBr1.end(), 0); + std::generate(expectedStepZeroBr2.begin(), expectedStepZeroBr2.end(), [n = 0]() mutable { + auto res = n * 2; + n++; + return res; + }); + std::iota(expectedStepOneBr1.begin(), expectedStepOneBr1.end(), 100); + std::iota(expectedTopLevelFriendBr1.begin(), expectedTopLevelFriendBr1.end(), 200); + std::generate(expectedTopLevelFriendBr2.begin(), expectedTopLevelFriendBr2.end(), [n = 200]() mutable { + auto res = n * 2; + n++; + return res; + }); + + for (Long64_t i = 0; i < stepOneChain->GetEntriesFast(); ++i) { + stepOneChain->GetEntry(i); + EXPECT_EQ(expectedStepZeroBr1[i], stepZeroBr1); + EXPECT_EQ(expectedStepZeroBr2[i], stepZeroBr2); + EXPECT_EQ(expectedStepOneBr1[i], stepOneBr1); + EXPECT_EQ(expectedTopLevelFriendBr1[i], friendBr1); + EXPECT_EQ(expectedTopLevelFriendBr2[i], friendBr2); + } + + // Now test with TTree::Scan + std::ostringstream strCout; + { + if (auto *treePlayer = static_cast(stepOneChain->GetPlayer())) { + struct FileRAII { + const char *fPath; + FileRAII(const char *name) : fPath(name) {} + ~FileRAII() { std::remove(fPath); } + } redirectFile{"regression_16805_redirect.txt"}; + treePlayer->SetScanRedirect(true); + treePlayer->SetScanFileName(redirectFile.fPath); + stepOneChain->Scan("stepZeroBr1:stepZeroBr2:stepOneBr1:friendBr1:friendBr2", "", "colsize=12"); + + std::ifstream redirectStream(redirectFile.fPath); + std::stringstream redirectOutput; + redirectOutput << redirectStream.rdbuf(); + + const static std::string expectedScanOut{ + R"Scan(*************************************************************************************** +* Row * stepZeroBr1 * stepZeroBr2 * stepOneBr1 * friendBr1 * friendBr2 * +*************************************************************************************** +* 0 * 0 * 0 * 100 * 200 * 400 * +* 1 * 1 * 2 * 101 * 201 * 402 * +* 2 * 2 * 4 * 102 * 202 * 404 * +* 3 * 3 * 6 * 103 * 203 * 406 * +* 4 * 4 * 8 * 104 * 204 * 408 * +* 5 * 5 * 10 * 105 * 205 * 410 * +* 6 * 6 * 12 * 106 * 206 * 412 * +* 7 * 7 * 14 * 107 * 207 * 414 * +* 8 * 8 * 16 * 108 * 208 * 416 * +* 9 * 9 * 18 * 109 * 209 * 418 * +* 10 * 10 * 20 * 110 * 210 * 420 * +* 11 * 11 * 22 * 111 * 211 * 422 * +* 12 * 12 * 24 * 112 * 212 * 424 * +* 13 * 13 * 26 * 113 * 213 * 426 * +* 14 * 14 * 28 * 114 * 214 * 428 * +* 15 * 15 * 30 * 115 * 215 * 430 * +* 16 * 16 * 32 * 116 * 216 * 432 * +* 17 * 17 * 34 * 117 * 217 * 434 * +* 18 * 18 * 36 * 118 * 218 * 436 * +* 19 * 19 * 38 * 119 * 219 * 438 * +*************************************************************************************** +)Scan"}; + EXPECT_EQ(redirectOutput.str(), expectedScanOut); + } else + throw std::runtime_error("Could not retrieve TTreePlayer from main tree!"); + } +}