RootPersistenceService.h

Go to the documentation of this file.

/*
 * RootPersistenceService.h
 *
 *  Created on: 12 May 2017
 *      Author: Nicola Mori
 */

#ifndef ROOTPERSISTENCESERVICE_H_
#define ROOTPERSISTENCESERVICE_H_

#include "persistence/PersistenceService.h"

class TFile;
class TObject;
class TTree;
class TBranch;
class TClass;
class RootOutputStruct;

#include <set>
#include <unordered_map>

namespace EA {

class RootPersistenceService : public PersistenceService {
public:
  RootPersistenceService(const std::string &name, const std::string &output);

  bool Connect() override;
  bool Disconnect() override;

  bool BookEventObject(const std::string &objName, const std::string &objStore) override;
  bool BookPassObject(const std::string &objName, const std::string &objStore) override;
  bool BookGlobalObject(const std::string &objName, const std::string &objStore) override;

  std::vector<BookedObject> GetBookedEventObjects() override;

  std::vector<BookedObject> GetBookedPassObjects() override;

  std::vector<BookedObject> GetBookedGlobalObjects() override;

  bool BeginOfEvent() override;
  bool EndOfEvent() override;
  bool BeginOfPass() override;
  bool EndOfPass() override;
  bool BeginOfProcessing() override;
  bool EndOfProcessing() override;

  std::unique_ptr<PersistenceService> Copy(const std::string &name, const std::string &output) override;

protected:
  observer_ptr<TFile> GetOutputFile() { return _outputFile; }

  observer_ptr<TTree> GetEventTree() { return _eventTree; }

private:
  /* ******* Bookkeeping structs for output data objects ******* */

  // Base
  struct ObjInfo {
    std::string dataStoreName;
    StorePtr dataStore;
    bool isGlob;              // True if the object name is a glob expression
    bool isMatched;           // True if the glob has been matched
    bool isAlias;             // True if the object name is an alias
    bool aliasToBeChecked;    // Set to false after checking if the object name is an alias
    TBranch *nameForEvBranch; // Branch storing the name of aliased global object for each event
    std::string nameForThisEv;
    std::set<std::string> actualNames; // Set of all the names assumed by the aliased object

    ObjInfo(const std::string &storeName)
        : dataStoreName(storeName), dataStore(nullptr), isGlob(false), isMatched(false), isAlias(false),
          aliasToBeChecked(true), nameForEvBranch(nullptr) {}
  };

  // Event
  struct EvObjInfo;
  using EvObjectsBook = std::unordered_map<std::string, EvObjInfo>;
  struct EvObjInfo : public ObjInfo {
    TBranch *outputBranch; // Output branch
    bool flag;             // True if the object is present for current event
    void *object;
    void *defaultObject;
    TClass *objClass;
    void *prevObject; // Address used for previous fill

    std::unique_ptr<EvObjectsBook> aliasedObjsBook; // Sub-book for aliased objects

    EvObjInfo(const std::string &storeName)
        : ObjInfo(storeName), outputBranch(nullptr), flag(true), object(nullptr), defaultObject(nullptr),
          objClass(nullptr), prevObject(nullptr), aliasedObjsBook(new EvObjectsBook) {}

    EvObjInfo(const EvObjInfo &otherObjInfo)
        : ObjInfo(otherObjInfo), outputBranch(otherObjInfo.outputBranch), flag(otherObjInfo.flag),
          object(otherObjInfo.object), defaultObject(otherObjInfo.defaultObject), objClass(otherObjInfo.objClass),
          prevObject(otherObjInfo.prevObject), aliasedObjsBook(new EvObjectsBook(*(otherObjInfo.aliasedObjsBook))) {}
  };

  EvObjectsBook _eventBook;

  // Global
  struct GlobObjInfo : public ObjInfo {
    bool saved; // Flag for saved global objects

