From 5f04132935cb157ca78cbc0a36cbe1ff6d1f5c10 Mon Sep 17 00:00:00 2001
From: Thomas Madlener <thomas.madlener@desy.de>
Date: Fri, 19 Jul 2024 10:58:35 +0200
Subject: [PATCH] Add conversion of associations to LCRelations (#189)

* Add conversion of associations to LCRelations

* Move collection retrieval into separate function
---
 .../k4MarlinWrapper/converters/EDM4hep2Lcio.h |  5 ++-
 .../src/components/EDM4hep2Lcio.cpp           | 35 ++++++++++++++-----
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/k4MarlinWrapper/k4MarlinWrapper/converters/EDM4hep2Lcio.h b/k4MarlinWrapper/k4MarlinWrapper/converters/EDM4hep2Lcio.h
index 3b515451..d1671d5b 100644
--- a/k4MarlinWrapper/k4MarlinWrapper/converters/EDM4hep2Lcio.h
+++ b/k4MarlinWrapper/k4MarlinWrapper/converters/EDM4hep2Lcio.h
@@ -113,6 +113,9 @@ class EDM4hep2LcioTool : public GaudiTool, virtual public IEDMConverter {
   void convertAdd(const std::string& e4h_coll_name, const std::string& lcio_coll_name, lcio::LCEventImpl* lcio_event,
                   CollectionPairMappings&                            collection_pairs,
                   std::vector<EDM4hep2LCIOConv::ParticleIDConvData>& pidCollections);
-};
 
+  /// Get an EDM4hep collection by name, consulting either the podio based data
+  /// svc or the IOSvc
+  podio::CollectionBase* getEDM4hepCollection(const std::string& name) const;
+};
 #endif
diff --git a/k4MarlinWrapper/src/components/EDM4hep2Lcio.cpp b/k4MarlinWrapper/src/components/EDM4hep2Lcio.cpp
index f6b3487d..652d6d38 100644
--- a/k4MarlinWrapper/src/components/EDM4hep2Lcio.cpp
+++ b/k4MarlinWrapper/src/components/EDM4hep2Lcio.cpp
@@ -278,14 +278,10 @@ void EDM4hep2LcioTool::convertEventHeader(const std::string& e4h_coll_name, lcio
   EDM4hep2LCIOConv::convertEventHeader(header_coll, lcio_event);
 }
 
-// Select the appropiate method to convert a collection given its type
-void EDM4hep2LcioTool::convertAdd(const std::string& e4h_coll_name, const std::string& lcio_coll_name,
-                                  lcio::LCEventImpl* lcio_event, CollectionPairMappings& collection_pairs,
-                                  std::vector<EDM4hep2LCIOConv::ParticleIDConvData>& pidCollections) {
-  const auto&            metadata = m_podioDataSvc->getMetaDataFrame();
-  podio::CollectionBase* collPtr  = nullptr;
+podio::CollectionBase* EDM4hep2LcioTool::getEDM4hepCollection(const std::string& collName) const {
+  podio::CollectionBase* collPtr{nullptr};
   DataObject*            p;
-  auto                   sc = m_podioDataSvc->retrieveObject(e4h_coll_name, p);
+  auto                   sc = m_podioDataSvc->retrieveObject(collName, p);
   if (sc.isFailure()) {
     throw GaudiException("Collection not found", name(), StatusCode::FAILURE);
   }
@@ -304,7 +300,17 @@ void EDM4hep2LcioTool::convertAdd(const std::string& e4h_coll_name, const std::s
       collPtr = dynamic_cast<podio::CollectionBase*>(ptr->getData().get());
     }
   }
-  const auto fulltype = collPtr->getValueTypeName();
+
+  return collPtr;
+}
+
+// Select the appropiate method to convert a collection given its type
+void EDM4hep2LcioTool::convertAdd(const std::string& e4h_coll_name, const std::string& lcio_coll_name,
+                                  lcio::LCEventImpl* lcio_event, CollectionPairMappings& collection_pairs,
+                                  std::vector<EDM4hep2LCIOConv::ParticleIDConvData>& pidCollections) {
+  const auto& metadata = m_podioDataSvc->getMetaDataFrame();
+  const auto  collPtr  = getEDM4hepCollection(e4h_coll_name);
+  const auto  fulltype = collPtr->getValueTypeName();
 
   debug() << "Converting type " << fulltype << " from input " << e4h_coll_name << endmsg;
 
@@ -373,7 +379,15 @@ StatusCode EDM4hep2LcioTool::convertCollections(lcio::LCEventImpl* lcio_event) {
   CollectionPairMappings                            collection_pairs{};
   std::vector<EDM4hep2LCIOConv::ParticleIDConvData> pidCollections{};
 
+  std::vector<std::tuple<std::string, const podio::CollectionBase*>> associations{};
+
   for (const auto& [edm4hepName, lcioName] : collsToConvert) {
+    const auto coll = getEDM4hepCollection(edm4hepName);
+    if (coll->getTypeName().find("Association") != std::string_view::npos) {
+      debug() << edm4hepName << " is an association collection, converting it later" << endmsg;
+      associations.emplace_back(lcioName, coll);
+      continue;
+    }
     debug() << "Converting collection " << edm4hepName << " (storing it as " << lcioName << ")" << endmsg;
     if (!EDM4hep2LCIOConv::collectionExist(lcioName, lcio_event)) {
       convertAdd(edm4hepName, lcioName, lcio_event, collection_pairs, pidCollections);
@@ -413,5 +427,10 @@ StatusCode EDM4hep2LcioTool::convertCollections(lcio::LCEventImpl* lcio_event) {
 
   EDM4hep2LCIOConv::resolveRelations(collection_pairs, globalObjMap);
 
+  // Now we can convert the assocations and add them to the event
+  for (auto& [name, coll] : EDM4hep2LCIOConv::createLCRelationCollections(associations, globalObjMap)) {
+    lcio_event->addCollection(coll.release(), name);
+  }
+
   return StatusCode::SUCCESS;
 }