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..fb99b7290bd3f 100644 --- 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(*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(*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(*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!"); + } +}