[tree] Fix interaction between GetEntries and GetListOfFriends

Fixes #17820

Full analysis of the issue follows:

1. TChain::GetEntries calls TChain::LoadTree(TTree::kMaxEntries - 1)

2. Internally, this will set fReadEntry == -2, due to how LoadTree works and the
   fact it's trying to read beyond the last entry in the chain

3. This has an effect here https://github.com/root-project/root/blob/f33985dca2d6bc505e83498f7b4561ddd969be17/tree/tree/src/TChain.cxx#L1675-L1679
   where `t` is the current friend being traversed in the list of friends of the *chain*.
   We call `t->LoadTree(TTree::kMaxEntries - 1)`.

4. Inside the call to LoadTree on the friend, we reach this line
   https://github.com/root-project/root/blob/f33985dca2d6bc505e83498f7b4561ddd969be17/tree/tree/src/TChain.cxx#L1484
   where the fTree data member of the friend is invalidated.

5. Back to the LoadTree of the main chain, these lines
   https://github.com/root-project/root/blob/f33985dca2d6bc505e83498f7b4561ddd969be17/tree/tree/src/TChain.cxx#L1681-L1684
   would normally connect the fTree data member of the main chain to the current friend, so the list of friends
   is properly populated. But due to 4., `friend->GetTree()` will return nullptr so the connection won't happen.

This commit changes the implementation of TChain::GetEntries, so that
after computing the total number of entries the chain is brought back
to an invalid but recoverable state, and if the previous read entry was
at least zero, the cursor is moved to it.

Co-authored-by: Philippe Canal <[email protected]>

Author: vepadulano

tree/tree/src/TChain.cxx

@@ -985,7 +985,12 @@ Long64_t TChain::GetEntries() const
       return fProofChain->GetEntries();
    }
    if (fEntries == TTree::kMaxEntries) {
+      auto readEntry = fReadEntry;
       const_cast<TChain*>(this)->LoadTree(TTree::kMaxEntries-1);
+      const_cast<TChain*>(this)->fTreeNumber = -1;
+      const_cast<TChain*>(this)->fReadEntry = -1;
+      if (readEntry >= 0)
+         const_cast<TChain*>(this)->LoadTree(readEntry);
    }
    return fEntries;
 }

tree/treeplayer/test/CMakeLists.txt

@@ -23,3 +23,5 @@ endif()
 ROOT_ADD_GTEST(ttreeindex_clone ttreeindex_clone.cxx LIBRARIES TreePlayer)
 
 ROOT_ADD_GTEST(ttreereader_friends ttreereader_friends.cxx LIBRARIES TreePlayer)
+
+ROOT_ADD_GTEST(ttreeindex_getlistoffriends ttreeindex_getlistoffriends.cxx LIBRARIES TreePlayer)

tree/treeplayer/test/ttreeindex_getlistoffriends.cxx

@@ -0,0 +1,86 @@
+#include "TChain.h"
+#include "TFile.h"
+#include "TFriendElement.h"
+#include "TTree.h"
+
+#include "gtest/gtest.h"
+
+#include <string_view>
+
+void write_data(std::string_view treename, std::string_view filename)
+{
+   TFile f{filename.data(), "update"};
+   TTree t{treename.data(), treename.data()};
+   int runNumber{};
+   int eventNumber{};
+   float val{};
+   t.Branch("runNumber", &runNumber, "runNumber/I");
+   t.Branch("eventNumber", &eventNumber, "eventNumber/I");
+   t.Branch("val", &val, "val/F");
+   if (treename == "main") {
+      for (auto rn = 0; rn < 3; rn++) {
+         runNumber = rn;
+         for (auto en = 0; en < 5; en++) {
+            eventNumber = en;
+            val = en * rn;
+            t.Fill();
+         }
+      }
+   } else {
+      for (auto rn = 0; rn < 3; rn++) {
+         runNumber = rn;
+         for (auto en = 4; en >= 0; en--) {
+            eventNumber = en;
+            val = en * rn;
+            t.Fill();
+         }
+      }
+   }
+
+   f.Write();
+}
+
+struct TTreeIndexGH_17820 : public ::testing::Test {
+   constexpr static auto fFileName{"ttreeindex_getlistoffriends_gh_17820.root"};
+   constexpr static auto fMainName{"main"};
+   constexpr static auto fFriendName{"friend"};
+
+   static void SetUpTestCase()
+   {
+      write_data(fMainName, fFileName);
+      write_data(fFriendName, fFileName);
+   }
+
+   static void TearDownTestCase() { std::remove(fFileName); }
+};
+
+// Regression test for https://github.com/root-project/root/issues/17820
+TEST_F(TTreeIndexGH_17820, RunTest)
+{
+   TChain mainChain{fMainName};
+   mainChain.AddFile(fFileName);
+
+   TChain friendChain{fFriendName};
+   friendChain.AddFile(fFileName);
+   friendChain.BuildIndex("runNumber", "eventNumber");
+
+   mainChain.AddFriend(&friendChain);
+
+   // Calling GetEntries used to mess with the fTree data member of the main
+   // chain, not connecting it to the friend chain and thus losing the list
+   // of friends
+   mainChain.GetEntries();
+
+   const auto *listOfFriends = mainChain.GetTree()->GetListOfFriends();
+
+   ASSERT_TRUE(listOfFriends);
+   EXPECT_EQ(listOfFriends->GetEntries(), 1);
+
+   const auto *friendTree = dynamic_cast<TTree *>(dynamic_cast<TFriendElement *>(listOfFriends->At(0))->GetTree());
+   ASSERT_TRUE(friendTree);
+
+   EXPECT_STREQ(friendTree->GetName(), fFriendName);
+   const auto *friendFile = friendTree->GetCurrentFile();
+   ASSERT_TRUE(friendFile);
+   EXPECT_STREQ(friendFile->GetName(), fFileName);
+}