Particles: support removing particles during the simulation

This still cannot be controlled by the user. Currently, all particles are
killed after two seconds

4 files changed, 193 insertions, 76 deletions

diff --git a/source/blender/simulation/intern/particle_allocator.hh b/source/blender/simulation/intern/particle_allocator.hh
index b742459b3c2..1e7578a75ed 100644
--- a/source/blender/simulation/intern/particle_allocator.hh
+++ b/source/blender/simulation/intern/particle_allocator.hh
@@ -77,6 +77,11 @@ class ParticleAllocator : NonCopyable, NonMovable {
   {
   }
 
+  const fn::AttributesInfo &attributes_info() const
+  {
+    return attributes_allocator_.attributes_info();
+  }
+
   Span<fn::MutableAttributesRef> get_allocations() const
   {
     return attributes_allocator_.get_allocations();
diff --git a/source/blender/simulation/intern/simulation_collect_influences.cc b/source/blender/simulation/intern/simulation_collect_influences.cc
index c1b936d9aa3..d4161b06a00 100644
--- a/source/blender/simulation/intern/simulation_collect_influences.cc
+++ b/source/blender/simulation/intern/simulation_collect_influences.cc
@@ -278,10 +278,12 @@ class MyBasicEmitter : public ParticleEmitter {
 
     MutableSpan<float3> positions = attributes.get<float3>("Position");
     MutableSpan<float3> velocities = attributes.get<float3>("Velocity");
+    MutableSpan<float> birth_times = attributes.get<float>("Birth Time");
 
     for (int i : IndexRange(attributes.size())) {
       positions[i] = rng.get_unit_float3();
       velocities[i] = rng.get_unit_float3();
+      birth_times[i] = context.simulation_time_interval().start();
     }
   }
 };
@@ -307,6 +309,9 @@ static void prepare_particle_attribute_builders(nodes::MFNetworkTreeMap &network
     builder.add<float3>("Position", {0, 0, 0});
     builder.add<float3>("Velocity", {0, 0, 0});
     builder.add<int>("ID", 0);
+    /* TODO: Use bool property, but need to add CD_PROP_BOOL first. */
+    builder.add<int>("Dead", 0);
+    builder.add<float>("Birth Time", 0.0f);
     r_influences.particle_attributes_builder.add_new(std::move(name), &builder);
   }
 }
diff --git a/source/blender/simulation/intern/simulation_solver.cc b/source/blender/simulation/intern/simulation_solver.cc
index 8460b9e148e..f542d6ab2c8 100644
--- a/source/blender/simulation/intern/simulation_solver.cc
+++ b/source/blender/simulation/intern/simulation_solver.cc
@@ -125,6 +125,124 @@ static void ensure_attributes_exist(ParticleSimulationState *state, const fn::At
   }
 }
 
