use gravity node in simulation

4 files changed, 57 insertions, 9 deletions

diff --git a/release/scripts/startup/nodes/socket_builder.py b/release/scripts/startup/nodes/socket_builder.py
index 4154605e138..2b7552c619d 100644
--- a/release/scripts/startup/nodes/socket_builder.py
+++ b/release/scripts/startup/nodes/socket_builder.py
@@ -9,6 +9,7 @@ from . declaration import (
     EmitterSocketDecl,
     EventSocketDecl,
     ControlFlowSocketDecl,
+    ParticleModifierSocketDecl,
 )
 
 class SocketBuilder:
@@ -198,11 +199,11 @@ class SocketBuilder:
         self._add_output(decl)
 
     def particle_modifier_input(self, identifier, name):
-        decl = ControlFlowSocketDecl(self.node, identifier, name)
+        decl = ParticleModifierSocketDecl(self.node, identifier, name)
         self._add_input(decl)
 
     def particle_modifier_output(self, identifier, name):
-        decl = ControlFlowSocketDecl(self.node, identifier, name)
+        decl = ParticleModifierSocketDecl(self.node, identifier, name)
         self._add_output(decl)
 
 
diff --git a/source/blender/simulations/bparticles/c_wrapper.cpp b/source/blender/simulations/bparticles/c_wrapper.cpp
index e366b4fb5ab..4996e426f35 100644
--- a/source/blender/simulations/bparticles/c_wrapper.cpp
+++ b/source/blender/simulations/bparticles/c_wrapper.cpp
@@ -248,6 +248,7 @@ using EmitterInserter = std::function<void(bNode *bnode,
                                            FN::DataFlowNodes::GeneratedGraph &data_graph,
                                            ModifierStepDescription &step_description)>;
 using EventInserter = EmitterInserter;
+using ModifierInserter = EmitterInserter;
 
 static bool is_particle_type_node(bNode *bnode)
 {
@@ -442,6 +443,30 @@ static void INSERT_EVENT_mesh_collision(bNode *event_node,
   }
 }
 
+static void INSERT_FORCE_gravity(bNode *force_node,
+                                 IndexedNodeTree &indexed_tree,
+                                 FN::DataFlowNodes::GeneratedGraph &data_graph,
+                                 ModifierStepDescription &step_description)
+{
+  BLI_assert(STREQ(force_node->idname, "bp_GravityForceNode"));
+  bSocketList node_inputs(force_node->inputs);
+  bSocketList node_outputs(force_node->outputs);
+
+  for (SocketWithNode linked : indexed_tree.linked(node_outputs.get(0))) {
+    if (!is_particle_type_node(linked.node)) {
+      continue;
+    }
+
+    SharedFunction fn = create_function(
+        indexed_tree, data_graph, {node_inputs.get(0)}, force_node->name);
+
+    Force *force = FORCE_gravity(fn);
+
+    bNode *type_node = linked.node;
+    step_description.m_types.lookup_ref(type_node->name)->m_integrator->m_forces.append(force);
+  }
+}
+
 static ModifierStepDescription *step_description_from_node_tree(bNodeTree *btree)
 {
   SCOPED_TIMER(__func__);
@@ -454,6 +479,9 @@ static ModifierStepDescription *step_description_from_node_tree(bNodeTree *btree
   event_inserters.add_new("bp_AgeReachedEventNode", INSERT_EVENT_age_reached);
   event_inserters.add_new("bp_MeshCollisionEventNode", INSERT_EVENT_mesh_collision);
 
+  SmallMap<std::string, ModifierInserter> modifier_inserters;
+  event_inserters.add_new("bp_GravityForceNode", INSERT_FORCE_gravity);
+
   ModifierStepDescription *step_description = new ModifierStepDescription();
 
   IndexedNodeTree indexed_tree(btree);
@@ -481,6 +509,12 @@ static ModifierStepDescription *step_description_from_node_tree(bNodeTree *btree
     }
   }
 
+  for (auto item : modifier_inserters.items()) {
+    for (bNode *modifier_node : indexed_tree.nodes_with_idname(item.key)) {
+      item.value(modifier_node, indexed_tree, generated_graph, *step_description);
+    }
+  }
+
   return step_description;
 }
 
diff --git a/source/blender/simulations/bparticles/forces.cpp b/source/blender/simulations/bparticles/forces.cpp
index ea5fc83a79b..4b3215720eb 100644
--- a/source/blender/simulations/bparticles/forces.cpp
+++ b/source/blender/simulations/bparticles/forces.cpp
@@ -8,19 +8,31 @@ Force::~Force()
 {
 }
 
-class DirectionalForce : public Force {
+class GravityForce : public Force {
  private:
-  float3 m_force;
+  SharedFunction m_compute_acceleration_fn;
+  TupleCallBody *m_compute_acceleration_body;
 
  public:
-  DirectionalForce(float3 force) : m_force(force)
+  GravityForce(SharedFunction &compute_acceleration_fn)
+      : m_compute_acceleration_fn(compute_acceleration_fn)
   {
+    m_compute_acceleration_body = m_compute_acceleration_fn->body<TupleCallBody>();
   }
 
   void add_force(ParticlesBlock &block, ArrayRef<float3> r_force) override
   {
+    FN_TUPLE_CALL_ALLOC_TUPLES(m_compute_acceleration_body, fn_in, fn_out);
+
+    FN::ExecutionStack stack;
+    FN::ExecutionContext execution_context(stack);
+
+    m_compute_acceleration_body->call(fn_in, fn_out, execution_context);
+
+    float3 acceleration = fn_out.get<float3>(0);
+
     for (uint i = 0; i < block.active_amount(); i++) {
-      r_force[i] += m_force;
+      r_force[i] += acceleration;
     }
   };
 };
@@ -46,9 +58,9 @@ class TurbulenceForce : public BParticles::Force {
   }
 };
 
-Force *FORCE_directional(float3 force)
+Force *FORCE_gravity(SharedFunction &compute_acceleration_fn)
 {
-  return new DirectionalForce(force);
+  return new GravityForce(compute_acceleration_fn);
 }
 
 Force *FORCE_turbulence(float strength)
diff --git a/source/blender/simulations/bparticles/forces.hpp b/source/blender/simulations/bparticles/forces.hpp
index c3045e09b18..08680b7da18 100644
--- a/source/blender/simulations/bparticles/forces.hpp
+++ b/source/blender/simulations/bparticles/forces.hpp
@@ -1,6 +1,7 @@
 #pragma once
 
 #include "core.hpp"
+#include "actions.hpp"
 
 namespace BParticles {
 
@@ -10,7 +11,7 @@ class Force {
   virtual void add_force(ParticlesBlock &block, ArrayRef<float3> r_force) = 0;
 };
 
-Force *FORCE_directional(float3 force);
+Force *FORCE_gravity(SharedFunction &compute_acceleration_fn);
 Force *FORCE_turbulence(float strength);
 
 }  // namespace BParticles