TreeBranchHelpers.h

Go to the documentation of this file.

// Dear emacs, this is -*- c++ -*-
//
// Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
//
#ifndef ASGANALYSISALGORITHMS_TREEBRANCHHELPERS_H
#define ASGANALYSISALGORITHMS_TREEBRANCHHELPERS_H
// System include(s):
#include <unordered_map>
#include <unordered_set>
#include <set>
#include <optional>
#include <string>
#include <vector>
#include <memory>
 
// Framework include(s):
#include "AsgMessaging/AsgMessaging.h"
#include "AsgMessaging/AsgMessagingForward.h"
#include "AsgServices/ServiceHandle.h"
#include "AnaAlgorithm/AnaAlgorithm.h"
#include "CxxUtils/checker_macros.h"
#include "SystematicsHandles/SysListHandle.h"
#include <AsgTools/PropertyWrapper.h>
 
// ROOT include(s):
#include <ROOT/RNTupleModel.hxx>
 
// EDM include(s):
#include "AthContainersInterfaces/IAuxTypeVector.h"
#include "AthContainers/AuxElement.h"
 
// Forward declaration(s):
class TClass;
class TTree;
class TVirtualCollectionProxy;
namespace SG {
   class AuxVectorBase;
   class IAuxTypeVectorFactory;
}
 
namespace CP
{
 
  namespace TreeBranchHelpers
  {
#ifdef XAOD_STANDALONE
    using StoreType = asg::SgEvent;
#else
    using StoreType = StoreGateSvc;
#endif
 
 

    struct BranchConfig
    {
      std::string branchDecl;

      std::string sgName;

      std::string auxName;

      std::string branchName;

      std::string typeName;

      bool nominalOnly = false;

      std::optional<int> basketSize;

      std::string metTermName;
 

      SG::auxid_t nominalAuxId = SG::null_auxid;

      const std::type_info* auxType = nullptr;

      const std::type_info* auxVecType = nullptr;

      const SG::IAuxTypeVectorFactory* auxFactory = nullptr;
 

      CP::SystematicSet sgNameFilterSys;

      CP::SystematicSet auxNameFilterSys;

      CP::SystematicSet branchNameFilterSys;
 

      StatusCode parse (const std::string& branchDecl, MsgStream& msg);

      StatusCode configureTypes (std::set<std::string>& decosWithoutType, MsgStream& msg);

      StatusCode configureSystematics (ISystematicsSvc& sysSvc, MsgStream& msg);
    };
 
 

    struct OutputBranchData
    {
      const BranchConfig *branchConfig = nullptr;

      std::size_t sysIndex = 0u;

      bool isNominal = false;

      std::string sgName;

      std::string auxName;

      std::string branchName;
 

      StatusCode configureNames (const BranchConfig& branchConfig, const CP::SystematicSet& sys, ISystematicsSvc& sysSvc, MsgStream& msg);
    };
 
 

    class IComponentProcessor
    {
    public:
      virtual ~IComponentProcessor() = default;
 
      virtual StatusCode setup ( ROOT::RNTupleModel& /* model */, const BranchConfig& /* branchConfig */, OutputBranchData& /* outputData */, MsgStream& msg ) = 0;
 
      virtual StatusCode setup ( TTree& /* tree */, const BranchConfig& /* branchConfig */, OutputBranchData& /* outputData */, MsgStream& msg ) = 0;
    };
 
 

    class IObjectProcessor
    {
    public:
      virtual ~IObjectProcessor() = default;

      virtual StatusCode retrieveProcess (StoreType& evtStore) = 0;

 
      virtual StatusCode addBranch(TTree&, const BranchConfig&, OutputBranchData&) = 0;
 
      virtual StatusCode addBranch(ROOT::RNTupleModel&, const BranchConfig&, OutputBranchData&) = 0;
    };

    class ElementBranchProcessor : public IComponentProcessor {
 
    public:
      ElementBranchProcessor() = default;
      virtual ~ElementBranchProcessor() = default;
      ElementBranchProcessor( const ElementBranchProcessor& ) = delete;
      ElementBranchProcessor& operator=( const ElementBranchProcessor& ) = delete;
 
 
      virtual StatusCode setup( TTree& tree, const BranchConfig& branchConfig, OutputBranchData& outputData, MsgStream& msg ) override;
 
      virtual StatusCode setup( ROOT::RNTupleModel& /* model */, const BranchConfig& /* branchConfig */, OutputBranchData& /* outputData */, MsgStream& msg ) override;
      StatusCode process( const SG::AuxElement& element,
                          MsgStream& msg );
 

      std::string m_branchName;
      std::unique_ptr< SG::TypelessConstAccessor > m_acc;
      const SG::IAuxTypeVectorFactory* m_factory = nullptr;
      std::unique_ptr< SG::IAuxTypeVector > m_data;
      void* m_dataPtr = nullptr;
 
    }; // class ElementBranchProcessor
 
 

    class ContainerBranchProcessor : public IComponentProcessor {
 