+BLI_NOINLINE static void simulate_existing_particles(SimulationSolveContext &solve_context,
+                                                     ParticleSimulationState &state,
+                                                     const fn::AttributesInfo &attributes_info)
+{
+  CustomDataAttributesRef custom_data_attributes{
+      state.attributes, state.tot_particles, attributes_info};
+  fn::MutableAttributesRef attributes = custom_data_attributes;
+
+  Array<float3> force_vectors{state.tot_particles, {0, 0, 0}};
+  const Vector<const ParticleForce *> *forces =
+      solve_context.influences().particle_forces.lookup_ptr(state.head.name);
+
+  if (forces != nullptr) {
+    ParticleChunkContext particle_chunk_context{IndexMask(state.tot_particles), attributes};
+    ParticleForceContext particle_force_context{
+        solve_context, particle_chunk_context, force_vectors};
+
+    for (const ParticleForce *force : *forces) {
+      force->add_force(particle_force_context);
+    }
+  }
+
+  MutableSpan<float3> positions = attributes.get<float3>("Position");
+  MutableSpan<float3> velocities = attributes.get<float3>("Velocity");
+  MutableSpan<float> birth_times = attributes.get<float>("Birth Time");
+  MutableSpan<int> dead_states = attributes.get<int>("Dead");
+  float end_time = solve_context.solve_interval().end();
+  float time_step = solve_context.solve_interval().duration();
+  for (int i : positions.index_range()) {
+    velocities[i] += force_vectors[i] * time_step;
+    positions[i] += velocities[i] * time_step;
+
+    if (end_time - birth_times[i] > 2) {
+      dead_states[i] = true;
+    }
+  }
+}
+
+BLI_NOINLINE static void run_emitters(SimulationSolveContext &solve_context,
+                                      ParticleAllocators &particle_allocators)
+{
+  for (const ParticleEmitter *emitter : solve_context.influences().particle_emitters) {
+    ParticleEmitterContext emitter_context{
+        solve_context, particle_allocators, solve_context.solve_interval()};
+    emitter->emit(emitter_context);
+  }
+}
+
+BLI_NOINLINE static int count_particles_after_time_step(ParticleSimulationState &state,
+                                                        ParticleAllocator &allocator)
+{
+  CustomDataAttributesRef custom_data_attributes{
+      state.attributes, state.tot_particles, allocator.attributes_info()};
+  fn::MutableAttributesRef attributes = custom_data_attributes;
+  int new_particle_amount = attributes.get<int>("Dead").count(0);
+
+  for (fn::MutableAttributesRef emitted_attributes : allocator.get_allocations()) {
+    new_particle_amount += emitted_attributes.get<int>("Dead").count(0);
+  }
+
+  return new_particle_amount;
+}
+
+BLI_NOINLINE static void remove_dead_and_add_new_particles(ParticleSimulationState &state,
+                                                           ParticleAllocator &allocator)
+{
+  const int new_particle_amount = count_particles_after_time_step(state, allocator);
+
+  CustomDataAttributesRef custom_data_attributes{
+      state.attributes, state.tot_particles, allocator.attributes_info()};
+
+  Vector<fn::MutableAttributesRef> particle_sources;
+  particle_sources.append(custom_data_attributes);
+  particle_sources.extend(allocator.get_allocations());
+
+  CustomDataLayer *dead_layer = nullptr;
+
+  for (CustomDataLayer &layer : MutableSpan(state.attributes.layers, state.attributes.totlayer)) {
+    StringRefNull name = layer.name;
+    if (name == "Dead") {
+      dead_layer = &layer;
+      continue;
+    }
+    const fn::CPPType &cpp_type = custom_to_cpp_data_type((CustomDataType)layer.type);
+    fn::GMutableSpan new_buffer{
+        cpp_type,
+        MEM_mallocN_aligned(new_particle_amount * cpp_type.size(), cpp_type.alignment(), AT),
+        new_particle_amount};
+
+    int current = 0;
+    for (fn::MutableAttributesRef attributes : particle_sources) {
+      Span<int> dead_states = attributes.get<int>("Dead");
+      fn::GSpan source_buffer = attributes.get(name);
+      BLI_assert(source_buffer.type() == cpp_type);
+      for (int i : attributes.index_range()) {
+        if (dead_states[i] == 0) {
+          cpp_type.copy_to_uninitialized(source_buffer[i], new_buffer[current]);
+          current++;
+        }
+      }
+    }
+
+    if (layer.data != nullptr) {
+      MEM_freeN(layer.data);
+    }
+    layer.data = new_buffer.buffer();
+  }
+
+  BLI_assert(dead_layer != nullptr);
+  if (dead_layer->data != nullptr) {
+    MEM_freeN(dead_layer->data);
+  }
+  dead_layer->data = MEM_callocN(sizeof(int) * new_particle_amount, AT);
+
+  state.tot_particles = new_particle_amount;
+  state.next_particle_id += allocator.total_allocated();
+}
+
 void initialize_simulation_states(Simulation &simulation,
                                   Depsgraph &UNUSED(depsgraph),
                                   const SimulationInfluences &UNUSED(influences))
