[rfile] Add support for reading TTree

core/base/inc/TObject.h

@@ -36,9 +36,11 @@ class TTimer;
 namespace ROOT {
 namespace Internal {
    bool DeleteChangesMemoryImpl();
+   void MarkTObjectAsNotOnHeap(TObject &obj);
 }}
 
 class TObject {
+   friend void ROOT::Internal::MarkTObjectAsNotOnHeap(TObject &obj);
 
 private:
    UInt_t         fUniqueID;   ///< object unique identifier

core/base/src/TObject.cxx

@@ -1265,3 +1265,8 @@ void TObject::operator delete[](void *ptr, void *vp)
 {
    TStorage::ObjectDealloc(ptr, vp);
 }
+
+void ROOT::Internal::MarkTObjectAsNotOnHeap(TObject &obj)
+{
+   obj.fBits &= ~TObject::kIsOnHeap;
+}

io/io/src/RFile.cxx

@@ -304,9 +304,18 @@ void *RFile::GetUntyped(std::string_view path,
    void *obj = key ? key->ReadObjectAny(cls) : nullptr;
 
    if (obj) {
-      // Disavow any ownership on `obj`
-      if (auto autoAddFunc = cls->GetDirectoryAutoAdd(); autoAddFunc) {
-         autoAddFunc(obj, nullptr);
+      // Disavow any ownership on `obj` unless the object is a TTree, in which case we need to link it to our internal
+      // file for it to be usable.
+      if (auto autoAddFunc = cls->GetDirectoryAutoAdd()) {
+         if (cls->InheritsFrom("TTree")) {
+            autoAddFunc(obj, fFile.get());
+            // NOTE(gparolini): this is a hacky but effective way of preventing the Tree from being deleted by
+            // the internal TFile once we close it. We need to avoid that because this TTree will be returned inside
+            // a unique_ptr and would end up being double-freed if we allowed ROOT to do its own memory management.
+            ROOT::Internal::MarkTObjectAsNotOnHeap(*static_cast<TObject *>(obj));
+         } else {
+            autoAddFunc(obj, nullptr);
+         }
       }
    } else if (key) {
       R__LOG_INFO(RFileLog()) << "Tried to get object '" << path << "' of type " << cls->GetName()

io/io/test/CMakeLists.txt

@@ -18,7 +18,19 @@ if(uring AND NOT DEFINED ENV{ROOTTEST_IGNORE_URING})
   ROOT_ADD_GTEST(RIoUring RIoUring.cxx LIBRARIES RIO)
 endif()
 
-ROOT_ADD_GTEST(rfile rfile.cxx LIBRARIES RIO Hist ROOTNTuple)
+ROOT_STANDARD_LIBRARY_PACKAGE(RFileTestIncludes
+                              NO_INSTALL_HEADERS
+                              HEADERS ${CMAKE_CURRENT_SOURCE_DIR}/RFileTestIncludes.hxx
+                              SOURCES RFileTestIncludes.cxx
+                              LINKDEF RFileTestIncludesLinkDef.h
+                              DEPENDENCIES RIO Tree)
+configure_file(RFileTestIncludes.hxx . COPYONLY)
+if(MSVC AND NOT CMAKE_GENERATOR MATCHES Ninja)
+  add_custom_command(TARGET RFileTestIncludes POST_BUILD
+    COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_CURRENT_BINARY_DIR}/$<CONFIG>/libRFileTestIncludes.dll
+                                     ${CMAKE_CURRENT_BINARY_DIR}/libRFileTestIncludes.dll)
+endif()
+ROOT_ADD_GTEST(rfile rfile.cxx LIBRARIES RIO Hist ROOTNTuple Tree RFileTestIncludes)
 if(pyroot)
   ROOT_ADD_PYUNITTEST(rfile_py rfile.py)
 endif()

io/io/test/RFileTestIncludes.cxx

@@ -0,0 +1,17 @@
+#include "RFileTestIncludes.hxx"
+
+TTreeDestructorCounter::~TTreeDestructorCounter() {
+  ++fgTimesDestructed;
+}
+
+void TTreeDestructorCounter::ResetTimesDestructed()
+{
+   fgTimesDestructed = 0;
+}
+
+int TTreeDestructorCounter::GetTimesDestructed()
+{
+   return fgTimesDestructed;
+}
+
+int TTreeDestructorCounter::fgTimesDestructed = 0;

io/io/test/RFileTestIncludes.hxx

@@ -0,0 +1,21 @@
+#ifndef ROOT_RFILE_TEST_INCLUDES
+#define ROOT_RFILE_TEST_INCLUDES
+
+#include <TTree.h>
+
+// WARNING: this class is used in some tests that check its `fgTimesDestructed`.
+// ResetTimesDestructed() must be called before using this class for consistency.
+class TTreeDestructorCounter : public TTree {
+   static int fgTimesDestructed;
+
+public:
+   static void ResetTimesDestructed();
+   static int GetTimesDestructed();
+
+   using TTree::TTree;
+   ~TTreeDestructorCounter();
+
+   ClassDefOverride(TTreeDestructorCounter, 2);
+};
+
+#endif

io/io/test/RFileTestIncludesLinkDef.h

@@ -0,0 +1,5 @@
+#ifdef __CLING__
+
+#pragma link C++ class TTreeDestructorCounter+;
+
+#endif

io/io/test/rfile.cxx

@@ -6,6 +6,7 @@
 #include <TRandom3.h>
 #include <TROOT.h>
 #include <TSystem.h>
+#include <TTree.h>
 #include <RZip.h>
 #include <ROOT/RError.hxx>
 #include <ROOT/RFile.hxx>
@@ -15,6 +16,8 @@
 #include <ROOT/RNTupleReader.hxx>
 #include <ROOT/RNTupleWriter.hxx>
 
+#include "RFileTestIncludes.hxx"
+
 using ROOT::Experimental::RFile;
 using ROOT::TestSupport::FileRaii;
 
@@ -721,3 +724,89 @@ TEST(RFile, RNTuple)
    EXPECT_EQ(reader->GetNEntries(), 1);
    EXPECT_FLOAT_EQ(reader->GetView<float>("x")(0), 42);
 }
+
+TEST(RFile, TTreeRead)
+{
+   FileRaii fileGuard("test_rfile_ttree_read.root");
+
+   {
+      auto file = std::unique_ptr<TFile>(TFile::Open(fileGuard.GetPath().c_str(), "RECREATE"));
+      auto tree = std::make_unique<TTree>("tree", "tree");
+      int x;
+      tree->Branch("x", &x);
+      for (int i = 0; i < 10; ++i) {
+         x = i;
+         tree->Fill();
+      }
+      tree->Write();
+   }
+
+   {
+      auto file = ROOT::Experimental::RFile::Open(fileGuard.GetPath());
+      auto tree = file->Get<TTree>("tree");
+      ASSERT_NE(tree, nullptr);
+      EXPECT_EQ(tree->GetEntries(), 10);
+      int x;
+      tree->SetBranchAddress("x", &x);
+      for (auto i = 0u; i < tree->GetEntries(); ++i) {
+         tree->GetEntry(i);
+         EXPECT_EQ(x, i);
+      }
+   }
+}
+
+TEST(RFile, TTreeReadAfterClose)
+{
+   FileRaii fileGuard("test_rfile_ttree_read_after_close.root");
+
+   {
+      auto file = std::unique_ptr<TFile>(TFile::Open(fileGuard.GetPath().c_str(), "RECREATE"));
+      auto tree = std::make_unique<TTree>("tree", "tree");
+      int x;
+      tree->Branch("x", &x);
+      for (int i = 0; i < 10; ++i) {
+         x = i;
+         tree->Fill();
+      }
+      tree->Write();
+   }
+
+   {
+      auto file = ROOT::Experimental::RFile::Open(fileGuard.GetPath());
+      auto tree = file->Get<TTree>("tree");
+      file.reset(); // close the file
+      ASSERT_NE(tree, nullptr);
+      EXPECT_EQ(tree->GetEntries(), 10);
+      EXPECT_EQ(tree->GetEntry(0), -1); // -1 means "I/O error"
+   }
+}
+
+TEST(RFile, TTreeNoDoubleFree)
+{
+   FileRaii fileGuard("test_rfile_ttree_read_no_double_free.root");
+
+   TTreeDestructorCounter::ResetTimesDestructed();
+
+   {
+      auto file = std::unique_ptr<TFile>(TFile::Open(fileGuard.GetPath().c_str(), "RECREATE"));
+      auto tree = std::make_unique<TTreeDestructorCounter>("tree", "tree");
+      int x;
+      tree->Branch("x", &x);
+      for (int i = 0; i < 10; ++i) {
+         x = i;
+         tree->Fill();
+      }
+      tree->Write();
+   }
+   EXPECT_EQ(TTreeDestructorCounter::GetTimesDestructed(), 1);
+
+   {
+      auto file = ROOT::Experimental::RFile::Open(fileGuard.GetPath());
+      auto tree = file->Get<TTreeDestructorCounter>("tree");
+      file.reset(); // close the file (does not delete the three)
+      // tree is deleted here
+   }
+
+   // destructed once during writing, once during reading.
+   EXPECT_EQ(TTreeDestructorCounter::GetTimesDestructed(), 2);
+}

io/io/test/rfile.py

@@ -144,5 +144,29 @@
         finally:
             os.remove(fileName)
 
+    def test_getTree(self):
+        from array import array
+
+        fileName = "test_rfile_gettree_py.root"
+        try:
+            with ROOT.TFile.Open(fileName, "RECREATE") as file:
+                tree = ROOT.TTree("tree", "tree")
+                x = array("i", [0])
+                tree.Branch("myColumn", x, "myColumn/I")
+                for i in range(10):
+                    x[0] = i
+                    tree.Fill()
+                tree.Write()
+
+            with RFile.Open(fileName) as rfile:
+                tree = rfile.Get("tree")
+                self.assertIsNot(tree, None)
+                self.assertEqual(tree.GetEntries(), 10)
+                xs = [entry.myColumn for entry in tree]
+                self.assertEqual(xs, [x for x in range(10)])
+        finally:
+            os.remove(fileName)
+
+
 if __name__ == "__main__":
     unittest.main()