SparseTree.h

Go to the documentation of this file.

/*
 * SparseTree.h
 *
 *  Created on: 7 Aug 2019
 *      Author: Nicola Mori
 */

#ifndef SPARSETREE_H_
#define SPARSETREE_H_

// Root headers
#include "TLeaf.h"
#include "TTree.h"

// C/C++ headers
#include <functional>

namespace EA {

class SparseTree : public TTree {
public:
  using TTree::TTree;

  template <typename T> Int_t SetBranchAddress(const char *bname, T **add, TBranch **ptr = 0);

  ClassDef(SparseTree, 1)
};

template <typename T> Int_t SparseTree::SetBranchAddress(const char *bname, T **add, TBranch **ptr) {

  // Deactivate all branches on 1st call (i.e. when all the branches are active)
  auto *branches = GetListOfBranches();
  bool firstCall = true;
  for (int iBranch = 0; iBranch < branches->GetEntries(); ++iBranch) {
    auto *branch = (TBranch *)(branches->UncheckedAt(iBranch));
    if (branch->TestBit(kDoNotProcess)) {
      firstCall = false;
      break;
    }
  }
  if (firstCall) {
    SetBranchStatus("*", 0);
  }

  auto *branch = GetBranch(bname);
  if (branch) {
    // Recursive lambda that activates a branch and all its sub-branches and leaves
    std::function<void(TBranch *)> ActivateBranch = [&ActivateBranch](TBranch *branch) {
      auto *subBranches = branch->GetListOfBranches();
      for (int iBranch = 0; iBranch < subBranches->GetEntries(); ++iBranch) {
        ActivateBranch((TBranch *)(subBranches->UncheckedAt(iBranch)));
      }
      branch->SetStatus(1);

      auto *leaves = branch->GetListOfLeaves();
      for (int iLeaf = 0; iLeaf < leaves->GetEntries(); ++iLeaf) {
        auto *leaf = (TLeaf *)(leaves->UncheckedAt(iLeaf));
        leaf->GetBranch()->SetStatus(1);
        auto *leafCount = leaf->GetLeafCount();
        if (leafCount) {
          leafCount->GetBranch()->SetStatus(1);
        }
      }
    };

    // Activate the branch and its sub-branches
    ActivateBranch(branch);

    // Set the branch address
    return TTree::SetBranchAddress<T>(bname, add, ptr);
  }
  return 0;
}

} // namespace EA

#endif /* SPARSETREE_H_ */