Particles: give emitter its own state

High quality emitters need to maintain state themselves. For example,
this it needs to remember when it spawned the last particle.
This is especially important when the birth rate is changing over time.
Otherwise, there will be very visible artifacts.

It is quite likely that other components of the simulation need their own
state as well. Therefore, I refactored the `SimulationState` type a bit,
to make it more extensible. Instead of using hardcoded type numbers, a
string is used to identify the state type. Also, instead of having switch
statements in many places, there is a new `SimulationStateType` that
encapsulates information about how a specific state is created/freed/copied/...

I removed the integration with the point cache for now, because it was
not used anyway in it's current state.

5 files changed, 159 insertions, 84 deletions

diff --git a/source/blender/simulation/intern/simulation_collect_influences.cc b/source/blender/simulation/intern/simulation_collect_influences.cc
index 4760063daca..764e587d157 100644
--- a/source/blender/simulation/intern/simulation_collect_influences.cc
+++ b/source/blender/simulation/intern/simulation_collect_influences.cc
@@ -280,17 +280,30 @@ static void collect_forces(nodes::MFNetworkTreeMap &network_map,
 class MyBasicEmitter : public ParticleEmitter {
  private:
   Array<std::string> names_;
+  std::string my_state_;
   const fn::MultiFunction &inputs_fn_;
   uint32_t seed_;
 
  public:
-  MyBasicEmitter(Array<std::string> names, const fn::MultiFunction &inputs_fn, uint32_t seed)
-      : names_(std::move(names)), inputs_fn_(inputs_fn), seed_(seed)
+  MyBasicEmitter(Array<std::string> names,
+                 std::string my_state,
+                 const fn::MultiFunction &inputs_fn,
+                 uint32_t seed)
+      : names_(std::move(names)),
+        my_state_(std::move(my_state)),
+        inputs_fn_(inputs_fn),
+        seed_(seed)
   {
   }
 
   void emit(ParticleEmitterContext &context) const override
   {
+    auto *state = context.solve_context().state_map().lookup<ParticleMeshEmitterSimulationState>(
+        my_state_);
+    if (state == nullptr) {
+      return;
+    }
+
     fn::MFContextBuilder mf_context;
     mf_context.add_global_context("PersistentDataHandleMap",
                                   &context.solve_context().handle_map());
@@ -311,15 +324,17 @@ class MyBasicEmitter : public ParticleEmitter {
     Vector<float3> new_velocities;
     Vector<float> new_birth_times;
 
-    float start_time = context.simulation_time_interval().start();
+    TimeInterval time_interval = context.simulation_time_interval();
+    float start_time = time_interval.start();
     RandomNumberGenerator rng{(*(uint32_t *)&start_time) ^ seed_};
 
-    int amount = rate * 10;
-    for (int i : IndexRange(amount)) {
-      UNUSED_VARS(i);
+    const float time_between_particles = 1.0f / rate;
+    while (state->last_birth_time + time_between_particles < time_interval.end()) {
       new_positions.append(rng.get_unit_float3() * 0.3 + float3(object->loc));
       new_velocities.append(rng.get_unit_float3());
-      new_birth_times.append(context.simulation_time_interval().start());
+      const float birth_time = state->last_birth_time + time_between_particles;
+      new_birth_times.append(birth_time);
+      state->last_birth_time = birth_time;
     }
 
     for (StringRef name : names_) {
@@ -328,6 +343,7 @@ class MyBasicEmitter : public ParticleEmitter {
         return;
       }
 
+      int amount = new_positions.size();
       fn::MutableAttributesRef attributes = allocator->allocate(amount);
 
       initialized_copy_n(new_positions.data(), amount, attributes.get<float3>("Position").data());
@@ -352,7 +368,8 @@ static Vector<const nodes::DNode *> find_linked_particle_simulations(
 
 static ParticleEmitter *create_particle_emitter(const nodes::DNode &dnode,
                                                 ResourceCollector &resources,
-                                                nodes::MFNetworkTreeMap &network_map)
+                                                nodes::MFNetworkTreeMap &network_map,
+                                                RequiredStates &r_required_states)
 {
   Vector<const nodes::DNode *> simulation_dnodes = find_linked_particle_simulations(
       dnode.output(0));
@@ -377,19 +394,23 @@ static ParticleEmitter *create_particle_emitter(const nodes::DNode &dnode,
   fn::MultiFunction &inputs_fn = resources.construct<fn::MFNetworkEvaluator>(
       AT, Span<const fn::MFOutputSocket *>(), input_sockets.as_span());
 
-  uint32_t seed = DefaultHash<std::string>{}(dnode_to_path(dnode));
+  std::string my_state_name = dnode_to_path(dnode);
+  r_required_states.add(my_state_name, SIM_TYPE_NAME_PARTICLE_MESH_EMITTER);
+  uint32_t seed = DefaultHash<std::string>{}(my_state_name);
   ParticleEmitter &emitter = resources.construct<MyBasicEmitter>(
-      AT, std::move(names), inputs_fn, seed);
+      AT, std::move(names), std::move(my_state_name), inputs_fn, seed);
   return &emitter;
 }
 
 static void collect_emitters(nodes::MFNetworkTreeMap &network_map,
                              ResourceCollector &resources,
-                             SimulationInfluences &r_influences)
+                             SimulationInfluences &r_influences,
+                             RequiredStates &r_required_states)
 {
   for (const nodes::DNode *dnode :
        network_map.tree().nodes_by_type("SimulationNodeParticleMeshEmitter")) {
-    ParticleEmitter *emitter = create_particle_emitter(*dnode, resources, network_map);
+    ParticleEmitter *emitter = create_particle_emitter(
+        *dnode, resources, network_map, r_required_states);
     if (emitter != nullptr) {
       r_influences.particle_emitters.append(emitter);
     }
@@ -433,7 +454,7 @@ static void find_used_data_blocks(const nodes::DerivedNodeTree &tree,
 void collect_simulation_influences(Simulation &simulation,
                                    ResourceCollector &resources,
                                    SimulationInfluences &r_influences,
-                                   SimulationStatesInfo &r_states_info)
+                                   RequiredStates &r_required_states)
 {
   nodes::NodeTreeRefMap tree_refs;
   const nodes::DerivedNodeTree tree{simulation.nodetree, tree_refs};
@@ -452,10 +473,10 @@ void collect_simulation_influences(Simulation &simulation,
   // WM_clipboard_text_set(network.to_dot().c_str(), false);
 
   collect_forces(network_map, resources, data_sources, r_influences);
-  collect_emitters(network_map, resources, r_influences);
+  collect_emitters(network_map, resources, r_influences, r_required_states);
 
   for (const nodes::DNode *dnode : get_particle_simulation_nodes(tree)) {
-    r_states_info.particle_simulation_names.add(dnode_to_path(*dnode));
+    r_required_states.add(dnode_to_path(*dnode), SIM_TYPE_NAME_PARTICLE_SIMULATION);
   }
 
   find_used_data_blocks(tree, r_influences);
diff --git a/source/blender/simulation/intern/simulation_collect_influences.hh b/source/blender/simulation/intern/simulation_collect_influences.hh
index cca396ffa66..42cbea6977e 100644
--- a/source/blender/simulation/intern/simulation_collect_influences.hh
+++ b/source/blender/simulation/intern/simulation_collect_influences.hh
@@ -25,14 +25,43 @@
 
 namespace blender::sim {
 
-struct SimulationStatesInfo {
-  VectorSet<std::string> particle_simulation_names;
+class RequiredStates {
+ private:
+  Map<std::string, const char *> state_type_by_state_name_;
+
+ public:
+  void add(std::string state_name, const char *state_type)
+  {
+    BLI_assert(state_type != nullptr);
+    const char *type_name = state_type_by_state_name_.lookup_default(state_name, nullptr);
+    if (type_name != nullptr) {
+      if (!STREQ(state_type, type_name)) {
+        std::cout << "Warning: Tried to have two different states with the same name.\n";
+        std::cout << "    Name: " << state_name << "\n";
+        std::cout << "    Type 1: " << state_type << "\n";
+        std::cout << "    Type 2: " << type_name << "\n";
+      }
+      return;
+    }
+
+    state_type_by_state_name_.add(std::move(state_name), state_type);
+  }
+
+  const Map<std::string, const char *> &states() const
+  {
+    return state_type_by_state_name_;
+  }
+
+  bool is_required(StringRef state_name, StringRef state_type) const
+  {
+    return state_type_by_state_name_.lookup_default_as(state_name, "") == state_type;
+  }
 };
 
 void collect_simulation_influences(Simulation &simulation,
                                    ResourceCollector &resources,
                                    SimulationInfluences &r_influences,
-                                   SimulationStatesInfo &r_states_info);
+                                   RequiredStates &r_required_states);
 
 }  // namespace blender::sim
 
diff --git a/source/blender/simulation/intern/simulation_solver.cc b/source/blender/simulation/intern/simulation_solver.cc
index 158c50957df..ee7a8d40035 100644
--- a/source/blender/simulation/intern/simulation_solver.cc
+++ b/source/blender/simulation/intern/simulation_solver.cc
@@ -312,20 +312,21 @@ void solve_simulation_time_step(Simulation &simulation,
     handle_map.add(handle->handle, *handle->id);
   }
 
+  SimulationStateMap state_map;
+  LISTBASE_FOREACH (SimulationState *, state, &simulation.states) {
+    state_map.add(state);
+  }
+
   SimulationSolveContext solve_context{simulation,
                                        depsgraph,
                                        influences,
                                        TimeInterval(simulation.current_simulation_time, time_step),
+                                       state_map,
                                        handle_map};
   TimeInterval simulation_time_interval{simulation.current_simulation_time, time_step};
 
-  Vector<SimulationState *> simulation_states{simulation.states};
-  Vector<ParticleSimulationState *> particle_simulation_states;
-  for (SimulationState *state : simulation_states) {
-    if (state->type == SIM_STATE_TYPE_PARTICLES) {
-      particle_simulation_states.append((ParticleSimulationState *)state);
-    }
-  }
+  Span<ParticleSimulationState *> particle_simulation_states =
+      state_map.lookup<ParticleSimulationState>();
 
   Map<std::string, std::unique_ptr<fn::AttributesInfo>> attribute_infos;
   Map<std::string, std::unique_ptr<ParticleAllocator>> particle_allocators_map;
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 {
diff --git a/source/blender/simulation/intern/simulation_update.cc b/source/blender/simulation/intern/simulation_update.cc
index 5328304863b..09219e0238f 100644
--- a/source/blender/simulation/intern/simulation_update.cc
+++ b/source/blender/simulation/intern/simulation_update.cc
@@ -37,88 +37,55 @@
 
 namespace blender::sim {
 
-static void copy_states_to_cow(Simulation *simulation_orig, Simulation *simulation_cow)
+static void copy_states_to_cow(const Simulation *simulation_orig, Simulation *simulation_cow)
 {
   BKE_simulation_state_remove_all(simulation_cow);
   simulation_cow->current_frame = simulation_orig->current_frame;
 
-  LISTBASE_FOREACH (SimulationState *, state_orig, &simulation_orig->states) {
-    switch ((eSimulationStateType)state_orig->type) {
-      case SIM_STATE_TYPE_PARTICLES: {
-        ParticleSimulationState *particle_state_orig = (ParticleSimulationState *)state_orig;
-        ParticleSimulationState *particle_state_cow = (ParticleSimulationState *)
-            BKE_simulation_state_add(simulation_cow, SIM_STATE_TYPE_PARTICLES, state_orig->name);
-        particle_state_cow->tot_particles = particle_state_orig->tot_particles;
-        CustomData_copy(&particle_state_orig->attributes,
-                        &particle_state_cow->attributes,
-                        CD_MASK_ALL,
-                        CD_DUPLICATE,
-                        particle_state_orig->tot_particles);
-        break;
-      }
-    }
+  LISTBASE_FOREACH (const SimulationState *, state_orig, &simulation_orig->states) {
+    SimulationState *state_cow = BKE_simulation_state_add(
+        simulation_cow, state_orig->type, state_orig->name);
+    BKE_simulation_state_copy_data(state_orig, state_cow);
   }
 }
 
-static void remove_unused_states(Simulation *simulation, const VectorSet<std::string> &state_names)
+static void remove_unused_states(Simulation *simulation, const RequiredStates &required_states)
 {
   LISTBASE_FOREACH_MUTABLE (SimulationState *, state, &simulation->states) {
-    if (!state_names.contains(state->name)) {
+    if (!required_states.is_required(state->name, state->type)) {
       BKE_simulation_state_remove(simulation, state);
     }
   }
 }
 
-static void reset_states(Simulation *simulation)
+static void add_missing_states(Simulation *simulation, const RequiredStates &required_states)
 {
-  LISTBASE_FOREACH (SimulationState *, state, &simulation->states) {
-    switch ((eSimulationStateType)state->type) {
-      case SIM_STATE_TYPE_PARTICLES: {
-        ParticleSimulationState *particle_state = (ParticleSimulationState *)state;
-        CustomData_free(&particle_state->attributes, particle_state->tot_particles);
-        particle_state->tot_particles = 0;
-        break;
-      }
-    }
-  }
-}
+  for (auto &&item : required_states.states().items()) {
+    const char *name = item.key.c_str();
+    const char *type = item.value;
 
-static SimulationState *try_find_state_by_name(Simulation *simulation, StringRef name)
-{
-  LISTBASE_FOREACH (SimulationState *, state, &simulation->states) {
-    if (state->name == name) {
-      return state;
-    }
-  }
-  return nullptr;
-}
+    SimulationState *state = BKE_simulation_state_try_find_by_name_and_type(
+        simulation, name, type);
 
-static void add_missing_particle_states(Simulation *simulation, Span<std::string> state_names)
-{
-  for (StringRefNull name : state_names) {
-    SimulationState *state = try_find_state_by_name(simulation, name);
-    if (state != nullptr) {
-      BLI_assert(state->type == SIM_STATE_TYPE_PARTICLES);
-      continue;
+    if (state == nullptr) {
+      BKE_simulation_state_add(simulation, type, name);
     }
-
-    BKE_simulation_state_add(simulation, SIM_STATE_TYPE_PARTICLES, name.c_str());
   }
 }
 
 static void reinitialize_empty_simulation_states(Simulation *simulation,
-                                                 const SimulationStatesInfo &states_info)
+                                                 const RequiredStates &required_states)
 {
-  remove_unused_states(simulation, states_info.particle_simulation_names);
-  reset_states(simulation);
-  add_missing_particle_states(simulation, states_info.particle_simulation_names);
+  remove_unused_states(simulation, required_states);
+  BKE_simulation_state_reset_all(simulation);
+  add_missing_states(simulation, required_states);
 }
 
 static void update_simulation_state_list(Simulation *simulation,
-                                         const SimulationStatesInfo &states_info)
+                                         const RequiredStates &required_states)
 {
-  remove_unused_states(simulation, states_info.particle_simulation_names);
-  add_missing_particle_states(simulation, states_info.particle_simulation_names);
+  remove_unused_states(simulation, required_states);
+  add_missing_states(simulation, required_states);
 }
 
 void update_simulation_in_depsgraph(Depsgraph *depsgraph,
@@ -139,13 +106,13 @@ void update_simulation_in_depsgraph(Depsgraph *depsgraph,
 
   ResourceCollector resources;
   SimulationInfluences influences;
-  SimulationStatesInfo states_info;
+  RequiredStates required_states;
 
   /* TODO: Use simulation_cow, but need to add depsgraph relations before that. */
-  collect_simulation_influences(*simulation_orig, resources, influences, states_info);
+  collect_simulation_influences(*simulation_orig, resources, influences, required_states);
 
   if (current_frame == 1) {
-    reinitialize_empty_simulation_states(simulation_orig, states_info);
+    reinitialize_empty_simulation_states(simulation_orig, required_states);
 
     initialize_simulation_states(*simulation_orig, *depsgraph, influences);
     simulation_orig->current_frame = 1;
@@ -153,7 +120,7 @@ void update_simulation_in_depsgraph(Depsgraph *depsgraph,
     copy_states_to_cow(simulation_orig, simulation_cow);
   }
   else if (current_frame == simulation_orig->current_frame + 1) {
-    update_simulation_state_list(simulation_orig, states_info);
+    update_simulation_state_list(simulation_orig, required_states);
 
     float time_step = 1.0f / 24.0f;
     solve_simulation_time_step(*simulation_orig, *depsgraph, influences, time_step);