Fix a few more issues

AIDASoft · Jul 1, 2024 · d4edbf9 · d4edbf9
1 parent 46480b5
commit d4edbf9
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 19 deletions.
diff --git a/include/podio/GenericParameters.h b/include/podio/GenericParameters.h
@@ -13,6 +13,7 @@
 #include <map>
 #include <memory>
 #include <mutex>
+#include <optional>
 #include <string>
 #include <tuple>
 #include <type_traits>

diff --git a/src/ROOTLegacyReader.cc b/src/ROOTLegacyReader.cc
@@ -25,7 +25,6 @@
 #include "TClass.h"
 #include "TFile.h"
 #include "TTree.h"
-#include "podio/CollectionBranches.h"
 #include "podio/podioVersion.h"
 
 namespace podio {

diff --git a/src/ROOTWriter.cc b/src/ROOTWriter.cc
@@ -7,6 +7,7 @@
 #include "podio/Frame.h"
 #include "podio/GenericParameters.h"
 #include "podio/podioVersion.h"
+#include "podio/utilities/RootHelpers.h"
 #include "rootUtils.h"
 
 #include <TBranch.h>
@@ -43,7 +44,7 @@ void ROOTWriter::writeFrame(const podio::Frame& frame, const std::string& catego
     catInfo.tree->SetDirectory(m_file.get());
   }
 
-  std::vector<StoreCollection> collections;
+  std::vector<root_utils::StoreCollection> collections;
   collections.reserve(catInfo.collsToWrite.size());
   for (const auto& name : catInfo.collsToWrite) {
     auto* coll = frame.getCollectionForWrite(name);
@@ -67,7 +68,8 @@ void ROOTWriter::writeFrame(const podio::Frame& frame, const std::string& catego
       throw std::runtime_error("Trying to write category '" + category + "' with inconsistent collection content. " +
                                root_utils::getInconsistentCollsMsg(catInfo.collsToWrite, collsToWrite));
     }
-    resetBranches(catInfo.branches, collections, &const_cast<podio::GenericParameters&>(frame.getParameters()));
+    fillParams(catInfo, frame.getParameters());
+    resetBranches(catInfo, collections);
   }
 
   catInfo.tree->Fill();
@@ -82,9 +84,9 @@ ROOTWriter::CategoryInfo& ROOTWriter::getCategoryInfo(const std::string& categor
   return it->second;
 }
 
-void ROOTWriter::initBranches(CategoryInfo& catInfo, const std::vector<StoreCollection>& collections,
+void ROOTWriter::initBranches(CategoryInfo& catInfo, const std::vector<root_utils::StoreCollection>& collections,
                               /*const*/ podio::GenericParameters& parameters) {
-  catInfo.branches.reserve(collections.size() + 1); // collections + parameters
+  catInfo.branches.reserve(collections.size() + root_utils::nParamBranches); // collections + parameters
 
   // First collections
   for (auto& [name, coll] : collections) {
@@ -123,28 +125,54 @@ void ROOTWriter::initBranches(CategoryInfo& catInfo, const std::vector<StoreColl
       }
     }
 
-    catInfo.branches.push_back(branches);
-    catInfo.collInfo.emplace_back(catInfo.idTable.collectionID(name).value(), coll->getTypeName(),
+    catInfo.branches.emplace_back(std::move(branches));
+    catInfo.collInfo.emplace_back(catInfo.idTable.collectionID(name).value(), std::string(coll->getTypeName()),
                                   coll->isSubsetCollection(), coll->getSchemaVersion());
   }
 
-  // Also make branches for the parameters
-  root_utils::CollectionBranches branches;
-  branches.data = catInfo.tree->Branch(root_utils::paramBranchName, &parameters);
-  catInfo.branches.push_back(branches);
+  fillParams(catInfo, parameters);
+  // NOTE: The order in which these are created is codified for later use in
+  // root_utils::getGPBranchOffsets
+  catInfo.branches.emplace_back(catInfo.tree->Branch(root_utils::intKeyName, &catInfo.intParams.keys));
+  catInfo.branches.emplace_back(catInfo.tree->Branch(root_utils::intValueName, &catInfo.intParams.values));
+
+  catInfo.branches.emplace_back(catInfo.tree->Branch(root_utils::floatKeyName, &catInfo.floatParams.keys));
+  catInfo.branches.emplace_back(catInfo.tree->Branch(root_utils::floatValueName, &catInfo.floatParams.values));
+
+  catInfo.branches.emplace_back(catInfo.tree->Branch(root_utils::doubleKeyName, &catInfo.doubleParams.keys));
+  catInfo.branches.emplace_back(catInfo.tree->Branch(root_utils::doubleValueName, &catInfo.doubleParams.values));
+
+  catInfo.branches.emplace_back(catInfo.tree->Branch(root_utils::stringKeyName, &catInfo.stringParams.keys));
+  catInfo.branches.emplace_back(catInfo.tree->Branch(root_utils::stringValueName, &catInfo.stringParams.values));
 }
 
-void ROOTWriter::resetBranches(std::vector<root_utils::CollectionBranches>& branches,
-                               const std::vector<ROOTWriter::StoreCollection>& collections,
-                               /*const*/ podio::GenericParameters* parameters) {
+void ROOTWriter::resetBranches(CategoryInfo& categoryInfo,
+                               const std::vector<root_utils::StoreCollection>& collections) {
   size_t iColl = 0;
-  for (auto& coll : collections) {
-    const auto& collBranches = branches[iColl];
-    root_utils::setCollectionAddresses(coll.second->getBuffers(), collBranches);
+  for (auto& [_, coll] : collections) {
+    const auto& collBranches = categoryInfo.branches[iColl];
+    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--;
+
+  constexpr auto intOffset = root_utils::getGPBranchOffsets<int>();
+  categoryInfo.branches[iColl + intOffset.keys].data->SetAddress(categoryInfo.intParams.keysPtr());
+  categoryInfo.branches[iColl + intOffset.values].data->SetAddress(categoryInfo.intParams.valuesPtr());
 
-  branches.back().data->SetAddress(&parameters);
+  constexpr auto floatOffset = root_utils::getGPBranchOffsets<float>();
+  categoryInfo.branches[iColl + floatOffset.keys].data->SetAddress(categoryInfo.floatParams.keysPtr());
+  categoryInfo.branches[iColl + floatOffset.values].data->SetAddress(categoryInfo.floatParams.valuesPtr());
+
+  constexpr auto doubleOffset = root_utils::getGPBranchOffsets<double>();
+  categoryInfo.branches[iColl + doubleOffset.keys].data->SetAddress(categoryInfo.doubleParams.keysPtr());
+  categoryInfo.branches[iColl + doubleOffset.values].data->SetAddress(categoryInfo.doubleParams.valuesPtr());
+
+  constexpr auto stringOffset = root_utils::getGPBranchOffsets<std::string>();
+  categoryInfo.branches[iColl + stringOffset.keys].data->SetAddress(categoryInfo.stringParams.keysPtr());
+  categoryInfo.branches[iColl + stringOffset.values].data->SetAddress(categoryInfo.stringParams.valuesPtr());
 }
 
 void ROOTWriter::finish() {
@@ -181,4 +209,11 @@ ROOTWriter::checkConsistency(const std::vector<std::string>& collsToWrite, const
   return {std::vector<std::string>{}, collsToWrite};
 }
 
+void ROOTWriter::fillParams(CategoryInfo& catInfo, const GenericParameters& params) {
+  catInfo.intParams = params.getKeysAndValues<int>();
+  catInfo.floatParams = params.getKeysAndValues<float>();
+  catInfo.doubleParams = params.getKeysAndValues<double>();
+  catInfo.stringParams = params.getKeysAndValues<std::string>();
+}
+
 } // namespace podio
diff --git a/src/rootUtils.h b/src/rootUtils.h
@@ -3,7 +3,7 @@
 
 #include "TBranch.h"
 #include "TTree.h"
-#include "podio/CollectionBranches.h"
+#include "podio/utilities/RootHelpers.h"
 #include "podio/CollectionIDTable.h"
 
 #include <algorithm>