Add indexed access to the RelationRange (#194)

include/podio/RelationRange.h

@@ -16,19 +16,36 @@ namespace podio {
 
     RelationRange() = delete;
 
-    RelationRange(ConstIteratorType begin, ConstIteratorType end) : m_begin(begin), m_end(end) {}
+    RelationRange(ConstIteratorType begin, ConstIteratorType end) : m_begin(begin), m_end(end), m_size(std::distance(m_begin, m_end)) {}
 
     /// begin of the range (necessary for range-based for loop)
     ConstIteratorType begin() const { return m_begin; }
     /// end of the range (necessary for range-based for loop)
     ConstIteratorType end() const { return m_end; }
     /// convenience overload for size
-    size_t size() const { return std::distance(m_begin, m_end); }
+    size_t size() const { return m_size; }
     /// convenience overload to check if the range is empty
     bool empty() const { return m_begin == m_end; }
+    /// Indexed access
+    ReferenceType operator[](size_t i) const {
+      auto it = m_begin;
+      std::advance(it, i);
+      return *it;
+    }
+    /// Indexed access with range check
+    ReferenceType at(size_t i) const {
+      if (i < m_size) {
+        auto it = m_begin;
+        std::advance(it, i);
+        return *it;
+      }
+      throw std::out_of_range("index out of bounds for RelationRange");
+    }
+
   private:
     ConstIteratorType m_begin;
     ConstIteratorType m_end;
+    size_t m_size{0};
   };
 }
 

python/test_EventStore.py

@@ -59,6 +59,18 @@ def test_one_to_many(self):
     for hit in ref_hits:
       self.assertTrue(hit in hits)
 
+  def test_relation_range(self):
+    """Test that the RelationRange functionality is also accessible in python"""
+    clusters = self.store.get("clusters")
+    hits = self.store.get("hits")
+
+    for cluster in clusters:
+      sume = 0
+      for hit in cluster.Hits():
+        self.assertTrue(hit in hits)
+        sume += hit.energy()
+      self.assertEqual(cluster.energy(), sume)
+
   def test_hash(self):
     clusters = self.store.get("clusters")
     ref_hits = []

tests/relation_range.cpp

@@ -64,6 +64,10 @@ void doTestExampleMC(ExampleMCCollection const& collection)
   // Empty
   ASSERT_CONDITION(collection[7].daughters().size() == 0 && collection[7].parents().size() == 0,
                    "RelationRange of empty collection is not empty");
+  // Equivalent but potentially quicker way of checking an empty collection
+  ASSERT_CONDITION(collection[7].daughters().empty() && collection[7].parents().empty(),
+                   "RelationRange of empty collection is not empty");
+
   // alternatively check if a loop is entered
   for (const auto& p[[maybe_unused]]: collection[7].daughters()) {
     throw std::runtime_error("Range based for loop entered on a supposedly empty range");
@@ -79,6 +83,14 @@ void doTestExampleMC(ExampleMCCollection const& collection)
     index++;
   }
 
+  // Check indexed access
+  const auto daughters = collection[0].daughters();
+  for (size_t i = 0; i < expectedPDG.size(); ++i) {
+    const auto daughter = daughters[i];
+    ASSERT_EQUAL(daughter.PDG(), expectedPDG[i],
+                 "ExampleMC daughter points to the wrong particle (by PDG)");
+  }
+
   // mothers and daughters
   ASSERT_EQUAL(collection[2].daughters().size(), 2, "Range has wrong size");
   ASSERT_EQUAL(collection[2].parents().size(), 2, "Range has wrong size");
@@ -98,6 +110,22 @@ void doTestExampleMC(ExampleMCCollection const& collection)
     index++;
   }
 
+  // Indexed access with range check
+  const auto parents = collection[2].parents();
+  for (size_t i = 0; i < expectedPDG.size(); ++i) {
+    const auto parent = parents.at(i);
+    ASSERT_EQUAL(parent.PDG(), expectedPDG[i],
+                 "ExampleMC parents points to the wrong particle (by PDG)");
+  }
+
+  try {
+    const auto parent = parents.at(3);
+    throw std::runtime_error("Trying to access out of bounds in a RelationRange::at should throw");
+  } catch (const std::out_of_range& err) {
+    ASSERT_EQUAL(err.what(), std::string("index out of bounds for RelationRange"),
+                 "Access out of bounds throws wrong exception");
+  }
+
   // realistic case
   auto mcp6 = collection[6];
   ASSERT_EQUAL(mcp6.daughters().size(), 1, "Wrong number of daughters");
@@ -121,7 +149,6 @@ void testExampleMC()
  doTestExampleMC(mcps);
 }
 
-
 void testExampleWithVectorMember()
 {
   ExampleWithVectorMember ex;