Cache Iteration indexes instead of computing them on the spot (#1860)

* attempt to fix issue #1859 by caching the index and avoid indexOf() call
for iterations

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Franz Pöschel <franz.poeschel@gmail.com>

Author: 3 people

include/openPMD/Iteration.hpp

@@ -129,6 +129,15 @@ namespace internal
          */
         StepStatus m_stepStatus = StepStatus::NoStep;
 
+        /**
+         * Cached copy of the key under which this Iteration lives in
+         * <code>Series::iterations</code>.  Populated when the iteration
+         * object is created/inserted.  This allows constant-time lookup
+         * of the owning map entry instead of a linear scan in
+         * <code>Series::indexOf()</code>.
+         */
+        std::optional<uint64_t> m_iterationIndex = 0;
+
         /**
          * Information on a parsing request that has not yet been executed.
          * Otherwise empty.
@@ -153,6 +162,8 @@ class Iteration : public Attributable
     friend class Writable;
     friend class StatefulIterator;
     friend class StatefulSnapshotsContainer;
+    template <typename>
+    friend struct traits::GenerationPolicy;
 
 public:
     Iteration(Iteration const &) = default;
@@ -276,6 +287,16 @@ class Iteration : public Attributable
 private:
     Iteration();
 
+    /**
+     * @brief Get the cached iteration index.
+     *        This is the key under which this iteration is stored in the
+     *        Series::iterations map. Used internally for testing the index
+     *        caching optimization.
+     *
+     * @return The cached iteration index.
+     */
+    uint64_t getCachedIterationIndex() const;
+
     using Data_t = internal::IterationData;
     std::shared_ptr<Data_t> m_iterationData;
 
@@ -431,6 +452,20 @@ class Iteration : public Attributable
     void runDeferredParseAccess();
 }; // Iteration
 
+namespace traits
+{
+    template <>
+    struct GenerationPolicy<Iteration>
+    {
+        constexpr static bool is_noop = false;
+        template <typename Iterator>
+        void operator()(Iterator &it)
+        {
+            it->second.get().m_iterationIndex = it->first;
+        }
+    };
+} // namespace traits
+
 extern template float Iteration::time<float>() const;
 
 extern template double Iteration::time<double>() const;

include/openPMD/ParticleSpecies.hpp

@@ -62,9 +62,9 @@ namespace traits
     {
         constexpr static bool is_noop = false;
         template <typename T>
-        void operator()(T &ret)
+        void operator()(T &it)
         {
-            ret.particlePatches.linkHierarchy(ret.writable());
+            it->second.particlePatches.linkHierarchy(it->second.writable());
         }
     };
 } // namespace traits

include/openPMD/backend/ContainerImpl.tpp

@@ -140,7 +140,8 @@ auto Container<T, T_key, T_container>::operator[](key_type const &key)
 
         T t = T();
         t.linkHierarchy(writable());
-        auto &ret = container().insert({key, std::move(t)}).first->second;
+        auto inserted_iterator = container().insert({key, std::move(t)}).first;
+        auto &ret = inserted_iterator->second;
         if constexpr (std::is_same_v<T_key, std::string>)
         {
             ret.writable().ownKeyWithinParent = key;
@@ -150,7 +151,7 @@ auto Container<T, T_key, T_container>::operator[](key_type const &key)
             ret.writable().ownKeyWithinParent = std::to_string(key);
         }
         traits::GenerationPolicy<T> gen;
-        gen(ret);
+        gen(inserted_iterator);
         return ret;
     }
 }
@@ -172,7 +173,8 @@ auto Container<T, T_key, T_container>::operator[](key_type &&key)
 
         T t = T();
         t.linkHierarchy(writable());
-        auto &ret = container().insert({key, std::move(t)}).first->second;
+        auto inserted_iterator = container().insert({key, std::move(t)}).first;
+        auto &ret = inserted_iterator->second;
         if constexpr (std::is_same_v<T_key, std::string>)
         {
             ret.writable().ownKeyWithinParent = std::move(key);
@@ -182,7 +184,7 @@ auto Container<T, T_key, T_container>::operator[](key_type &&key)
             ret.writable().ownKeyWithinParent = std::to_string(std::move(key));
         }
         traits::GenerationPolicy<T> gen;
-        gen(ret);
+        gen(inserted_iterator);
         return ret;
     }
 }

src/Iteration.cpp

@@ -59,6 +59,16 @@ Iteration::Iteration() : Attributable(NoInit())
     particles.writable().ownKeyWithinParent = "particles";
 }
 
+uint64_t Iteration::getCachedIterationIndex() const
+{
+    auto idx = get().m_iterationIndex;
+    if (!idx.has_value())
+    {
+        throw error::Internal("Iteration index not known.");
+    }
+    return *idx;
+}
+
 template <typename T>
 Iteration &Iteration::setTime(T newTime)
 {

src/Series.cpp

@@ -2604,16 +2604,25 @@ std::string Series::iterationFilename(IterationIndex_t i)
 Series::iterations_iterator Series::indexOf(Iteration const &iteration)
 {
     auto &series = get();
-    for (auto it = series.iterations.begin(); it != series.iterations.end();
-         ++it)
+    // first try the cached index; if it points to the correct entry return it
+    auto idx = iteration.get().m_iterationIndex;
+    if (!idx.has_value())
     {
-        if (&it->second.Attributable::get() == &iteration.Attributable::get())
-        {
-            return it;
-        }
+        throw error::Internal("Iteration index not known.");
+    }
+
+    auto it = series.iterations.find(*idx);
+    if (it != series.iterations.end() &&
+        &it->second.Attributable::get() == &iteration.Attributable::get())
+    {
+        return it;
+    }
+    else
+    {
+        throw error::Internal(
+            "Iteration " + std::to_string(*idx) +
+            " no longer known by the Series?");
     }
-    throw std::runtime_error(
-        "[Iteration::close] Iteration not found in Series.");
 }
 
 AdvanceStatus Series::advance(