NTreeBranch.cxx

#include <TTree.h>
#include <TList.h>
#include <string>
#include <vector>
#include <TH1.h>
#include <THnSparse.h>
#include <TBranch.h>
#include <TObject.h>
#include <TROOT.h>
#include <TSystem.h>
#include "NLogger.h"
#include "NUtils.h"
 
#include "NTreeBranch.h"

ClassImp(Ndmspc::NTreeBranch);
 
namespace Ndmspc {
NTreeBranch::NTreeBranch(TTree * tree, const std::string & name, void * address, const std::string & objClassName)
    : TObject(), fName(name), fObjectClassName(objClassName)
{
  // SetAddress(address);
  if (tree) {
    Branch(tree, address);
  }
}
NTreeBranch::~NTreeBranch()
{
  // delete fObject;
  // delete fBranch;
}
TBranch * NTreeBranch::Branch(TTree * tree, void * address)
{
  if (!tree) {
    NLogError("Tree is nullptr !!!");
    return nullptr;
  }
 
  if (tree->GetBranch(fName.c_str())) {
    fBranch = tree->GetBranch(fName.c_str());
    return fBranch;
  }
  // fBranch = tree->Branch(fName.c_str(), fObjectClassName.c_str(), &address);
  tree->Branch(fName.c_str(), fObjectClassName.c_str(), &address);
  fBranch = tree->GetBranch(fName.c_str());
  return fBranch;
}
void NTreeBranch::SetAddress(void * address, bool deleteExisting)
{
  NLogTrace("NTreeBranch::SetAddress: Setting address %p for branch '%s' ...", address, fName.c_str());
 
  // if (fObject && deleteExisting) {
  if (fObject && deleteExisting) {
    NLogDebug("NTreeBranch::SetAddress: Deleting existing object %p for branch '%s' ...", fObject, fName.c_str());
    fObject->Delete(); // Delete existing object to avoid memory leaks when setting new address
    // NOTE: fObject->Delete() calls 'delete this' internally (TObject::Delete()),
    // so do NOT also call 'delete fObject' here — that would be a double-free.
    fObject = nullptr;
  }
  fObject = (TObject *)address;
  fBranch->SetAddress(&fObject);
}
void NTreeBranch::SetBranchAddress(TTree * tree)
{
 
  if (!tree) {
    NLogError("Tree is nullptr !!!");
    return;
  }
 
  NLogTrace("NTreeBranch::SetBranchAddress: Setting branch address '%s' ...", fName.c_str());
 
  if (fObject) {
    NLogDebug("NTreeBranch::SetBranchAddress: Deleting existing object %p for branch '%s' ...", fObject, fName.c_str());
    fObject->Delete(); // Delete existing object to avoid memory leaks when setting new address
    // NOTE: fObject->Delete() calls 'delete this' internally (TObject::Delete()),
    // so do NOT also call 'delete fObject' here — that would be a double-free.
    fObject = nullptr;
  }
  tree->SetBranchStatus(fName.c_str(), fBranchStatus);
  tree->SetBranchAddress(fName.c_str(), &fObject);
  fBranch = tree->GetBranch(fName.c_str());
}
 
Long64_t NTreeBranch::GetEntry(TTree * tree, Long64_t entry)
{
 
  if (tree == nullptr) {
    NLogError("Tree is not initialized !!!");
    return -1;
  }
 
  NLogTrace("Getting entry for branch '%s' %lld status=%d ...", fBranch->GetName(), entry,
            tree->GetBranchStatus(fBranch->GetName()));
 
  Long64_t bytes = 0;
  if (fBranch && tree->GetBranchStatus(fBranch->GetName()) == 1) {
    // Temporarily disable ROOT's recursive cleanup to prevent crashes when
    // TCanvas objects (stored in TList branches) are destroyed during
    // deserialization of a new entry (TCanvas::Streamer -> TObjArray::~TObjArray
    // -> RecursiveRemove).
    // Bool_t prevMustClean = gROOT->MustClean();
    // gROOT->SetMustClean(kFALSE);
 
    if (fObject && fObject->InheritsFrom(TList::Class())) {
      NLogTrace("NTreeBranch::GetEntry: Deleting existing object %p for branch '%s' ...", fObject, fName.c_str());
      // NUtils::SafeDeleteObject(fObject);
      TList * lst = static_cast<TList *>(fObject);
      NUtils::SafeDeleteTList(lst);
      fObject = nullptr;
    }
    bytes = fBranch->GetEntry(entry);
    // ProcInfo_t info;
    // gSystem->GetProcInfo(&info);
    // NLogDebug("NTreeBranch::GetEntry:[RSS]: %ld kB", info.fMemResident);
 
    // gROOT->SetMustClean(prevMustClean);
    NLogTrace("NTreeBranch::GetEntry: Getting content from %s with size %.3f MB", fBranch->GetName(),
              (double)bytes / (1024 * 1024));
  }
 
  return bytes;
}
void NTreeBranch::SaveEntry(NTreeBranch * hnstBranchIn, bool useProjection, const std::string projOpt)
{
  NLogTrace("Saving entry for branch=%s ...", fName.c_str());
 
  TString classNameStr = hnstBranchIn->GetObjectClassName().c_str();
  // NLogDebug("NTreeBranch::SaveEntry: Obj class name %s ...", classNameStr.Data());
  if (classNameStr.BeginsWith("THnSparse")) {
 
    THnSparse * in = (THnSparse *)hnstBranchIn->GetObject();
    if (in) {
      if (in->GetNdimensions() > 0) {
        if (useProjection) {
          // in->Print();
          auto dims = std::make_unique<Int_t[]>(in->GetNdimensions());
          // Int_t dims[in->GetNdimensions()];
          for (Int_t iDim = 0; iDim < in->GetNdimensions(); iDim++) {
            dims[iDim] = iDim;
          }
          THnSparse * out = (THnSparse *)in->ProjectionND(in->GetNdimensions(), dims.get(), projOpt.c_str());
          // Loop over all bins
          double sum = 0;
          NLogTrace("Projection of %s with filled bins %lld ...", in->GetName(), out->GetNbins());
          for (Int_t i = 0; i < out->GetNbins(); i++) {
            NLogTrace("Bin %d content=%f", i, out->GetBinContent(i));
            sum += out->GetBinContent(i);
          }
          // out->Projection(0)->Print();
          out->SetEntries(sum);
          out->SetNameTitle(in->GetName(), in->GetTitle());
          // out->Print();
          SetAddress(out);
        }
        else {
          SetAddress(in);
        }
      }
    }
    else {
      SetAddress(nullptr);
    }
  }
  else {
    NLogTrace("Class '%s' is stored default method !!!", classNameStr.Data());
    // return;
  }
}
 
void NTreeBranch::Print(Option_t * /*option*/) const
{
  NLogInfo("Branch name='%s' objClassName='%s' address=%p branch=%p status=%d", fName.c_str(), fObjectClassName.c_str(),
           fObject, fBranch, fBranchStatus);
}
 
} // namespace Ndmspc