Flow.h

Go to the documentation of this file.

// SPDX-License-Identifier: GPLv3-or-later
//
// SPDX-FileCopyrightText: Patrick L.S. Connor <patrick.connor@desy.de>
// SPDX-FileCopyrightText: Louis Moureaux <louis.moureaux@cern.ch>
 
#pragma once
 
#include "exceptions.h"
#include "FriendUtils.h"
#include "Looper.h"
 
#include <TChain.h>
#include <TFile.h>
#include <TH1.h>
#include <TTree.h>
 
#include <any>
#include <filesystem>
#include <memory>
#include <ranges>
#include <source_location>
#include <string>
#include <vector>
 
namespace Darwin::Tools {
 
enum BranchMode {
    mandatory  , 
    facultative  
}; 
 
using enum BranchMode;
 
class Flow {
 
    int steering; 
 
    std::unique_ptr<ChainSlice> tIn; 
    std::shared_ptr<TFile> fOut; 
    std::unique_ptr<TTree> tOut; 
 
    std::map<std::string, std::any> branches; 
 
    template<typename T> 
    std::shared_ptr<T*> GetBranch (const std::string& name) const 
    try {
        using namespace std;
        if (steering & verbose)
            cout << "Flow: retrieving an existing branch" << endl;
        auto& branch = branches.at(name);
        return any_cast<shared_ptr<T*>>(branch);
    }
    catch (const std::bad_any_cast& e) {
        BOOST_THROW_EXCEPTION(e);
    }
 
    template<typename T> T * NoBranch (const std::string& name, BranchMode mode)
    {
        if (mode == facultative)
            return nullptr;
 
        namespace DE = Darwin::Exceptions;
        std::string what = name + " branch could not be found";
        BOOST_THROW_EXCEPTION( DE::BadInput(what.c_str(), *tIn) );
    }
 
public:
    inline std::string DumpActiveBranches () const
    {
        auto keys = branches | std::views::keys;
        return std::accumulate(keys.begin(), keys.end(), std::string(),
                [](const std::string& str, const std::string& branch_name) {
                    return str + ' ' + branch_name;
                });
    }
 
    Flow (int = none); 
 
    ~Flow ();
 
    ChainSlice * GetInputTree
            (std::vector<std::filesystem::path> = {}, 
             const Slice = {1, 0}, 
             const std::string& = "events"); 
 
    TTree * GetOutputTree (std::shared_ptr<TFile> = {}, 
            const std::source_location = std::source_location::current()); 
 
    TTree * GetOutputTree (const std::filesystem::path&, 
            const std::source_location = std::source_location::current()); 
 
    inline void SetOutputFile (std::shared_ptr<TFile> fOut) { this->fOut = fOut; }
 
    inline TFile * GetOutputFile () { return fOut.get(); }
 
    std::pair<TFile *, TTree *> GetOutput (const std::filesystem::path&, 
            const std::source_location = std::source_location::current()); 
 
    template<typename T> 
    T * GetBranchReadOnly (
                const std::string& name, 
                BranchMode mode = mandatory) 
    {
        using namespace std;
 
        if (!tIn)
            BOOST_THROW_EXCEPTION( invalid_argument("`GetInputTree()` should "
                            "be called before declaring a read-only branch") );
 
        namespace DE = Darwin::Exceptions;
 
        shared_ptr<T*> t;
 
        if (branches.contains(name))
            t = GetBranch<T>(name);
        else {
            t = make_shared<T*>();
            if (steering & verbose)
                cout << "Flow: loading branch `" << name << "`" << endl;
 
            if (tIn->GetBranch(name.c_str()) == nullptr)
                return NoBranch<T>(name, mode);
 
            int err = tIn->SetBranchAddress(name.c_str(), t.get());
            if (steering & verbose)
                cout << "Flow: `TTree::SetBranchAddress()` returned " << to_string(err)
                     << " (check `TTree::ESetBranchAddressStatus` for the meaning)."
                     << endl;
            if (err < 0) {
                string what = "`"s + name + "` branch could not be set. "s;
                if (mode == facultative) {
                    if (steering & verbose)
                        cout << orange << "Flow: " << what << def << endl;
                    return nullptr;
                }
                BOOST_THROW_EXCEPTION( DE::BadInput(what.c_str(), *tIn) );
            }
            branches.insert({name, t});
        }
        return *t;
    }
 
    template<typename T> 
    T * GetBranchWriteOnly (const std::string& name) 
    {
        using namespace std;
 
        if (!tOut)
            BOOST_THROW_EXCEPTION( invalid_argument("`GetOutputTree()` should "
                                                    "be called before") );
 
        shared_ptr<T*> t;
 
        if (branches.contains(name))
            t = GetBranch<T>(name);
        else
            t = make_shared<T*>();
 
        if (steering & verbose)
            cout << "Flow: setting up new branch for `" << name << "`" << endl;
        if (tOut->Branch(name.c_str(), t.get()) == nullptr) {
            namespace DE = Darwin::Exceptions;
            string what = name + " branch could not be set up";
            BOOST_THROW_EXCEPTION( DE::BadInput(what.c_str(), *tOut) );
        }
 
        if (!branches.contains(name))
            branches.insert({name, t});
        return *t;
    }
 
    template<typename T> 
    T * GetBranchReadWrite (const std::string& name, 
                            BranchMode mode = mandatory) 
    {
        using namespace std;
 
        if (!tIn)
            BOOST_THROW_EXCEPTION( invalid_argument("`GetInputTree()` should "
                            "be called before declaring a read-write branch") );
 
        if (!tOut)
            BOOST_THROW_EXCEPTION( invalid_argument("`GetOutputTree()` should "
                            "be called before declaring a read-write branch") );
 
        shared_ptr<T*> t;
 
        if (branches.contains(name))
            t = GetBranch<T>(name);
        else {
            if (GetBranchReadOnly<T>(name, mode) == nullptr)
                return NoBranch<T>(name, mode);
            if (steering & Friend)
                GetBranchWriteOnly<T>(name);
            t = any_cast<shared_ptr<T*>>(branches[name]);
        }
        return *t;
    }
};
 
} // namespace Darwin::Tools