diff --git a/src/ROOTReader.cc b/src/ROOTReader.cc index 60fa37018..747235e86 100644 --- a/src/ROOTReader.cc +++ b/src/ROOTReader.cc @@ -26,48 +26,21 @@ std::tuple, std::vector& collInfo); -/// Helper struct to get the negative offsets from the end of the branches -/// vector for the different types of generic parameters. -template -struct TypeToBranchIndexOffset; - -template <> -struct TypeToBranchIndexOffset { - constexpr static int keys = 8; - constexpr static int values = 7; -}; - -template <> -struct TypeToBranchIndexOffset { - constexpr static int keys = 6; - constexpr static int values = 5; -}; - -template <> -struct TypeToBranchIndexOffset { - constexpr static int keys = 4; - constexpr static int values = 3; -}; - -template <> -struct TypeToBranchIndexOffset { - constexpr static int keys = 2; - constexpr static int values = 1; -}; - template void ROOTReader::readParams(ROOTReader::CategoryInfo& catInfo, podio::GenericParameters& params, bool reloadBranches, unsigned int localEntry) { - const auto nBranches = catInfo.branches.size(); + const auto collBranchIdx = catInfo.branches.size() - root_utils::nParamBranches - 1; + constexpr auto brOffset = root_utils::getGPBranchOffsets(); + if (reloadBranches) { - auto& keyBranch = catInfo.branches[nBranches - TypeToBranchIndexOffset::keys].data; + auto& keyBranch = catInfo.branches[collBranchIdx + brOffset.keys].data; keyBranch = root_utils::getBranch(catInfo.chain.get(), root_utils::getGPKeyName()); - auto& valueBranch = catInfo.branches[nBranches - TypeToBranchIndexOffset::values].data; + auto& valueBranch = catInfo.branches[collBranchIdx + brOffset.values].data; valueBranch = root_utils::getBranch(catInfo.chain.get(), root_utils::getGPValueName()); } - auto keyBranch = catInfo.branches[nBranches - TypeToBranchIndexOffset::keys].data; - auto valueBranch = catInfo.branches[nBranches - TypeToBranchIndexOffset::values].data; + auto keyBranch = catInfo.branches[collBranchIdx + brOffset.keys].data; + auto valueBranch = catInfo.branches[collBranchIdx + brOffset.values].data; root_utils::ParamStorage storage; keyBranch->SetAddress(storage.keysPtr()); diff --git a/src/ROOTWriter.cc b/src/ROOTWriter.cc index f0aab0acf..b43ac5f92 100644 --- a/src/ROOTWriter.cc +++ b/src/ROOTWriter.cc @@ -80,8 +80,7 @@ ROOTWriter::CategoryInfo& ROOTWriter::getCategoryInfo(const std::string& categor void ROOTWriter::initBranches(CategoryInfo& catInfo, const std::vector& collections, /*const*/ podio::GenericParameters& parameters) { - catInfo.branches.reserve(collections.size() + - std::tuple_size_v * 2); // collections + parameters + catInfo.branches.reserve(collections.size() + root_utils::nParamBranches); // collections + parameters // First collections for (auto& [name, coll] : collections) { @@ -126,6 +125,8 @@ void ROOTWriter::initBranches(CategoryInfo& catInfo, const std::vectorBranch(root_utils::intKeyName, &catInfo.intParams.keys)); catInfo.branches.emplace_back(catInfo.tree->Branch(root_utils::intValueName, &catInfo.intParams.values)); @@ -147,18 +148,25 @@ void ROOTWriter::resetBranches(CategoryInfo& categoryInfo, root_utils::setCollectionAddresses(coll->getBuffers(), collBranches); iColl++; } + // Correct index to point to the last branch of collection data for symmetric + // handling of the offsets in reading and writing + iColl--; - categoryInfo.branches[iColl].data->SetAddress(categoryInfo.intParams.keysPtr()); - categoryInfo.branches[iColl + 1].data->SetAddress(categoryInfo.intParams.valuesPtr()); + constexpr auto intOffset = root_utils::getGPBranchOffsets(); + categoryInfo.branches[iColl + intOffset.keys].data->SetAddress(categoryInfo.intParams.keysPtr()); + categoryInfo.branches[iColl + intOffset.values].data->SetAddress(categoryInfo.intParams.valuesPtr()); - categoryInfo.branches[iColl + 2].data->SetAddress(categoryInfo.floatParams.keysPtr()); - categoryInfo.branches[iColl + 3].data->SetAddress(categoryInfo.floatParams.valuesPtr()); + constexpr auto floatOffset = root_utils::getGPBranchOffsets(); + categoryInfo.branches[iColl + floatOffset.keys].data->SetAddress(categoryInfo.floatParams.keysPtr()); + categoryInfo.branches[iColl + floatOffset.values].data->SetAddress(categoryInfo.floatParams.valuesPtr()); - categoryInfo.branches[iColl + 4].data->SetAddress(categoryInfo.doubleParams.keysPtr()); - categoryInfo.branches[iColl + 5].data->SetAddress(categoryInfo.doubleParams.valuesPtr()); + constexpr auto doubleOffset = root_utils::getGPBranchOffsets(); + categoryInfo.branches[iColl + doubleOffset.keys].data->SetAddress(categoryInfo.doubleParams.keysPtr()); + categoryInfo.branches[iColl + doubleOffset.values].data->SetAddress(categoryInfo.doubleParams.valuesPtr()); - categoryInfo.branches[iColl + 6].data->SetAddress(categoryInfo.stringParams.keysPtr()); - categoryInfo.branches[iColl + 7].data->SetAddress(categoryInfo.stringParams.valuesPtr()); + constexpr auto stringOffset = root_utils::getGPBranchOffsets(); + categoryInfo.branches[iColl + stringOffset.keys].data->SetAddress(categoryInfo.stringParams.keysPtr()); + categoryInfo.branches[iColl + stringOffset.values].data->SetAddress(categoryInfo.stringParams.valuesPtr()); } void ROOTWriter::finish() { diff --git a/src/rootUtils.h b/src/rootUtils.h index c9e7c1238..131e0ef13 100644 --- a/src/rootUtils.h +++ b/src/rootUtils.h @@ -2,6 +2,7 @@ #define PODIO_ROOT_UTILS_H // NOLINT(llvm-header-guard): internal headers confuse clang-tidy #include "podio/CollectionIDTable.h" +#include "podio/GenericParameters.h" #include "podio/utilities/RootHelpers.h" #include "TBranch.h" @@ -80,6 +81,35 @@ constexpr auto getGPValueName() { } } +/// Small helper struct to get info on the offsets of the branches holding +/// GenericParameter keys and values for a given parameter type +struct GPBranchOffsets { + int keys{-1}; + int values{-1}; +}; + +/// The number of branches that we create on top of the collection branches per +/// category +constexpr auto nParamBranches = std::tuple_size_v * 2; + +/// Get the branch offsets for a given parameter type. In this case it is +/// assumed that the integer branches start immediately after the branche for +/// the collections +template +constexpr auto getGPBranchOffsets() { + if constexpr (std::is_same_v) { + return GPBranchOffsets{1, 2}; + } else if constexpr (std::is_same_v) { + return GPBranchOffsets{3, 4}; + } else if constexpr (std::is_same_v) { + return GPBranchOffsets{5, 6}; + } else if constexpr (std::is_same_v) { + return GPBranchOffsets{7, 8}; + } else { + static_assert(sizeof(T) == 0, "Unsupported type for generic parameters"); + } +} + /** * Name of the field with the list of categories for RNTuples */