Make subset collections behave the same as other relations

Includes a fix to the setReferences behavior, which previously was only
available for types with at least one relation (this slipped through
until now for some reason).

python/templates/Collection.cc.jinja2

@@ -89,12 +89,8 @@ void {{ collection_type }}::prepareAfterRead() {
   m_isReadFromFile = true;
 }
 
-bool {{ collection_type }}::setReferences(const podio::ICollectionProvider* {% if OneToManyRelations or OneToOneRelations -%}collectionProvider{%- endif -%}) {
-{% if OneToManyRelations or OneToOneRelations %}
-  m_storage.setReferences(collectionProvider, m_isSubsetColl);
-{% endif %}
-
-  return true; //TODO: check success
+bool {{ collection_type }}::setReferences(const podio::ICollectionProvider* collectionProvider) {
+  return m_storage.setReferences(collectionProvider, m_isSubsetColl);
 }
 
 void {{ collection_type }}::push_back({{ class.bare_type }} object) {

python/templates/CollectionData.cc.jinja2

@@ -160,14 +160,13 @@ void {{ class_type }}::createRelations({{ class.bare_type }}Obj* obj) {
 }
 {% endif %}
 
-{% if OneToManyRelations or OneToOneRelations %}
-void {{ class_type }}::setReferences(const podio::ICollectionProvider* collectionProvider, bool isSubsetColl) {
+bool {{ class_type }}::setReferences(const podio::ICollectionProvider* collectionProvider, bool isSubsetColl) {
   if (isSubsetColl) {
     for (const auto& id : *m_refCollections[0]) {
-{{ macros.get_collection(class.full_type) }}
-      entries.push_back(tmp_coll->m_storage.entries[id.index]);
+{{ macros.get_obj_ptr(class.full_type) }}
+      entries.push_back(obj);
     }
-    return;
+    return true; // TODO: check success, how?
   }
 
   // Normal collections have to resolve all relations
@@ -178,8 +177,8 @@ void {{ class_type }}::setReferences(const podio::ICollectionProvider* collectio
 {{ macros.set_reference_single_relation(relation, loop.index0, OneToManyRelations | length) }}
 {% endfor %}
 
+  return true; // TODO: check success, how?
 }
-{% endif %}
 
 void {{ class_type }}::makeSubsetCollection() {
   // Subset collections do not need all the data buffers that normal

python/templates/CollectionData.h.jinja2

@@ -63,9 +63,7 @@ public:
   void createRelations({{ class.bare_type }}Obj* obj);
 {% endif %}
 
-{% if OneToManyRelations or OneToOneRelations %}
-  void setReferences(const podio::ICollectionProvider* collectionProvider, bool isSubsetColl);
-{% endif %}
+  bool setReferences(const podio::ICollectionProvider* collectionProvider, bool isSubsetColl);
 
 private:
   // members to handle 1-to-N-relations

python/templates/macros/collections.jinja2

@@ -51,10 +51,13 @@ const std::array<{{ member.full_type }}, arraysize> {{ class.bare_type }}Collect
   }
 {% endmacro %}
 
-{% macro get_collection(type) %}
+{% macro get_obj_ptr(type) %}
       podio::CollectionBase* coll = nullptr;
-      collectionProvider->get(id.collectionID, coll);
-      {{ type }}Collection* tmp_coll = static_cast<{{ type }}Collection*>(coll);
+      {{ type }}Obj* obj = nullptr;
+      if (collectionProvider->get(id.collectionID, coll)) {
+        auto* tmp_coll = static_cast<{{ type }}Collection*>(coll);
+        obj = tmp_coll->m_storage.entries[id.index];
+      }
 {%- endmacro %}
 
 {% macro set_references_multi_relation(relation, index) %}

tests/read_and_write_associated.cpp

@@ -12,13 +12,14 @@
 void writeCollection() {
 
   auto store = podio::EventStore();
-  podio::ROOTWriter writer("example1.root", &store);
+  podio::ROOTWriter writer("associations.root", &store);
 
   std::cout << "start writting collections...\n";
   auto& info       = store.create<EventInfoCollection>("info");
   auto& hits       = store.create<ExampleHitCollection>("hits");
   auto& clusters   = store.create<ExampleClusterCollection>("clusters");
   auto& hits_subset   = store.create<ExampleHitCollection>("hits_subset");
+  hits_subset.setSubsetCollection(true);
 
   writer.registerForWrite("clusters");
   // writer.registerForWrite("hits");
@@ -59,7 +60,6 @@ void writeCollection() {
     cluster.addClusters( clu1 );
 
     // Add tracked hits to subset hits collection
-    hits_subset.setSubsetCollection(true);
     hits_subset.push_back(hit1);
     hits_subset.push_back(hit2);
 
@@ -78,7 +78,7 @@ void readCollection() {
 
   // Start reading the input
   auto reader = podio::ROOTReader();
-  reader.openFile("example1.root");
+  reader.openFile("associations.root");
 
   auto store = podio::EventStore();
   store.setReader(&reader);
@@ -105,11 +105,21 @@ void readCollection() {
     // Test for subset collections
     auto& hits_subset = store.get<ExampleHitCollection>("hits_subset");
     if(hits_subset.isValid()) {
+      if (!hits_subset.isSubsetCollection()) {
+        throw std::runtime_error("hits_subset should be a subset collection");
+      }
+
+      if (hits_subset.size() != 2) {
+        throw std::runtime_error("subset collection should have original size");
+      }
+
       for (const auto& hit : hits_subset) {
         if (hit.isAvailable()) {
           throw std::runtime_error("Hit is available, although it has not been written");
         }
       }
+    } else {
+      throw std::runtime_error("Collection 'hits_subset' should be present");
     }
 
     store.clear();