From 6c7f23a59a1d415d1d3c16c4ccd834bbf7173a04 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Franz=20P=C3=B6schel?= <franz.poeschel@gmail.com>
Date: Wed, 3 May 2023 13:56:06 +0200
Subject: [PATCH] Support custom hierarchies with attributes

---
 include/openPMD/CustomHierarchy.hpp | 15 +++++++-
 src/CustomHierarchy.cpp             | 53 +++++++++++++++++++++++++++--
 src/Iteration.cpp                   | 23 +++++++++++++
 3 files changed, 88 insertions(+), 3 deletions(-)

diff --git a/include/openPMD/CustomHierarchy.hpp b/include/openPMD/CustomHierarchy.hpp
index 59fb368e07..109c7fcd41 100644
--- a/include/openPMD/CustomHierarchy.hpp
+++ b/include/openPMD/CustomHierarchy.hpp
@@ -20,9 +20,12 @@
  */
 #pragma once
 
+#include "openPMD/IO/AbstractIOHandler.hpp"
 #include "openPMD/backend/Container.hpp"
 
 #include <iostream>
+#include <set>
+#include <string>
 #include <type_traits>
 
 namespace openPMD
@@ -30,12 +33,18 @@ namespace openPMD
 class CustomHierarchy;
 namespace internal
 {
+    struct MeshesParticlesPath
+    {
+        std::set<std::string> paths;
+        [[nodiscard]] bool ignore(std::string const &name) const;
+    };
     using CustomHierarchyData = ContainerData<CustomHierarchy>;
-}
+} // namespace internal
 
 class CustomHierarchy : public Container<CustomHierarchy>
 {
     friend class Iteration;
+    friend class Container<CustomHierarchy>;
 
 private:
     using Container_t = Container<CustomHierarchy>;
@@ -51,6 +60,10 @@ class CustomHierarchy : public Container<CustomHierarchy>
         Container_t::setData(std::move(data));
     }
 
+    void read(internal::MeshesParticlesPath const &);
+
+    void flush(std::string const &path, internal::FlushParams const &) override;
+
 public:
     CustomHierarchy(CustomHierarchy const &other) = default;
     CustomHierarchy(CustomHierarchy &&other) = default;
diff --git a/src/CustomHierarchy.cpp b/src/CustomHierarchy.cpp
index e04301b07c..b55f2f0aa5 100644
--- a/src/CustomHierarchy.cpp
+++ b/src/CustomHierarchy.cpp
@@ -20,9 +20,58 @@
  */
 
 #include "openPMD/CustomHierarchy.hpp"
+#include "openPMD/IO/AbstractIOHandler.hpp"
+#include "openPMD/IO/Access.hpp"
+#include "openPMD/IO/IOTask.hpp"
+#include "openPMD/backend/Attributable.hpp"
 
 namespace openPMD
 {
+namespace internal
+{
+    bool MeshesParticlesPath::ignore(const std::string &name) const
+    {
+        return paths.find(name) != paths.end();
+    }
+} // namespace internal
+
 CustomHierarchy::CustomHierarchy() = default;
-CustomHierarchy::CustomHierarchy(NoInit): Container_t(NoInit()) {}
-}
\ No newline at end of file
+CustomHierarchy::CustomHierarchy(NoInit) : Container_t(NoInit())
+{}
+
+void CustomHierarchy::read(internal::MeshesParticlesPath const &mpp)
+{
+    Attributable::readAttributes(ReadMode::FullyReread);
+    Parameter<Operation::LIST_PATHS> pList;
+    IOHandler()->enqueue(IOTask(this, pList));
+    IOHandler()->flush(internal::defaultFlushParams);
+    for (auto const &path : *pList.paths)
+    {
+        if (mpp.ignore(path))
+        {
+            continue;
+        }
+        Parameter<Operation::OPEN_PATH> pOpen;
+        pOpen.path = path;
+        auto subpath = this->operator[](path);
+        IOHandler()->enqueue(IOTask(&subpath, pOpen));
+        subpath.read(mpp);
+    }
+}
+
+void CustomHierarchy::flush(
+    std::string const &path, internal::FlushParams const &flushParams)
+{
+    if (!written() && !path.empty())
+    {
+        Parameter<Operation::CREATE_PATH> pCreate;
+        pCreate.path = path;
+        IOHandler()->enqueue(IOTask(this, pCreate));
+    }
+    for (auto &[name, subpath] : *this)
+    {
+        subpath.flush(name, flushParams);
+    }
+    flushAttributes(flushParams);
+}
+} // namespace openPMD
diff --git a/src/Iteration.cpp b/src/Iteration.cpp
index 5d7f8690e5..43183114e3 100644
--- a/src/Iteration.cpp
+++ b/src/Iteration.cpp
@@ -19,6 +19,7 @@
  * If not, see <http://www.gnu.org/licenses/>.
  */
 #include "openPMD/Iteration.hpp"
+#include "openPMD/CustomHierarchy.hpp"
 #include "openPMD/Dataset.hpp"
 #include "openPMD/Datatype.hpp"
 #include "openPMD/Series.hpp"
@@ -290,8 +291,12 @@ void Iteration::flushVariableBased(
     }
 }
 
+/*
+ * @todo move much of this logic to CustomHierarchy::flush
+ */
 void Iteration::flushIteration(internal::FlushParams const &flushParams)
 {
+    CustomHierarchy::flush("", flushParams);
     if (access::readOnly(IOHandler()->m_frontendAccess))
     {
         for (auto &m : meshes)
@@ -388,6 +393,9 @@ void Iteration::readGorVBased(std::string const &groupPath, bool doBeginStep)
     read_impl(groupPath);
 }
 
+/*
+ * @todo move lots of this to CustomHierarchy::read
+ */
 void Iteration::read_impl(std::string const &groupPath)
 {
     Parameter<Operation::OPEN_PATH> pOpen;
@@ -484,6 +492,21 @@ void Iteration::read_impl(std::string const &groupPath)
         hasParticles = s.containsAttribute("particlesPath");
     }
 
+    {
+        internal::MeshesParticlesPath mpp;
+        if (hasMeshes)
+        {
+            mpp.paths.emplace(auxiliary::trim(
+                s.meshesPath(), [](char const &c) { return c == '/'; }));
+        }
+        if (hasParticles)
+        {
+            mpp.paths.emplace(auxiliary::trim(
+                s.particlesPath(), [](char const &c) { return c == '/'; }));
+        }
+        CustomHierarchy::read(mpp);
+    }
+
     if (hasMeshes)
     {
         try