[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 cursor on the chain is
brought back to its previous entry, or at least 0 if it was in an
invalid state before. TChain::LoadTree is also modified so that if the
current tree is not connected to the list of friends, the connection is
refreshed.

Author: vepadulano

tree/tree/src/TChain.cxx

@@ -985,7 +985,9 @@ Long64_t TChain::GetEntries() const
       return fProofChain->GetEntries();
    }
    if (fEntries == TTree::kMaxEntries) {
+      const auto readEntry = GetReadEntry();
       const_cast<TChain*>(this)->LoadTree(TTree::kMaxEntries-1);
+      const_cast<TChain *>(this)->LoadTree(std::max(readEntry, 0ll));
    }
    return fEntries;
 }
@@ -1373,13 +1375,21 @@ Long64_t TChain::LoadTree(Long64_t entry)
          TIter next(fFriends);
          TFriendLock lock(this, kLoadTree);
          TFriendElement* fe = nullptr;
+         const bool reconstructLocalTreeListOfFriends = fTree->GetListOfFriends() ? false : true;
          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);
+            if (reconstructLocalTreeListOfFriends) {
+               TTree *friend_t = at->GetTree();
+               if (friend_t) {
+                  auto localfe = fTree->AddFriend(at, fe->GetName());
+                  localfe->SetBit(TFriendElement::kFromChain);
+               }
+            }
          }
          bool needUpdate = false;
          if (fTree->GetListOfFriends()) {

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);
+}