1 files changed, 57 insertions, 0 deletions
diff --git a/source/blender/simulation/intern/simulation_solver.hh b/source/blender/simulation/intern/simulation_solver.hh
index 1d1b9935661..2dc0921ad9e 100644
--- a/source/blender/simulation/intern/simulation_solver.hh
+++ b/source/blender/simulation/intern/simulation_solver.hh
@@ -25,6 +25,7 @@
 #include "FN_attributes_ref.hh"
 
 #include "BKE_persistent_data_handle.hh"
+#include "BKE_simulation.h"
 
 #include "particle_allocator.hh"
 #include "time_interval.hh"
@@ -55,12 +56,61 @@ struct SimulationInfluences {
   VectorSet<ID *> used_data_blocks;
 };
 
+class SimulationStateMap {
+ private:
+  Map<StringRefNull, SimulationState *> states_by_name_;
+  Map<StringRefNull, Vector<SimulationState *>> states_by_type_;
+
+ public:
+  void add(SimulationState *state)
+  {
+    states_by_name_.add_new(state->name, state);
+    states_by_type_.lookup_or_add_default(state->type).append(state);
+  }
+
+  template<typename StateType> StateType *lookup(StringRef name) const
+  {
+    const char *type = BKE_simulation_get_state_type_name<StateType>();
+    return (StateType *)this->lookup_name_type(name, type);
+  }
+
+  template<typename StateType> Span<StateType *> lookup() const
+  {
+    const char *type = BKE_simulation_get_state_type_name<StateType>();
+    return this->lookup_type(type).cast<StateType *>();
+  }
+
+  SimulationState *lookup_name_type(StringRef name, StringRef type) const
+  {
+    SimulationState *state = states_by_name_.lookup_default_as(name, nullptr);
+    if (state == nullptr) {
+      return nullptr;
+    }
+    if (state->type == type) {
+      return state;
+    }
+    return nullptr;
+  }
+
+  Span<SimulationState *> lookup_type(StringRef type) const
+  {
+    const Vector<SimulationState *> *states = states_by_type_.lookup_ptr_as(type);
+    if (states == nullptr) {
+      return {};
+    }
+    else {
+      return states->as_span();
+    }
+  }
+};
+
 class SimulationSolveContext {
  private:
   Simulation &simulation_;
   Depsgraph &depsgraph_;
   const SimulationInfluences &influences_;
   TimeInterval solve_interval_;
+  const SimulationStateMap &state_map_;
   const bke::PersistentDataHandleMap &id_handle_map_;
 
  public:
@@ -68,11 +118,13 @@ class SimulationSolveContext {
                          Depsgraph &depsgraph,
                          const SimulationInfluences &influences,
                          TimeInterval solve_interval,
+                         const SimulationStateMap &state_map,
                          const bke::PersistentDataHandleMap &handle_map)
       : simulation_(simulation),
         depsgraph_(depsgraph),
         influences_(influences),
         solve_interval_(solve_interval),
+        state_map_(state_map),
         id_handle_map_(handle_map)
   {
   }
@@ -91,6 +143,11 @@ class SimulationSolveContext {
   {
     return id_handle_map_;
   }
+
+  const SimulationStateMap &state_map() const
+  {
+    return state_map_;
+  }
 };
 
 class ParticleAllocators {