RootTreeReader.h

Go to the documentation of this file.

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
#ifndef FRAMEWORK_ROOTTREEREADER_H
#define FRAMEWORK_ROOTTREEREADER_H
 
 
#include "Framework/RootSerializationSupport.h"
#include "Framework/Output.h"
#include "Framework/ProcessingContext.h"
#include "Framework/DataAllocator.h"
#include "Framework/RootMessageContext.h"
#include "Framework/Logger.h"
#include "Headers/DataHeader.h"
#include <TChain.h>
#include <TTree.h>
#include <TBranch.h>
#include <TClass.h>
#include <vector>
#include <string>
#include <stdexcept> // std::runtime_error
#include <type_traits>
#include <memory>     // std::make_unique
#include <functional> // std::function
#include <utility>    // std::forward
 
namespace o2::framework
{
 
namespace rtr
{
// The class 'Output' has a unique header stack member which makes the class not
// usable as key, in particular there is no copy constructor. Thats why a dedicated
// structure is used which can be transformed into Output. Moreover it also handles the
// conversion from OutputStack to Output
struct DefaultKey {
  header::DataOrigin origin;
  header::DataDescription description;
  header::DataHeader::SubSpecificationType subSpec = 0;
  enum Lifetime lifetime = Lifetime::Timeframe;
 
  DefaultKey(const Output& desc)
    : origin(desc.origin), description(desc.description), subSpec(desc.subSpec)
  {
  }
 
  operator Output() const { return Output{origin, description, subSpec}; }
};
} // namespace rtr
 
template <typename KeyType>
class GenericRootTreeReader
{
 public:
  using self_type = GenericRootTreeReader<KeyType>;
  using key_type = KeyType;
  // The RootTreeWriter/Reader support storage of binary data as a vector of char.
  // This allows to store e.g. raw data or some intermediate binary data alongside
  // with ROOT objects like the MC labels.
  using BinaryDataStoreType = std::vector<char>;
 
  enum struct PublishingMode {
    Single,
    Loop,
  };
 
  struct SpecialPublishHook {
    std::function<bool(std::string_view name, ProcessingContext& context, Output const&, char* data)> hook;
  };
 
  // the key must not be of type const char* to make sure that the variable argument
  // list of the constructor can be parsed
  static_assert(std::is_same<KeyType, const char*>::value == false, "the key type must not be const char*");
 
  struct ConstructorArg {
    ConstructorArg(key_type _key, const char* _name)
      : key(_key), name(_name) {}
    ConstructorArg(key_type _key, std::string const& _name)
      : key(_key), name(_name) {}
    key_type key;
    std::string name;
  };
 
  // the vector of arguments is filled during the build up of the branch configuration and
  // passed to the construction of the constructed mixin class
  using ConstructorArgs = std::vector<ConstructorArg>;
 
  class BranchConfigurationInterface
  {
   public:
    static const size_t STAGE = 0;
    BranchConfigurationInterface() = default;
    BranchConfigurationInterface(ConstructorArgs const&){};
    virtual ~BranchConfigurationInterface() = default;
 
    virtual void setup(TTree&, SpecialPublishHook* h = nullptr) {}
    virtual void exec(ProcessingContext& ctx, int entry, std::function<o2::header::Stack()> stackcreator) {}
 
   private:
  };
 
  template <typename DataT, typename BASE>
  class BranchConfigurationElement : public BASE
  {
   public:
    using PrevT = BASE;
    using value_type = DataT;
    using publish_type = void;
    static const size_t STAGE = BASE::STAGE + 1;
    BranchConfigurationElement() = default;
    BranchConfigurationElement(ConstructorArgs const& args)
      : PrevT(args), mKey(args[STAGE - 1].key), mName(args[STAGE - 1].name)
    {
    }
    ~BranchConfigurationElement() override = default;
 
    void exec(ProcessingContext& ctx, int entry, std::function<o2::header::Stack()> stackcreator) override
    {
      process(ctx, entry, stackcreator);
    }
 
    void setup(TTree& tree, SpecialPublishHook* publishhook = nullptr) override
    {
      setupInstance(tree, publishhook);
    }
 