    GlobObjInfo(const std::string &storeName) : ObjInfo(storeName), saved(false) {}
  };
  using GlobObjectsBook = std::unordered_map<std::string, GlobObjInfo>;
  GlobObjectsBook _globalBook;

  /* ******* Pointers to Root output objects ******* */
  TFile *_outputFile;
  TTree *_eventTree;

  template <typename T>
  bool _ResolveGlobAndBook(std::unordered_map<std::string, T> &objBook, ObjectCategory objCategory);
  bool _EndOfEventForEventObject(const std::string &objName, EvObjInfo &objInfo);
  bool _EndOfEventForGlobalObject(const std::string &objName, GlobObjInfo &objInfo);
  bool _SaveGlobalObject(const std::string &objName, GlobObjInfo &objInfo);
};

template <typename T>
bool RootPersistenceService::_ResolveGlobAndBook(std::unordered_map<std::string, T> &objBook,
                                                 ObjectCategory objCategory) {
  const std::string routineName("RootPersistenceService::_ResolveGlobAndBook");
  typename std::remove_reference<decltype(objBook)>::type newBook;
  for (auto &objInfoElem : objBook) {
    const std::string &objName = objInfoElem.first;
    T &objInfo = objInfoElem.second;
    if (!(objInfo.dataStore)) {
      observer_ptr<DataStoreManager> dsManager = GetDataStoreManager();
      try {
        if (objCategory == ObjectCategory::GLOBAL) {
          objInfo.dataStore = dsManager->GetGlobalDataStore(objInfo.dataStoreName);
        } else {
          objInfo.dataStore = dsManager->GetEventDataStore(objInfo.dataStoreName);
        }
      } catch (Retrieval::NotFound &exc) {
        COUT(ERROR) << objCategory << " object store " << objInfo.dataStoreName << " not found." << ENDL;
        return false;
      }
    }

    if (objInfo.isGlob) {
      // The name is a glob expression. Search all the matching names and aliases and book each of them separately
      std::regex regEx(StringUtils::RegexFromGlob(objName));
      std::smatch match;
      auto storeKnownObjs = objInfo.dataStore->GetKnownObjects();
      for (auto knownObj : storeKnownObjs) {
        if (std::regex_match(knownObj.name, match, regEx)) {
          objInfo.isMatched = true;
          auto insertionResult = newBook.insert(std::pair<std::string, T>(knownObj.name, objInfo));
          if (!(insertionResult.second)) {
            COUT(ERROR) << "Can't insert the " << std::nouppercase << objCategory << " object name " << knownObj.name
                        << " matching the booked glob " << objName << " among the booked objects" << ENDL;
            return false;
          }
          insertionResult.first->second.isGlob = false;
          insertionResult.first->second.isAlias = false;
          insertionResult.first->second.aliasToBeChecked = false;
        }
        for (auto &alias : knownObj.aliases) {
          if (std::regex_match(alias, match, regEx)) {
            objInfo.isMatched = true;
            auto insertionResult = newBook.insert(std::pair<std::string, T>(alias, objInfo));
            if (!(insertionResult.second)) {
              COUT(ERROR) << "Can't insert the " << objCategory << " object alias " << alias
                          << " matching the booked glob " << objName << " among the booked objects" << ENDL;
              return false;
            }
            insertionResult.first->second.isGlob = false;
            insertionResult.first->second.isAlias = true;
            insertionResult.first->second.aliasToBeChecked = false;
          }
        }
      }
    }

    // Keep also the original entry (even if it is a glob in order to save unknown objects with matching name, will be
    // handled in _Save{Global,Event}Object)
    auto insertionResult = newBook.insert(objInfoElem);
    if (!(insertionResult.second)) {
      COUT(ERROR) << "Can't insert the " << objCategory << " object " << objInfoElem.first
                  << " among the booked objects" << ENDL;
      return false;
    }
  }

  objBook.swap(newBook);
  return true;
}

} // namespace EA

#endif /* ROOTPERSISTENCESERVICE_H_ */