[wip] Make associations work with Root I/O

AIDASoft · Mar 22, 2023 · 8f534c4 · 8f534c4
1 parent 22256bf
commit 8f534c4
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 2 deletions.
diff --git a/include/podio/AssociationCollection.h b/include/podio/AssociationCollection.h
@@ -1,17 +1,20 @@
 #ifndef PODIO_ASSOCIATIONCOLLECTION_H
 #define PODIO_ASSOCIATIONCOLLECTION_H
 
+#include "podio/CollectionBufferFactory.h"
 #include "podio/detail/AssociationCollectionData.h"
 #include "podio/detail/AssociationFwd.h"
 #include "podio/detail/AssociationObj.h"
 
 #include "podio/Association.h"
 #include "podio/AssociationCollectionIterator.h"
 #include "podio/CollectionBase.h"
+#include "podio/CollectionBufferFactory.h"
 #include "podio/CollectionBuffers.h"
 #include "podio/DatamodelRegistry.h"
 #include "podio/ICollectionProvider.h"
 #include "podio/SchemaEvolution.h"
+#include "podio/utilities/TypeHelpers.h"
 
 #ifdef PODIO_JSON_OUTPUT
   #include "nlohmann/json.hpp"
@@ -26,6 +29,13 @@
 
 namespace podio {
 
+template <typename FromT, typename ToT>
+std::string associationCollTypeName() {
+  const static std::string typeName =
+      std::string("podio::AssociationCollection<") + FromT::TypeName + "," + ToT::TypeName + ">";
+  return typeName;
+}
+
 template <typename FromT, typename ToT>
 class AssociationCollection : public podio::CollectionBase {
   static_assert(std::is_same_v<FromT, detail::GetDefT<FromT>>,
@@ -301,6 +311,48 @@ void to_json(nlohmann::json& j, const AssociationCollection<FromT, ToT>& collect
 }
 #endif
 
+template <typename FromT, typename ToT>
+bool registerAssociationCollection(const std::string& assocTypeName) {
+  const static auto reg = [&assocTypeName]() {
+    auto& factory = CollectionBufferFactory::mutInstance();
+    factory.registerCreationFunc(assocTypeName, AssociationCollection<FromT, ToT>::schemaVersion, [](bool subsetColl) {
+      auto readBuffers = podio::CollectionReadBuffers{};
+      readBuffers.data = subsetColl ? nullptr : new AssociationDataContainer();
+
+      // Either it is a subset collection or we have two relations
+      const auto nRefs = subsetColl ? 1 : 2;
+      readBuffers.references = new podio::CollRefCollection(nRefs);
+      for (auto& ref : *readBuffers.references) {
+        // Make sure to place usable buffer pointers here
+        ref = std::make_unique<std::vector<podio::ObjectID>>();
+      }
+
+      readBuffers.createCollection = [](podio::CollectionReadBuffers buffers, bool isSubsetColl) {
+        AssociationCollectionData<FromT, ToT> data(buffers, isSubsetColl);
+        return std::make_unique<AssociationCollection<FromT, ToT>>(std::move(data), isSubsetColl);
+      };
+
+      readBuffers.recast = [](podio::CollectionReadBuffers& buffers) {
+        if (buffers.data) {
+          buffers.data = podio::CollectionWriteBuffers::asVector<float>(buffers.data);
+        }
+      };
+
+      return readBuffers;
+    });
+
+    return true;
+  }();
+  return reg;
+}
+
+#define PODIO_DECLARE_ASSOCIATION(TypeName, FromT, ToT)                                                                \
+  namespace {                                                                                                          \
+    using TypeName = podio::AssociationCollection<FromT, ToT>;                                                         \
+    const auto registerAssociation =                                                                                   \
+        podio::registerAssociationCollection<FromT, ToT>(podio::associationCollTypeName<FromT, ToT>());                \
+  } // namespace podio
+
 } // namespace podio
 
 #endif // PODIO_ASSOCIATIONCOLLECTION_H
diff --git a/tests/read_test.h b/tests/read_test.h
@@ -32,7 +32,8 @@
 #include <vector>
 
 // Define an association that is used for the I/O tests
-using TestAssocCollection = podio::AssociationCollection<ExampleMC, ex42::ExampleWithARelation>;
+// using TestAssocCollection = podio::AssociationCollection<ExampleMC, ex42::ExampleWithARelation>;
+PODIO_DECLARE_ASSOCIATION(TestAssocCollection, ExampleMC, ex42::ExampleWithARelation)
 
 template <typename FixedWidthT>
 bool check_fixed_width_value(FixedWidthT actual, FixedWidthT expected, const std::string& type) {

diff --git a/tests/write_frame.h b/tests/write_frame.h
@@ -26,7 +26,9 @@
 #include <vector>
 
 // Define an association that is used for the I/O tests
-using TestAssocCollection = podio::AssociationCollection<ExampleMC, ex42::ExampleWithARelation>;
+// using TestAssocCollection = podio::AssociationCollection<ExampleMC, ex42::ExampleWithARelation>;
+
+PODIO_DECLARE_ASSOCIATION(TestAssocCollection, ExampleMC, ex42::ExampleWithARelation)
 
 static const std::vector<std::string> collsToWrite = {"mcparticles",
                                                       "moreMCs",