    void setupInstance(TTree& tree, SpecialPublishHook* publishhook = nullptr)
    {
      // recursing through the tree structure by simply using method of the previous type,
      // i.e. the base class method.
      if constexpr (STAGE > 1) {
        PrevT::setupInstance(tree, publishhook);
      }
      mPublishHook = publishhook;
 
      // right now we allow the same key to appear for multiple branches
      mBranch = tree.GetBranch(mName.c_str());
      if (mBranch) {
        std::string sizebranchName = std::string(mName) + "Size";
        auto sizebranch = tree.GetBranch(sizebranchName.c_str());
        auto* classinfo = TClass::GetClass(mBranch->GetClassName());
        if constexpr (not std::is_void<value_type>::value) {
          // check if the configured type matches the stored type
          auto* storedclass = TClass::GetClass(typeid(value_type));
          if (classinfo != storedclass) {
            throw std::runtime_error(std::string("Configured type ") +
                                     (storedclass != nullptr ? storedclass->GetName() : typeid(value_type).name()) +
                                     " does not match the stored data type " +
                                     (classinfo != nullptr ? classinfo->GetName() : "") +
                                     " in branch " + mName);
          }
        }
        if (!sizebranch) {
          if (classinfo == nullptr) {
            throw std::runtime_error(std::string("can not find class description for branch ") + mName);
          }
          LOG(info) << "branch set up: " << mName;
        } else {
          if (classinfo == nullptr || classinfo != TClass::GetClass(typeid(BinaryDataStoreType))) {
            throw std::runtime_error("mismatching class type, expecting std::vector<char> for binary branch");
          }
          mSizeBranch = sizebranch;
          LOG(info) << "binary branch set up: " << mName;
        }
        mClassInfo = classinfo;
      } else {
        throw std::runtime_error(std::string("can not find branch ") + mName);
      }
    }
 
    void process(ProcessingContext& context, int entry, std::function<o2::header::Stack()>& stackcreator)
    {
      // recursing through the tree structure by simply using method of the previous type,
      // i.e. the base class method.
      if constexpr (STAGE > 1) {
        PrevT::process(context, entry, stackcreator);
      }
 
      auto snapshot = [&context, &stackcreator](const KeyType& key, const auto& object) {
        context.outputs().snapshot(Output{key.origin, key.description, key.subSpec, std::move(stackcreator())}, object);
      };
 
      char* data = nullptr;
      mBranch->SetAddress(&data);
      mBranch->GetEntry(entry);
 
      // execute hook if it was registered; if this return true do not proceed further
      if (mPublishHook != nullptr && (*mPublishHook).hook(mName, context, Output{mKey.origin, mKey.description, mKey.subSpec, std::move(stackcreator())}, data)) {
 
      }
      // try to figureout when we need to do something special
      else {
        if (mSizeBranch != nullptr) {
          size_t datasize = 0;
          mSizeBranch->SetAddress(&datasize);
          mSizeBranch->GetEntry(entry);
          auto* buffer = reinterpret_cast<BinaryDataStoreType*>(data);
          if (buffer->size() == datasize) {
            LOG(debug) << "branch " << mName << ": publishing binary chunk of " << datasize << " bytes(s)";
            snapshot(mKey, std::move(*buffer));
          } else {
            LOG(error) << "branch " << mName << ": inconsitent size of binary chunk "
                       << buffer->size() << " vs " << datasize;
            BinaryDataStoreType empty;
            snapshot(mKey, empty);
          }
        } else {
          if constexpr (std::is_void<value_type>::value == true) {
            // the default branch configuration publishes the object ROOT serialized
            snapshot(mKey, std::move(ROOTSerializedByClass(*data, mClassInfo)));
          } else {
            // if type is specified in the branch configuration, the allocator API decides
            // upon serialization
            snapshot(mKey, *reinterpret_cast<value_type*>(data));
          }
        }
      }
      // cleanup the memory
      auto* delfunc = mClassInfo->GetDelete();
      if (delfunc) {
        (*delfunc)(data);
      }
      mBranch->DropBaskets("all");
    }
 
   private:
    key_type mKey;
    std::string mName;
    TBranch* mBranch = nullptr;
    TBranch* mSizeBranch = nullptr;
    TClass* mClassInfo = nullptr;
    SpecialPublishHook* mPublishHook = nullptr;
  };
 
  template <typename T>
  struct BranchDefinition {
    using type = T;
    template <typename U>
    BranchDefinition(U _key, const char* _name)
      : key(_key), name(_name)
    {
    }
    template <typename U>
    BranchDefinition(U _key, std::string const& _name)
      : key(_key), name(_name)
    {
    }
    template <typename U>
    BranchDefinition(U _key, std::string&& _name)
      : key(_key), name(std::move(_name))
    {
    }
 
    key_type key;
    std::string name;
  };
 
  GenericRootTreeReader();
 
  template <typename... Args>
  GenericRootTreeReader(const char* treename, // name of the tree to read from
                        Args&&... args)       // file names, followed by branch info
    : mInput(treename)
  {
    mInput.SetCacheSize(0);
    parseConstructorArgs<0>(std::forward<Args>(args)...);
    mBranchConfiguration->setup(mInput, mPublishHook);
  }
 
  void addFile(const char* fileName)
  {
    mInput.AddFile(fileName);
    mNEntries = mInput.GetEntries();
  }
 
