SkimTreeHandle.h

Go to the documentation of this file.

#ifndef NAIA_SKIMTREEHANDLE_H
#define NAIA_SKIMTREEHANDLE_H

#include <memory>

#include "TFile.h"

namespace NAIA {

template <class T> class SkimTreeHandle {
public:
  SkimTreeHandle(std::unique_ptr<TFile> outFile, std::unique_ptr<T> chain, T *originalChain)
      : m_outFile{std::move(outFile)}, m_chain{std::move(chain)}, m_originalChain{originalChain} {};

  int Fill() {
    m_originalChain->m_event.ForceReadAllBranches();
    return m_chain->Fill();
  }

  int Write() { // NOTE: We copy the whole RTI and FileInfo tree in the output chain
    NAIA::Version versionHeader;
    m_outFile->WriteTObject(&versionHeader, "VersionHeader");

    // we need raw pointers cause SetBranchAddress expects the address to a pointer

    if (!m_originalChain->IsMC()) {
      auto *rtiInfo = new NAIA::RTIInfo();
      m_originalChain->GetRTITree()->SetBranchAddress("RTIInfo", &rtiInfo);
      m_chain->m_rti.tree->Branch("RTIInfo", rtiInfo);

      for (int ii = 0; ii < m_originalChain->GetRTITree()->GetEntries(); ii++) {
        m_originalChain->GetRTITree()->GetEntry(ii);
        m_chain->FillRTI();
      }
      m_chain->m_rti.tree->Write();

      delete rtiInfo;
    }

    auto *fileInfo = new NAIA::FileInfo();
    auto *fileInfoMC = new NAIA::MCFileInfo();
    m_originalChain->GetFileInfoTree()->SetBranchAddress("FileInfo", &fileInfo);
    m_chain->m_file.tree->Branch("FileInfo", fileInfo);
    if (m_originalChain->IsMC()) {
      m_originalChain->GetFileInfoTree()->SetBranchAddress("MCFileInfo", &fileInfoMC);
      m_chain->m_file.tree->Branch("MCFileInfo", fileInfoMC);
    }

    for (int ii = 0; ii < m_originalChain->GetFileInfoTree()->GetEntries(); ii++) {
      m_originalChain->GetFileInfoTree()->GetEntry(ii);
      m_chain->FillFileInfo();
    }
    m_chain->m_file.tree->Write();

    delete fileInfo;

    return m_chain->Write();
  }

  TFile *RootFile() { return m_outFile.get(); }

private:
  std::unique_ptr<TFile> m_outFile;
  std::unique_ptr<T> m_chain;
  T *m_originalChain;
};
} // namespace NAIA

#endif // NAIA_SKIMTREEHANDLE_H