    public:
      ContainerBranchProcessor() = default;
      virtual ~ContainerBranchProcessor() = default;
      ContainerBranchProcessor( const ContainerBranchProcessor& ) = delete;
      ContainerBranchProcessor& operator=( const ContainerBranchProcessor& ) = delete;
      virtual StatusCode setup( TTree& tree, const BranchConfig& branchConfig, OutputBranchData& outputData, MsgStream& msg ) override;
      StatusCode resize( size_t size, MsgStream& msg );
      StatusCode process( const SG::AuxElement& element, size_t index,
                          MsgStream& msg );
      virtual StatusCode setup( ROOT::RNTupleModel& /* model */, const BranchConfig& /* branchConfig */, OutputBranchData& /* outputData */, MsgStream& msg ) override;

      std::string m_branchName;
      std::unique_ptr< SG::TypelessConstAccessor > m_acc;
      const SG::IAuxTypeVectorFactory* m_factory = nullptr;
      std::unique_ptr< SG::IAuxTypeVector > m_data;
      void* m_dataPtr = nullptr;
 
    }; // class ContainerBranchProcessor
 
 
 

    class ElementProcessorRegular : public asg::AsgMessaging, public IObjectProcessor {
 
    public:
      ElementProcessorRegular(const std::string& sgName);
      virtual ~ElementProcessorRegular() = default;
      ElementProcessorRegular( const ElementProcessorRegular& ) = delete;
      ElementProcessorRegular& operator=( const ElementProcessorRegular& ) = delete;

      virtual StatusCode retrieveProcess (StoreType& evtStore) override;

      virtual StatusCode addBranch( TTree& tree, const BranchConfig& branchConfig, OutputBranchData& outputData ) override;
 
      virtual StatusCode addBranch( ROOT::RNTupleModel& model, const BranchConfig& branchConfig, OutputBranchData& outputData ) override;
 
    private:

      std::vector<std::unique_ptr<ElementBranchProcessor>> m_branches;

      std::string m_sgName;
 
    }; // class ElementProcessorRegular
 

    class ATLAS_NOT_THREAD_SAFE ContainerProcessorRegular : public asg::AsgMessaging, public IObjectProcessor {
 
    public:
      ContainerProcessorRegular(const std::string& sgName);
      virtual ~ContainerProcessorRegular() = default;
      ContainerProcessorRegular( const ContainerProcessorRegular& ) = delete;
      ContainerProcessorRegular& operator=( const ContainerProcessorRegular& ) = delete;

      virtual StatusCode retrieveProcess (StoreType& evtStore) override;

      virtual StatusCode addBranch( TTree& tree, const BranchConfig& branchConfig, OutputBranchData& outputData ) override;
 
      virtual StatusCode addBranch( ROOT::RNTupleModel& model, const BranchConfig& branchConfig, OutputBranchData& outputData ) override;
 
    private:
      std::vector<std::unique_ptr<ContainerBranchProcessor>> m_branches;
      TVirtualCollectionProxy* m_collProxy = nullptr;
      int m_auxElementOffset = -1;

      std::string m_sgName;
 
    }; // class ContainerProcessorRegular
 
 

    class ElementProcessorMet : public asg::AsgMessaging, public IObjectProcessor {
 
    public:
      ElementProcessorMet (const std::string& sgName, const std::string& termName);
      virtual ~ElementProcessorMet() = default;
      ElementProcessorMet( const ElementProcessorMet& ) = delete;
      ElementProcessorMet& operator=( const ElementProcessorMet& ) = delete;

      virtual StatusCode retrieveProcess (StoreType& evtStore) override;

      virtual StatusCode addBranch( TTree& tree, const BranchConfig& branchConfig, OutputBranchData& outputData ) override;
 
      virtual StatusCode addBranch( ROOT::RNTupleModel& model, const BranchConfig& branchConfig, OutputBranchData& outputData ) override;
 
    private:

      std::vector<std::unique_ptr<ElementBranchProcessor>> m_branches;

      std::string m_sgName;

      std::string m_termName;
 
    }; // class ElementProcessorMet
 
 
    class ProcessorList : public asg::AsgMessagingForward {
    public:
      using AsgMessagingForward::AsgMessagingForward;
      virtual ~ProcessorList() = default;
      ProcessorList( const ProcessorList& ) = delete;
      ProcessorList& operator=( const ProcessorList& ) = delete;

      StatusCode setupTree(const std::vector<std::string>& branches, std::unordered_set<std::string> nonContainers, ISystematicsSvc& sysSvc, TTree& tree );

      StatusCode setupBranch( const BranchConfig& branchConfig, OutputBranchData& outputData, TTree& tree );
 
      StatusCode process (StoreType& evtStore);
 
      IObjectProcessor& getObjectProcessor (const BranchConfig& branchConfig, const std::string& sgName);

      std::unordered_set<std::string> m_nonContainers;

      std::optional<int> defaultBasketSize;

      std::unordered_map< std::string, std::unique_ptr<IObjectProcessor>> m_processors;
    };
  }
}
 
#endif