@@ -137,89 +255,47 @@ void solve_simulation_time_step(Simulation &simulation,
                                 const SimulationInfluences &influences,
                                 float time_step)
 {
-  SimulationSolveContext solve_context{simulation, depsgraph, influences};
+  SimulationSolveContext solve_context{
+      simulation,
+      depsgraph,
+      influences,
+      TimeInterval(simulation.current_simulation_time, time_step)};
   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);
+    }
+  }
+
   Map<std::string, std::unique_ptr<fn::AttributesInfo>> attribute_infos;
-  Map<std::string, std::unique_ptr<ParticleAllocator>> particle_allocators;
-  LISTBASE_FOREACH (ParticleSimulationState *, state, &simulation.states) {
+  Map<std::string, std::unique_ptr<ParticleAllocator>> particle_allocators_map;
+  for (ParticleSimulationState *state : particle_simulation_states) {
     const fn::AttributesInfoBuilder &builder = *influences.particle_attributes_builder.lookup_as(
         state->head.name);
     auto info = std::make_unique<fn::AttributesInfo>(builder);
 
     ensure_attributes_exist(state, *info);
 
-    particle_allocators.add_new(
+    particle_allocators_map.add_new(
         state->head.name, std::make_unique<ParticleAllocator>(*info, state->next_particle_id));
     attribute_infos.add_new(state->head.name, std::move(info));
   }
 
-  LISTBASE_FOREACH (ParticleSimulationState *, state, &simulation.states) {
+  ParticleAllocators particle_allocators{particle_allocators_map};
 
+  for (ParticleSimulationState *state : particle_simulation_states) {
     const fn::AttributesInfo &attributes_info = *attribute_infos.lookup_as(state->head.name);
-    CustomDataAttributesRef custom_data_attributes{
-        state->attributes, state->tot_particles, attributes_info};
-    fn::MutableAttributesRef attributes = custom_data_attributes;
-
-    MutableSpan<float3> positions = attributes.get<float3>("Position");
-    MutableSpan<float3> velocities = attributes.get<float3>("Velocity");
-
-    Array<float3> force_vectors{state->tot_particles, {0, 0, 0}};
-    const Vector<const ParticleForce *> *forces = influences.particle_forces.lookup_ptr(
-        state->head.name);
-
-    if (forces != nullptr) {
-      ParticleChunkContext particle_chunk_context{IndexMask(state->tot_particles), attributes};
-      ParticleForceContext particle_force_context{
-          solve_context, particle_chunk_context, force_vectors};
-
-      for (const ParticleForce *force : *forces) {
-        force->add_force(particle_force_context);
-      }
-    }
-
-    for (int i : positions.index_range()) {
-      velocities[i] += force_vectors[i] * time_step;
-      positions[i] += velocities[i] * time_step;
-    }
+    simulate_existing_particles(solve_context, *state, attributes_info);
   }
 
-  for (const ParticleEmitter *emitter : influences.particle_emitters) {
-    ParticleEmitterContext emitter_context{
-        solve_context, particle_allocators, simulation_time_interval};
-    emitter->emit(emitter_context);
-  }
-
-  LISTBASE_FOREACH (ParticleSimulationState *, state, &simulation.states) {
-    const fn::AttributesInfo &attributes_info = *attribute_infos.lookup_as(state->head.name);
-    ParticleAllocator &allocator = *particle_allocators.lookup_as(state->head.name);
-
-    const int emitted_particle_amount = allocator.total_allocated();
-    const int old_particle_amount = state->tot_particles;
-    const int new_particle_amount = old_particle_amount + emitted_particle_amount;
-
-    CustomData_realloc(&state->attributes, new_particle_amount);
-
-    CustomDataAttributesRef custom_data_attributes{
-        state->attributes, new_particle_amount, attributes_info};
-    fn::MutableAttributesRef attributes = custom_data_attributes;
-
-    int offset = old_particle_amount;
-    for (fn::MutableAttributesRef emitted_attributes : allocator.get_allocations()) {
-      fn::MutableAttributesRef dst_attributes = attributes.slice(
-          IndexRange(offset, emitted_attributes.size()));
-      for (int attribute_index : attributes.info().index_range()) {
-        fn::GMutableSpan emitted_data = emitted_attributes.get(attribute_index);
-        fn::GMutableSpan dst = dst_attributes.get(attribute_index);
-        const fn::CPPType &type = dst.type();
-        type.copy_to_uninitialized_n(
-            emitted_data.buffer(), dst.buffer(), emitted_attributes.size());
-      }
-      offset += emitted_attributes.size();
-    }
+  run_emitters(solve_context, particle_allocators);
 
-    state->tot_particles = new_particle_amount;
-    state->next_particle_id += emitted_particle_amount;
+  for (ParticleSimulationState *state : particle_simulation_states) {
+    ParticleAllocator &allocator = *particle_allocators.try_get_allocator(state->head.name);
+    remove_dead_and_add_new_particles(*state, allocator);
   }
 
   simulation.current_simulation_time = simulation_time_interval.end();
diff --git a/source/blender/simulation/intern/simulation_solver.hh b/source/blender/simulation/intern/simulation_solver.hh
index 7b37ed0583d..b5e42b53846 100644
--- a/source/blender/simulation/intern/simulation_solver.hh
+++ b/source/blender/simulation/intern/simulation_solver.hh
@@ -57,14 +57,51 @@ class SimulationSolveContext {
   Simulation &simulation_;
   Depsgraph &depsgraph_;
   const SimulationInfluences &influences_;
+  TimeInterval solve_interval_;
 
  public:
   SimulationSolveContext(Simulation &simulation,
                          Depsgraph &depsgraph,
-                         const SimulationInfluences &influences)
-      : simulation_(simulation), depsgraph_(depsgraph), influences_(influences)
+                         const SimulationInfluences &influences,
+                         TimeInterval solve_interval)
+      : simulation_(simulation),
+        depsgraph_(depsgraph),
+        influences_(influences),
+        solve_interval_(solve_interval)
   {
   }
+
+  TimeInterval solve_interval() const
+  {
+    return solve_interval_;
+  }
+
+  const SimulationInfluences &influences() const
+  {
+    return influences_;
+  }
+};
+
+class ParticleAllocators {
+ private:
+  Map<std::string, std::unique_ptr<ParticleAllocator>> &allocators_;
+
+ public:
+  ParticleAllocators(Map<std::string, std::unique_ptr<ParticleAllocator>> &allocators)
+      : allocators_(allocators)
+  {
+  }
+
+  ParticleAllocator *try_get_allocator(StringRef particle_simulation_name)
+  {
+    auto *ptr = allocators_.lookup_ptr_as(particle_simulation_name);
+    if (ptr != nullptr) {
+      return ptr->get();
+    }
+    else {
+      return nullptr;
+    }
+  }
 };
 
 class ParticleChunkContext {
@@ -97,12 +134,12 @@ class ParticleChunkContext {
 class ParticleEmitterContext {
  private:
   SimulationSolveContext &solve_context_;
-  Map<std::string, std::unique_ptr<ParticleAllocator>> &particle_allocators_;
+  ParticleAllocators &particle_allocators_;
   TimeInterval simulation_time_interval_;
 
  public:
   ParticleEmitterContext(SimulationSolveContext &solve_context,
-                         Map<std::string, std::unique_ptr<ParticleAllocator>> &particle_allocators,
+                         ParticleAllocators &particle_allocators,
                          TimeInterval simulation_time_interval)
       : solve_context_(solve_context),
         particle_allocators_(particle_allocators),
@@ -112,13 +149,7 @@ class ParticleEmitterContext {
 
   ParticleAllocator *try_get_particle_allocator(StringRef particle_simulation_name)
   {
-    auto *ptr = particle_allocators_.lookup_ptr_as(particle_simulation_name);
-    if (ptr != nullptr) {
-      return ptr->get();
-    }
-    else {
-      return nullptr;
-    }
+    return particle_allocators_.try_get_allocator(particle_simulation_name);
   }
 
   TimeInterval simulation_time_interval() const