Particles: Make it easier to add attributes internally

3 files changed, 110 insertions, 55 deletions

diff --git a/source/blender/simulation/intern/simulation_collect_influences.cc b/source/blender/simulation/intern/simulation_collect_influences.cc
index ec1e7db0f72..98b055802c9 100644
--- a/source/blender/simulation/intern/simulation_collect_influences.cc
+++ b/source/blender/simulation/intern/simulation_collect_influences.cc
@@ -45,6 +45,11 @@ static std::string dnode_to_path(const nodes::DNode &dnode)
   return path;
 }
 
+static Span<const nodes::DNode *> get_particle_simulation_nodes(const nodes::DerivedNodeTree &tree)
+{
+  return tree.nodes_by_type("SimulationNodeParticleSimulation");
+}
+
 static std::optional<Array<std::string>> compute_global_string_inputs(
     nodes::MFNetworkTreeMap &network_map, Span<const fn::MFInputSocket *> sockets)
 {
@@ -243,8 +248,7 @@ static void collect_forces(nodes::MFNetworkTreeMap &network_map,
                            DummyDataSources &data_sources,
                            SimulationInfluences &r_influences)
 {
-  for (const nodes::DNode *dnode :
-       network_map.tree().nodes_by_type("SimulationNodeParticleSimulation")) {
+  for (const nodes::DNode *dnode : get_particle_simulation_nodes(network_map.tree())) {
     std::string name = dnode_to_path(*dnode);
     Vector<const ParticleForce *> forces = create_forces_for_particle_simulation(
         *dnode, network_map, resources, data_sources);
@@ -286,14 +290,27 @@ static void collect_emitters(nodes::MFNetworkTreeMap &network_map,
                              ResourceCollector &resources,
                              SimulationInfluences &r_influences)
 {
-  for (const nodes::DNode *dnode :
-       network_map.tree().nodes_by_type("SimulationNodeParticleSimulation")) {
+  for (const nodes::DNode *dnode : get_particle_simulation_nodes(network_map.tree())) {
     std::string name = dnode_to_path(*dnode);
     ParticleEmitter &emitter = resources.construct<MyBasicEmitter>(AT, name);
     r_influences.particle_emitters.append(&emitter);
   }
 }
 
+static void prepare_particle_attribute_builders(nodes::MFNetworkTreeMap &network_map,
+                                                ResourceCollector &resources,
+                                                SimulationInfluences &r_influences)
+{
+  for (const nodes::DNode *dnode : get_particle_simulation_nodes(network_map.tree())) {
+    std::string name = dnode_to_path(*dnode);
+    fn::AttributesInfoBuilder &builder = resources.construct<fn::AttributesInfoBuilder>(AT);
+    builder.add<float3>("Position", {0, 0, 0});
+    builder.add<float3>("Velocity", {0, 0, 0});
+    builder.add<int>("ID", 0);
+    r_influences.particle_attributes_builder.add_new(std::move(name), &builder);
+  }
+}
+
 void collect_simulation_influences(Simulation &simulation,
                                    ResourceCollector &resources,
                                    SimulationInfluences &r_influences,
@@ -305,6 +322,8 @@ void collect_simulation_influences(Simulation &simulation,
   fn::MFNetwork &network = resources.construct<fn::MFNetwork>(AT);
   nodes::MFNetworkTreeMap network_map = insert_node_tree_into_mf_network(network, tree, resources);
 
+  prepare_particle_attribute_builders(network_map, resources, r_influences);
+
   DummyDataSources data_sources;
   find_and_deduplicate_particle_attribute_nodes(network_map, data_sources);
 
@@ -316,7 +335,7 @@ void collect_simulation_influences(Simulation &simulation,
   collect_forces(network_map, resources, data_sources, r_influences);
   collect_emitters(network_map, resources, r_influences);
 
-  for (const nodes::DNode *dnode : tree.nodes_by_type("SimulationNodeParticleSimulation")) {
+  for (const nodes::DNode *dnode : get_particle_simulation_nodes(tree)) {
     r_states_info.particle_simulation_names.add(dnode_to_path(*dnode));
   }
 }
diff --git a/source/blender/simulation/intern/simulation_solver.cc b/source/blender/simulation/intern/simulation_solver.cc
index f158bd6e8bf..310261f6c95 100644
--- a/source/blender/simulation/intern/simulation_solver.cc
+++ b/source/blender/simulation/intern/simulation_solver.cc
@@ -30,57 +30,99 @@ ParticleEmitter::~ParticleEmitter()
 {
 }
 
+static CustomDataType cpp_to_custom_data_type(const fn::CPPType &type)
+{
+  if (type.is<float3>()) {
+    return CD_PROP_FLOAT3;
+  }
+  if (type.is<float>()) {
+    return CD_PROP_FLOAT;
+  }
+  if (type.is<int32_t>()) {
+    return CD_PROP_INT32;
+  }
+  BLI_assert(false);
+  return CD_PROP_FLOAT;
+}
+
+static const fn::CPPType &custom_to_cpp_data_type(CustomDataType type)
+{
+  switch (type) {
+    case CD_PROP_FLOAT3:
+      return fn::CPPType::get<float3>();
+    case CD_PROP_FLOAT:
+      return fn::CPPType::get<float>();
+    case CD_PROP_INT32:
+      return fn::CPPType::get<int32_t>();
+    default:
+      BLI_assert(false);
+      return fn::CPPType::get<float>();
+  }
+}
+
 class CustomDataAttributesRef {
  private:
-  Vector<void *> buffers_;
+  Array<void *> buffers_;
   uint size_;
-  std::unique_ptr<fn::AttributesInfo> info_;
+  const fn::AttributesInfo &info_;
 
  public:
-  CustomDataAttributesRef(CustomData &custom_data, uint size)
+  CustomDataAttributesRef(CustomData &custom_data, uint size, const fn::AttributesInfo &info)
+      : buffers_(info.size(), nullptr), size_(size), info_(info)
   {
-    fn::AttributesInfoBuilder builder;
-    for (const CustomDataLayer &layer : Span(custom_data.layers, custom_data.totlayer)) {
-      buffers_.append(layer.data);
-      switch (layer.type) {
-        case CD_PROP_INT32: {
-          builder.add<int32_t>(layer.name, 0);
-          break;
-        }
-        case CD_PROP_FLOAT3: {
-          builder.add<float3>(layer.name, {0, 0, 0});
-          break;
-        }
-      }
+    for (uint attribute_index : info.index_range()) {
+      StringRefNull name = info.name_of(attribute_index);
+      const fn::CPPType &cpp_type = info.type_of(attribute_index);
+      CustomDataType custom_type = cpp_to_custom_data_type(cpp_type);
+      void *data = CustomData_get_layer_named(&custom_data, custom_type, name.c_str());
+      buffers_[attribute_index] = data;
     }
-    info_ = std::make_unique<fn::AttributesInfo>(builder);
-    size_ = size;
   }
 
   operator fn::MutableAttributesRef()
   {
-    return fn::MutableAttributesRef(*info_, buffers_, size_);
+    return fn::MutableAttributesRef(info_, buffers_, size_);
   }
 
   operator fn::AttributesRef() const
   {
-    return fn::AttributesRef(*info_, buffers_, size_);
+    return fn::AttributesRef(info_, buffers_, size_);
   }
 };
 
-static void ensure_attributes_exist(ParticleSimulationState *state)
+static void ensure_attributes_exist(ParticleSimulationState *state, const fn::AttributesInfo &info)
 {
-  if (CustomData_get_layer_named(&state->attributes, CD_PROP_FLOAT3, "Position") == nullptr) {
-    CustomData_add_layer_named(
-        &state->attributes, CD_PROP_FLOAT3, CD_CALLOC, nullptr, state->tot_particles, "Position");
-  }
-  if (CustomData_get_layer_named(&state->attributes, CD_PROP_FLOAT3, "Velocity") == nullptr) {
-    CustomData_add_layer_named(
-        &state->attributes, CD_PROP_FLOAT3, CD_CALLOC, nullptr, state->tot_particles, "Velocity");
-  }
-  if (CustomData_get_layer_named(&state->attributes, CD_PROP_INT32, "ID") == nullptr) {
-    CustomData_add_layer_named(
-        &state->attributes, CD_PROP_INT32, CD_CALLOC, nullptr, state->tot_particles, "ID");
+  bool found_layer_to_remove;
+  do {
+    found_layer_to_remove = false;
+    for (int layer_index = 0; layer_index < state->attributes.totlayer; layer_index++) {
+      CustomDataLayer *layer = &state->attributes.layers[layer_index];
+      BLI_assert(layer->name != nullptr);
+      const fn::CPPType &cpp_type = custom_to_cpp_data_type((CustomDataType)layer->type);
+      StringRefNull name = layer->name;
+      if (!info.has_attribute(name, cpp_type)) {
+        found_layer_to_remove = true;
+        CustomData_free_layer(&state->attributes, layer->type, state->tot_particles, layer_index);
+        break;
+      }
+    }
+  } while (found_layer_to_remove);
+
+  for (uint attribute_index : info.index_range()) {
+    StringRefNull attribute_name = info.name_of(attribute_index);
+    const fn::CPPType &cpp_type = info.type_of(attribute_index);
+    CustomDataType custom_type = cpp_to_custom_data_type(cpp_type);
+    if (CustomData_get_layer_named(&state->attributes, custom_type, attribute_name.c_str()) ==
+        nullptr) {
+      void *data = CustomData_add_layer_named(&state->attributes,
+                                              custom_type,
+                                              CD_CALLOC,
+                                              nullptr,
+                                              state->tot_particles,
+                                              attribute_name.c_str());
+      cpp_type.fill_uninitialized(
+          info.default_of(attribute_index), data, (uint)state->tot_particles);
+    }
   }
 }
 
@@ -102,30 +144,21 @@ void solve_simulation_time_step(Simulation &simulation,
   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) {
-    ensure_attributes_exist(state);
-
-    fn::AttributesInfoBuilder builder;
-    CustomData &custom_data = state->attributes;
-    for (const CustomDataLayer &layer : Span(custom_data.layers, custom_data.totlayer)) {
-      switch (layer.type) {
-        case CD_PROP_INT32: {
-          builder.add<int32_t>(layer.name, 0);
-          break;
-        }
-        case CD_PROP_FLOAT3: {
-          builder.add<float3>(layer.name, {0, 0, 0});
-          break;
-        }
-      }
-    }
+    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(
         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) {
-    CustomDataAttributesRef custom_data_attributes{state->attributes, (uint)state->tot_particles};
+    const fn::AttributesInfo &attributes_info = *attribute_infos.lookup_as(state->head.name);
+    CustomDataAttributesRef custom_data_attributes{
+        state->attributes, (uint)state->tot_particles, attributes_info};
     fn::MutableAttributesRef attributes = custom_data_attributes;
 
     MutableSpan<float3> positions = attributes.get<float3>("Position");
@@ -159,6 +192,7 @@ void solve_simulation_time_step(Simulation &simulation,
   }
 
   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 uint emitted_particle_amount = allocator.total_allocated();
@@ -167,7 +201,8 @@ void solve_simulation_time_step(Simulation &simulation,
 
     CustomData_realloc(&state->attributes, new_particle_amount);
 
-    CustomDataAttributesRef custom_data_attributes{state->attributes, new_particle_amount};
+    CustomDataAttributesRef custom_data_attributes{
+        state->attributes, new_particle_amount, attributes_info};
     fn::MutableAttributesRef attributes = custom_data_attributes;
 
     uint offset = old_particle_amount;
diff --git a/source/blender/simulation/intern/simulation_solver.hh b/source/blender/simulation/intern/simulation_solver.hh
index 5a45a676ec7..7b37ed0583d 100644
--- a/source/blender/simulation/intern/simulation_solver.hh
+++ b/source/blender/simulation/intern/simulation_solver.hh
@@ -48,6 +48,7 @@ class ParticleForce {
 
 struct SimulationInfluences {
   Map<std::string, Vector<const ParticleForce *>> particle_forces;
+  Map<std::string, fn::AttributesInfoBuilder *> particle_attributes_builder;
   Vector<const ParticleEmitter *> particle_emitters;
 };