[rfile] add GetKeys() methods

Author: silverweed

io/io/inc/ROOT/RFile.hxx

@@ -11,11 +11,13 @@
 #include <ROOT/RError.hxx>
 
 #include <memory>
+#include <iostream>
 #include <string_view>
 #include <typeinfo>
 
-class TFile;
+class TIterator;
 class TKey;
+class TFile;
 
 namespace ROOT {
 namespace Experimental {
@@ -29,6 +31,93 @@ ROOT::RLogChannel &RFileLog();
 
 } // namespace Internal
 
+/**
+\class ROOT::Experimental::RFileKeyInfo
+\ingroup RFile
+\brief Information about an RFile object's Key.
+
+Every object inside a ROOT file has an associated "Key" which contains metadata on the object, such as its name, type
+etc.
+Querying this information can be done via RFile::GetKeys() or RFile::GetKeysNonRecursive. Reading an object's Key
+doesn't deserialize the full object, so it's a relatively lightweight operation.
+*/
+struct RFileKeyInfo {
+   std::string fName;
+   std::string fTitle;
+   std::string fClassName;
+   std::uint16_t fCycle;
+};
+
+/// The iterable returned by RFile::GetKeys() and RFile::GetKeysNonRecursive()
+class RFileKeyIterable final {
+   using Pattern_t = std::string;
+
+   TFile *fFile;
+   Pattern_t fPattern;
+   std::uint32_t fFlags = 0;
+
+public:
+   enum EFlags {
+      kNone = 0,
+      kRecursive = 1 << 0,
+   };
+
+   class RIterator {
+      friend class RFileKeyIterable;
+
+      struct RIterStackElem {
+         // This is ugly, but TList returns an (owning) pointer to a polymorphic TIterator...and we need this class
+         // to be copy-constructible.
+         std::shared_ptr<TIterator> fIter;
+         std::string fDirPath;
+
+         // Outlined to avoid including TIterator.h
+         RIterStackElem(TIterator *it, const std::string &path = "");
+         // Outlined to avoid including TIterator.h
+         ~RIterStackElem();
+
+         // fDirPath doesn't need to be compared because it's implied by fIter.
+         bool operator==(const RIterStackElem &other) const { return fIter == other.fIter; }
+      };
+
+      std::vector<RIterStackElem> fIterStack;
+      Pattern_t fPattern;
+      const TKey *fCurKey = nullptr;
+      std::uint16_t fRootDirNesting = 0;
+      std::uint32_t fFlags = 0;
+
+      void Advance();
+
+      // NOTE: `iter` here is an owning pointer (or null)
+      RIterator(TIterator *iter, Pattern_t pattern, std::uint32_t flags);
+
+   public:
+      using iterator = RIterator;
+      using iterator_category = std::forward_iterator_tag;
+      using difference_type = std::ptrdiff_t;
+      using value_type = RFileKeyInfo;
+      using pointer = const value_type *;
+      using reference = const value_type &;
+
+      iterator &operator++()
+      {
+         Advance();
+         return *this;
+      }
+      value_type operator*();
+      bool operator!=(const iterator &rh) const { return !(*this == rh); }
+      bool operator==(const iterator &rh) const { return fIterStack == rh.fIterStack; }
+   };
+
+   RFileKeyIterable(TFile *file, std::string_view rootDir, std::uint32_t flags)
+      : fFile(file), fPattern(std::string(rootDir)), fFlags(flags)
+   {
+   }
+
+   RIterator begin() const;
+   RIterator end() const;
+};
+
 /**
 \class ROOT::Experimental::RFile
 \ingroup RFile
@@ -196,6 +285,31 @@ public:
 
    /// Flushes the RFile if needed and closes it, disallowing any further reading or writing.
    void Close();
+
+   /// Returns an iterable over all paths of objects written into this RFile starting at path "rootPath".
+   /// The returned paths are always "absolute" paths: they are not relative to `rootPath`.
+   /// Keys relative to directories are not returned: only those relative to leaf objects are.
+   /// If `rootPath` is the path of a leaf object, only `rootPath` itself will be returned.
+   /// This recurses on all the subdirectories of `rootPath`. If you only want the immediate children of `rootPath`,
+   /// use GetKeysNonRecursive().
+   RFileKeyIterable GetKeys(std::string_view rootPath = "") const
+   {
+      return RFileKeyIterable(fFile.get(), rootPath, RFileKeyIterable::kRecursive);
+   }
+
+   /// Returns an iterable over all paths of objects written into this RFile contained in the directory "rootPath".
+   /// The returned paths are always "absolute" paths: they are not relative to `rootPath`.
+   /// Keys relative to directories are not returned: only those relative to leaf objects are.
+   /// If `rootPath` is the path of a leaf object, only `rootPath` itself will be returned.
+   /// This only returns the immediate children of `rootPath`. If you want to recurse into the subdirectories of
+   /// `rootPath`, use GetKeys().
+   RFileKeyIterable GetKeysNonRecursive(std::string_view rootPath = "") const
+   {
+      return RFileKeyIterable(fFile.get(), rootPath, RFileKeyIterable::kNone);
+   }
+
+   /// Prints the internal structure of this RFile to the given stream.
+   void Print(std::ostream &out = std::cout) const;
 };
 
 } // namespace Experimental

io/io/src/RFile.cxx

@@ -13,7 +13,9 @@
 #include <Byteswap.h>
 #include <TError.h>
 #include <TFile.h>
+#include <TIterator.h>
 #include <TKey.h>
+#include <TList.h>
 #include <TROOT.h>
 
 #include <algorithm>
@@ -184,10 +186,6 @@ static std::string ValidateAndNormalizePath(std::string &path)
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-RFile::RFile(std::unique_ptr<TFile> file) : fFile(std::move(file)) {}
-
-RFile::~RFile() = default;
-
 std::unique_ptr<RFile> RFile::Open(std::string_view path)
 {
    CheckExtension(path);
@@ -227,6 +225,10 @@ std::unique_ptr<RFile> RFile::Recreate(std::string_view path)
    return rfile;
 }
 
+RFile::RFile(std::unique_ptr<TFile> file) : fFile(std::move(file)) {}
+
+RFile::~RFile() = default;
+
 TKey *RFile::GetTKey(std::string_view path) const
 {
    // In RFile, differently from TFile, when dealing with a path like "a/b/c", we always consider it to mean
@@ -373,6 +375,125 @@ void RFile::PutUntyped(std::string_view pathSV, const std::type_info &type, cons
    }
 }
 
+ROOT::Experimental::RFileKeyIterable::RIterator::RIterStackElem::RIterStackElem(TIterator *it, const std::string &path)
+   : fIter(it), fDirPath(path)
+{
+}
+
+ROOT::Experimental::RFileKeyIterable::RIterator::RIterStackElem::~RIterStackElem() = default;
+
+ROOT::Experimental::RFileKeyIterable::RIterator::RIterator(TIterator *iter, Pattern_t pattern, std::uint32_t flags)
+   : fPattern(pattern), fFlags(flags)
+{
+   if (iter) {
+      fIterStack.emplace_back(iter);
+
+      if (!pattern.empty()) {
+         fRootDirNesting = std::count(pattern.begin(), pattern.end(), '/');
+         // `pattern` may or may not end with '/', but we consider it a directory regardless.
+         // In other words, like in virtually all filesystem operations, "dir" and "dir/" are equivalent.
+         fRootDirNesting += pattern.back() != '/';
+      }
+
+      // Advance the iterator to skip the first key, which is always the TFile key.
+      // This will also skip keys until we reach the first correct key we want to return.
+      Advance();
+   }
+}
+
+ROOT::Experimental::RFileKeyIterable::RIterator ROOT::Experimental::RFileKeyIterable::begin() const
+{
+   return {fFile->GetListOfKeys()->MakeIterator(), fPattern, fFlags};
+}
+
+ROOT::Experimental::RFileKeyIterable::RIterator ROOT::Experimental::RFileKeyIterable::end() const
+{
+   return {nullptr, fPattern, fFlags};
+}
+
+void ROOT::Experimental::RFileKeyIterable::RIterator::Advance()
+{
+   fCurKey = nullptr;
+
+   const bool recursive = fFlags & kRecursive;
+
+   // We only want to return keys that refer to user objects, not internal ones, therefore we skip
+   // all keys that have internal class names.
+   while (!fIterStack.empty()) {
+      auto &[iter, dirPath] = fIterStack.back();
+      assert(iter);
+      TObject *keyObj = iter->Next();
+      if (!keyObj) {
+         // reached end of the iteration
+         fIterStack.pop_back();
+         continue;
+      }
+
+      assert(keyObj->IsA() == TClass::GetClass<TKey>());
+      auto key = static_cast<TKey *>(keyObj);
+
+      const auto dirSep = (dirPath.empty() ? "" : "/");
+
+      if (strcmp(key->GetClassName(), "TDirectory") == 0 || strcmp(key->GetClassName(), "TDirectoryFile") == 0) {
+         TDirectory *dir = key->ReadObject<TDirectory>();
+         TIterator *innerIter = dir->GetListOfKeys()->MakeIterator();
+         assert(innerIter);
+         fIterStack.emplace_back(innerIter, dirPath + dirSep + dir->GetName());
+         continue;
+      }
+
+      // Reconstruct the full path of the key
+      const auto &fullPath = dirPath + dirSep + key->GetName();
+      const auto nesting = fIterStack.size() - 1;
+
+      // skip key if it's not a child of root dir
+      if (!ROOT::StartsWith(fullPath, fPattern))
+         continue;
+
+      // check that we are in the same directory as "rootDir".
+      if (!recursive && nesting != fRootDirNesting)
+         continue;
+
+      // All checks passed: return this key.
+      assert(!fullPath.empty());
+      fCurKey = key;
+      break;
+   }
+}
+
+ROOT::Experimental::RFileKeyInfo ROOT::Experimental::RFileKeyIterable::RIterator::operator*()
+{
+   if (fIterStack.empty())
+      throw ROOT::RException(R__FAIL("tried to dereference an invalid iterator"));
+
+   const TKey *key = fCurKey;
+   if (!key)
+      throw ROOT::RException(R__FAIL("tried to dereference an invalid iterator"));
+
+   const auto &dirPath = fIterStack.back().fDirPath;
+
+   RFileKeyInfo keyInfo;
+   keyInfo.fName = dirPath + (dirPath.empty() ? "" : "/") + key->GetName();
+   keyInfo.fClassName = key->GetClassName();
+   keyInfo.fCycle = key->GetCycle();
+   keyInfo.fTitle = key->GetTitle();
+   return keyInfo;
+}
+
+void RFile::Print(std::ostream &out) const
+{
+   std::vector<RFileKeyInfo> keys;
+   auto keysIter = GetKeys();
+   for (const auto &key : keysIter) {
+      keys.emplace_back(key);
+   }
+
+   std::sort(keys.begin(), keys.end(), [](const auto &a, const auto &b) { return a.fName < b.fName; });
+   for (const auto &key : keys) {
+      out << key.fClassName << " " << key.fName << ";" << key.fCycle << ": \"" << key.fTitle << "\"\n";
+   }
+}
+
 size_t RFile::Flush()
 {
    return fFile->Write();