[df] Move RNoCleanupNotifier usage to RTTreeDS

To keep the implementation aligned with the previously existing one from
RLoopManager. As a result of this, the header `InternalTreeUtils.hxx`
can be removed from the includes of `RLoopManager.hxx`, removing a
further header dependency on TTree. This change trickled some missing
header errors in other places that have been addressed.

Author: vepadulano

tree/dataframe/inc/ROOT/RDF/RInterface.hxx

@@ -1411,18 +1411,15 @@ public:
                                                  const RSnapshotOptions &options = RSnapshotOptions())
    {
       const auto definedColumns = fColRegister.GenerateColumnNames();
-      auto *tree = fLoopManager->GetTree();
 
-      const auto treeBranchNames = tree != nullptr ? ROOT::Internal::TreeUtils::GetTopLevelBranchNames(*tree) : ColumnNames_t{};
       const auto dsColumns = GetDataSource() ? ROOT::Internal::RDF::GetTopLevelFieldNames(*GetDataSource()) : ColumnNames_t{};
       // Ignore R_rdf_sizeof_* columns coming from datasources: we don't want to Snapshot those
       ColumnNames_t dsColumnsWithoutSizeColumns;
       std::copy_if(dsColumns.begin(), dsColumns.end(), std::back_inserter(dsColumnsWithoutSizeColumns),
                    [](const std::string &name) { return name.size() < 13 || name.substr(0, 13) != "R_rdf_sizeof_"; });
       ColumnNames_t columnNames;
-      columnNames.reserve(definedColumns.size() + treeBranchNames.size() + dsColumnsWithoutSizeColumns.size());
+      columnNames.reserve(definedColumns.size() + dsColumnsWithoutSizeColumns.size());
       columnNames.insert(columnNames.end(), definedColumns.begin(), definedColumns.end());
-      columnNames.insert(columnNames.end(), treeBranchNames.begin(), treeBranchNames.end());
       columnNames.insert(columnNames.end(), dsColumnsWithoutSizeColumns.begin(), dsColumnsWithoutSizeColumns.end());
 
       // The only way we can get duplicate entries is if a column coming from a tree or data-source is Redefine'd.
@@ -1560,18 +1557,14 @@ public:
    RInterface<RLoopManager> Cache(std::string_view columnNameRegexp = "")
    {
       const auto definedColumns = fColRegister.GenerateColumnNames();
-      auto *tree = fLoopManager->GetTree();
-      const auto treeBranchNames =
-         tree != nullptr ? ROOT::Internal::TreeUtils::GetTopLevelBranchNames(*tree) : ColumnNames_t{};
       const auto dsColumns = GetDataSource() ? GetDataSource()->GetColumnNames() : ColumnNames_t{};
       // Ignore R_rdf_sizeof_* columns coming from datasources: we don't want to Snapshot those
       ColumnNames_t dsColumnsWithoutSizeColumns;
       std::copy_if(dsColumns.begin(), dsColumns.end(), std::back_inserter(dsColumnsWithoutSizeColumns),
                    [](const std::string &name) { return name.size() < 13 || name.substr(0, 13) != "R_rdf_sizeof_"; });
       ColumnNames_t columnNames;
-      columnNames.reserve(definedColumns.size() + treeBranchNames.size() + dsColumns.size());
+      columnNames.reserve(definedColumns.size() + dsColumns.size());
       columnNames.insert(columnNames.end(), definedColumns.begin(), definedColumns.end());
-      columnNames.insert(columnNames.end(), treeBranchNames.begin(), treeBranchNames.end());
       columnNames.insert(columnNames.end(), dsColumns.begin(), dsColumns.end());
       const auto selectedColumns = RDFInternal::ConvertRegexToColumns(columnNames, columnNameRegexp, "Cache");
       return Cache(selectedColumns);

tree/dataframe/inc/ROOT/RDF/RLoopManager.hxx

@@ -11,7 +11,6 @@
 #ifndef ROOT_RLOOPMANAGER
 #define ROOT_RLOOPMANAGER
 
-#include "ROOT/InternalTreeUtils.hxx" // RNoCleanupNotifier
 #include "ROOT/RDataSource.hxx"
 #include "ROOT/RDF/RColumnReaderBase.hxx"
 #include "ROOT/RDF/RDatasetSpec.hxx"
@@ -129,6 +128,13 @@ class RLoopManager : public RNodeBase {
    friend struct RCallCleanUpTask;
    friend struct ROOT::Internal::RDF::RDSRangeRAII;
 
+   /*
+   A wrapped reference to a TTree dataset that can be shared by many computation graphs. Ensures lifetime management.
+   It needs to be destroyed *after* the fDataSource data member, in case it is holding the only Python reference to the
+   input dataset and the data source needs to access it before it is destroyed.
+   */
+   std::any fTTreeLifeline{};
+
    std::vector<RDFInternal::RActionBase *> fBookedActions; ///< Non-owning pointers to actions to be run
    std::vector<RDFInternal::RActionBase *> fRunActions;    ///< Non-owning pointers to actions already run
    std::vector<RFilterBase *> fBookedFilters;
@@ -172,7 +178,6 @@ class RLoopManager : public RNodeBase {
    /// Cache of the tree/chain branch names. Never access directy, always use GetBranchNames().
    ColumnNames_t fValidBranchNames;
 
-   ROOT::Internal::TreeUtils::RNoCleanupNotifier fNoCleanupNotifier;
    /// Pointer to a shared slot stack in case this instance runs concurrently with others:
    std::weak_ptr<ROOT::Internal::RSlotStack> fSlotStack;
 
@@ -201,10 +206,6 @@ class RLoopManager : public RNodeBase {
    std::set<std::pair<std::string_view, std::unique_ptr<ROOT::Internal::RDF::RVariationsWithReaders>>>
       fUniqueVariationsWithReaders;
 
-   /// A wrapped reference to a TTree dataset that can be shared by many computation graphs. Ensures lifetime
-   /// management.
-   std::any fTTreeLifeline{};
-
 public:
    RLoopManager(const ColumnNames_t &defaultColumns = {});
    RLoopManager(TTree *tree, const ColumnNames_t &defaultBranches);

tree/dataframe/inc/ROOT/RTTreeDS.hxx

@@ -46,6 +46,10 @@ namespace ROOT::TreeUtils {
 struct RFriendInfo;
 }
 
+namespace ROOT::Internal::TreeUtils {
+class RNoCleanupNotifier;
+}
+
 class TChain;
 class TDirectory;
 class TTree;
@@ -66,6 +70,10 @@ class RTTreeDS final : public ROOT::RDF::RDataSource {
 
    std::vector<std::unique_ptr<TChain>> fFriends;
 
+   // Should be needed mostly for MT runs, but we keep it here to document and align the existing functionality
+   // from RLoopManager. See https://github.com/root-project/root/pull/10729
+   std::unique_ptr<ROOT::Internal::TreeUtils::RNoCleanupNotifier> fNoCleanupNotifier;
+
    ROOT::RDF::RSampleInfo
    CreateSampleInfo(const std::unordered_map<std::string, ROOT::RDF::Experimental::RSample *> &sampleMap) const final;
 

tree/dataframe/src/RLoopManager.cxx

@@ -32,6 +32,7 @@
 #include "TROOT.h" // IsImplicitMTEnabled, gCoreMutex, R__*_LOCKGUARD
 #include "TTreeReader.h"
 #include "TTree.h" // For MaxTreeSizeRAII. Revert when #6640 will be solved.
+#include "ROOT/InternalTreeUtils.hxx" // GetTopLevelBranchNames
 
 #include "ROOT/RTTreeDS.hxx"
 
@@ -534,6 +535,7 @@ void RLoopManager::RunEmptySource()
    }
 }
 
+#ifdef R__USE_IMT
 namespace {
 /// Return true on succesful entry read.
 ///
@@ -558,6 +560,7 @@ bool validTTreeReaderRead(TTreeReader &treeReader)
    }
 }
 } // namespace
+#endif
 
 namespace {
 struct DSRunRAII {

tree/dataframe/src/RTTreeDS.cxx

@@ -1,17 +1,18 @@
 #include <sstream>
 
-#include <ROOT/InternalTreeUtils.hxx> // GetTopLevelBranchNames
-#include <ROOT/RDataFrame.hxx>
-#include <ROOT/RFriendInfo.hxx>
 #include <ROOT/RTTreeDS.hxx>
+#include <ROOT/RDataFrame.hxx>
 #include <ROOT/RDF/RLoopManager.hxx> // GetBranchNames
 #include <ROOT/RDF/RTreeColumnReader.hxx>
 #include <ROOT/RDF/Utils.hxx> // GetBranchOrLeafTypeName
 
 #include <TClassEdit.h>
 #include <TFile.h>
 #include <TTree.h>
+#include <TChain.h>
 #include <TTreeReader.h>
+#include <ROOT/RFriendInfo.hxx>
+#include <ROOT/InternalTreeUtils.hxx> // GetTopLevelBranchNames
 
 #ifdef R__USE_IMT
 #include <TROOT.h>
@@ -62,12 +63,24 @@ GetCollectionInfo(const std::string &typeName)
 } // namespace
 
 // Destructor is defined here, where the data member types are actually available
-ROOT::Internal::RDF::RTTreeDS::~RTTreeDS() = default;
+ROOT::Internal::RDF::RTTreeDS::~RTTreeDS()
+{
+   if (fNoCleanupNotifier && fTree)
+      // fNoCleanupNotifier was created only if the input TTree is a TChain
+      fNoCleanupNotifier->RemoveLink(*static_cast<TChain *>(fTree.get()));
+};
 
 void ROOT::Internal::RDF::RTTreeDS::Setup(std::shared_ptr<TTree> &&tree, const ROOT::TreeUtils::RFriendInfo *friendInfo)
 {
    fTree = tree;
 
+   // We keep the existing functionality from RLoopManager, until proven not necessary.
+   // See https://github.com/root-project/root/pull/10729
+   if (auto ch = dynamic_cast<TChain *>(fTree.get()); ch && !fNoCleanupNotifier) {
+      fNoCleanupNotifier = std::make_unique<ROOT::Internal::TreeUtils::RNoCleanupNotifier>();
+      fNoCleanupNotifier->RegisterChain(*ch);
+   }
+
    if (friendInfo) {
       fFriends = ROOT::Internal::TreeUtils::MakeFriends(*friendInfo);
       for (std::size_t i = 0ul; i < fFriends.size(); i++) {

tree/dataframe/test/dataframe_interface.cxx

@@ -9,6 +9,7 @@
 #include "TMemFile.h"
 #include "TSystem.h"
 #include "TTree.h"
+#include "TChain.h"
 
 #include "gtest/gtest.h"
 

tree/dataframe/test/dataframe_snapshot_emptyoutput.cxx

@@ -4,6 +4,7 @@
 // We only use this header to generate the dictionary for RVec<ROOT::Math::PtEtaPhiMVector>
 #include "DummyHeader.hxx"
 
+#include <ROOT/InternalTreeUtils.hxx>
 #include <ROOT/RDataFrame.hxx>
 
 #include <thread>