From 5486f1da8e5f94241a081a0bcdf2e5049b464b7a Mon Sep 17 00:00:00 2001
From: Axel Huebl <axel.huebl@plasma.ninja>
Date: Fri, 17 Jan 2025 16:58:22 -0800
Subject: [PATCH] PC: SoA Name Helpers

Add name to index getters and query (has) functions to SoA
names. Ported over from ImpactX.
---
 Src/Particle/AMReX_ParticleContainer.H  | 32 +++++++++++
 Src/Particle/AMReX_ParticleContainerI.H | 71 +++++++++++++++++++++++--
 2 files changed, 100 insertions(+), 3 deletions(-)

diff --git a/Src/Particle/AMReX_ParticleContainer.H b/Src/Particle/AMReX_ParticleContainer.H
index 2782e74335..82f05a11bd 100644
--- a/Src/Particle/AMReX_ParticleContainer.H
+++ b/Src/Particle/AMReX_ParticleContainer.H
@@ -1453,6 +1453,38 @@ public:
     /** Get the names for the int SoA components **/
     std::vector<std::string> GetIntSoANames () const {return m_soa_idata_names;}
 
+    /** Check if a container has a ParticleReal component
+     *
+     * @param name component name to check
+     * @return true if found, else false
+     */
+    bool HasRealComp (std::string const & name);
+
+    /** Check if a container has an Integer component
+     *
+     * @param name component name to check
+     * @return true if found, else false
+     */
+    bool HasIntComp (std::string const & name);
+
+    /** Get the ParticleReal SoA index of a component
+     *
+     * This throws a runtime exception if the component does not exist.
+     *
+     * @param name component name to query index for
+     * @return zero-based index
+     */
+    int GetRealCompIndex (std::string const & name);
+
+    /** Get the Integer SoA index of a component
+     *
+     * This throws a runtime exception if the component does not exist.
+     *
+     * @param name component name to query index for
+     * @return zero-based index
+     */
+    int GetIntCompIndex (std::string const & name);
+
 protected:
 
     template <class RTYPE>
diff --git a/Src/Particle/AMReX_ParticleContainerI.H b/Src/Particle/AMReX_ParticleContainerI.H
index 769ab997f3..616d2896e5 100644
--- a/Src/Particle/AMReX_ParticleContainerI.H
+++ b/Src/Particle/AMReX_ParticleContainerI.H
@@ -1,6 +1,9 @@
 #include <AMReX_MakeParticle.H>
 
+#include <algorithm>
+#include <iterator>
 #include <set>
+#include <stdexcept>
 #include <string>
 #include <type_traits>
 #include <vector>
@@ -80,10 +83,15 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
     }
 }
 
-template <typename ParticleType, int NArrayReal, int NArrayInt,
-        template<class> class Allocator, class CellAssignor>
+template<
+    typename ParticleType,
+    int NArrayReal,
+    int NArrayInt,
+    template<class> class Allocator,
+    class CellAssignor
+>
 void
-ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor> :: SetSoACompileTimeNames (
+ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>::SetSoACompileTimeNames (
     std::vector<std::string> const & rdata_name, std::vector<std::string> const & idata_name
 )
 {
@@ -106,6 +114,63 @@ ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssig
     }
 }
 
+template<
+    typename ParticleType,
+    int NArrayReal,
+    int NArrayInt,
+    template<class> class Allocator,
+    class CellAssignor
+>
+bool
+ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>::HasRealComp (std::string const & name)
+{
+    return std::find(m_soa_rdata_names.begin(), m_soa_rdata_names.end(), name) != std::end(m_soa_rdata_names);
+}
+
+template <typename ParticleType, int NArrayReal, int NArrayInt,
+        template<class> class Allocator, class CellAssignor>
+bool
+ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>::HasIntComp (std::string const & name)
+{
+    return std::find(m_soa_idata_names.begin(), m_soa_idata_names.end(), name) != std::end(m_soa_idata_names);
+}
+
+template<
+    typename ParticleType,
+    int NArrayReal,
+    int NArrayInt,
+    template<class> class Allocator,
+    class CellAssignor
+>
+int
+ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>::GetRealCompIndex (std::string const & name)
+{
+    const auto it = std::find(m_soa_rdata_names.begin(), m_soa_rdata_names.end(), name);
+
+    if (it == m_soa_rdata_names.end())
+        throw std::runtime_error("GetRealCompIndex: Component " + name + " does not exist!");
+    else
+        return std::distance(m_soa_rdata_names.begin(), it);
+}
+
+template<
+    typename ParticleType,
+    int NArrayReal,
+    int NArrayInt,
+    template<class> class Allocator,
+    class CellAssignor
+>
+int
+ParticleContainer_impl<ParticleType, NArrayReal, NArrayInt, Allocator, CellAssignor>::GetIntCompIndex (std::string const & name)
+{
+    const auto it = std::find(m_soa_idata_names.begin(), m_soa_idata_names.end(), name);
+
+    if (it == m_soa_idata_names.end())
+        throw std::runtime_error("GetIntCompIndex: Component " + name + " does not exist!");
+    else
+        return std::distance(m_soa_idata_names.begin(), it);
+}
+
 template <typename ParticleType, int NArrayReal, int NArrayInt,
           template<class> class Allocator, class CellAssignor>
 template <typename P, typename Assignor>