  bool next()
  {
    if ((mReadEntry + 1) >= mNEntries || mNEntries == 0) {
      if (mPublishingMode == PublishingMode::Single) {
        // stop here
        if (mReadEntry < mNEntries) {
          mReadEntry = mNEntries;
        }
        return false;
      }
      // start over in loop mode
      mReadEntry = -1;
    }
    if (mMaxEntries > 0 && (mNofPublished + 1) >= mMaxEntries) {
      if (mReadEntry < mNEntries) {
        mReadEntry = mNEntries;
      }
      return false;
    }
    ++mReadEntry;
    ++mNofPublished;
    return true;
  }
 
  self_type& operator++()
  {
    next();
    return *this;
  }
  self_type& operator++(int) = delete;
 
  using ROOTSerializedByClass = o2::framework::ROOTSerialized<char, TClass>;
 
  template <typename ContextType, typename... HeaderTypes>
  bool operator()(ContextType& context,
                  HeaderTypes&&... headers) const
  {
    if (mReadEntry >= mNEntries || mNEntries == 0 || (mMaxEntries > 0 && mNofPublished >= mMaxEntries)) {
      return false;
    }
 
    auto stackcreator = [&headers...]() {
      return o2::header::Stack{std::forward<HeaderTypes>(headers)...};
    };
 
    mBranchConfiguration->exec(context, mReadEntry, stackcreator);
    return true;
  }
 
  int getCount() const
  {
    return mNofPublished + 1;
  }
 
 private:
  // special helper to get the char argument from the argument pack
  template <typename T, typename... Args>
  const char* getCharArg(T arg, Args&&...)
  {
    static_assert(std::is_same<T, const char*>::value, "missing branch name after publishing key, use const char* argument right after the key");
    return arg;
  }
 
  template <size_t skip, typename U, typename... Args>
  void parseConstructorArgs(U key, Args&&... args)
  {
    // all argument parsing is done in the creation of the branch configuration
    mBranchConfiguration = createBranchConfiguration<0, BranchConfigurationInterface>({}, std::forward<U>(key), std::forward<Args>(args)...);
  }
 
  template <size_t skip, typename BASE, typename U, typename... Args>
  std::unique_ptr<BranchConfigurationInterface> createBranchConfiguration(ConstructorArgs&& cargs, U def, Args&&... args)
  {
    if constexpr (skip > 0) {
      return createBranchConfiguration<skip - 1, BASE>(std::move(cargs), std::forward<Args>(args)...);
    } else if constexpr (std::is_same<U, const char*>::value) {
      addFile(def);
    } else if constexpr (std::is_same<U, int>::value) {
      mMaxEntries = def;
    } else if constexpr (std::is_same<U, PublishingMode>::value) {
      mPublishingMode = def;
    } else if constexpr (std::is_same<U, SpecialPublishHook*>::value) {
      mPublishHook = def;
    } else if constexpr (is_specialization_v<U, BranchDefinition>) {
      cargs.emplace_back(key_type(def.key), def.name);
      using type = BranchConfigurationElement<typename U::type, BASE>;
      return std::move(createBranchConfiguration<0, type>(std::move(cargs), std::forward<Args>(args)...));
    } else if constexpr (sizeof...(Args) > 0) {
      const char* arg = getCharArg(std::forward<Args>(args)...);
      if (arg != nullptr && *arg != 0) {
        cargs.emplace_back(key_type(def), arg);
        using type = BranchConfigurationElement<void, BASE>;
        return std::move(createBranchConfiguration<1, type>(std::move(cargs), std::forward<Args>(args)...));
      }
      throw std::runtime_error("expecting valid branch name string after key");
    } else {
      static_assert(always_static_assert<U>::value, "argument mismatch, define branches either as argument pairs of key and branchname or using the BranchDefinition helper struct");
    }
    return createBranchConfiguration<0, BASE>(std::move(cargs), std::forward<Args>(args)...);
  }
 
  template <size_t skip, typename T>
  std::unique_ptr<BranchConfigurationInterface> createBranchConfiguration(ConstructorArgs&& cargs)
  {
    static_assert(skip == 0);
    return std::move(std::make_unique<T>(cargs));
  }
 
  TChain mInput;
  std::unique_ptr<BranchConfigurationInterface> mBranchConfiguration;
  int mNEntries = 0;
  int mReadEntry = -1;
  int mNofPublished = -1;
  int mMaxEntries = -1;
  PublishingMode mPublishingMode = PublishingMode::Single;
  SpecialPublishHook* mPublishHook = nullptr;
};
 
using RootTreeReader = GenericRootTreeReader<rtr::DefaultKey>;
 
} // namespace o2::framework
#endif