Functions: support wrapping lazy-function node execute function

This is a light weight solution to passing in some extra context into
a lazy-function that is invoked by the graph executor.
The new functionality is used by #112421.

source/blender/functions/FN_lazy_function_graph_executor.hh

@@ -56,10 +56,29 @@ class GraphExecutorSideEffectProvider {
   virtual Vector<const FunctionNode *> get_nodes_with_side_effects(const Context &context) const;
 };
 
+/**
+ * Can be used to pass extra context into the execution of a function. The main alternative to this
+ * is to create a wrapper `LazyFunction` for the `FunctionNode`s. Using this light weight wrapper
+ * is preferable if possible.
+ */
+class GraphExecutorNodeExecuteWrapper {
+ public:
+  virtual ~GraphExecutorNodeExecuteWrapper() = default;
+
+  /**
+   * Is expected to run `node.function().execute(params, context)` but might do some extra work,
+   * like adjusting the context.
+   */
+  virtual void execute_node(const FunctionNode &node,
+                            Params &params,
+                            const Context &context) const = 0;
+};
+
 class GraphExecutor : public LazyFunction {
  public:
   using Logger = GraphExecutorLogger;
   using SideEffectProvider = GraphExecutorSideEffectProvider;
+  using NodeExecuteWrapper = GraphExecutorNodeExecuteWrapper;
 
  private:
   /**
@@ -80,6 +99,10 @@ class GraphExecutor : public LazyFunction {
    * during evaluation.
    */
   const SideEffectProvider *side_effect_provider_;
+  /**
+   * Optional wrapper for node execution functions.
+   */
+  const NodeExecuteWrapper *node_execute_wrapper_;
 
   /**
    * When a graph is executed, various things have to be allocated (e.g. the state of all nodes).
@@ -100,7 +123,8 @@ class GraphExecutor : public LazyFunction {
                 Span<const OutputSocket *> graph_inputs,
                 Span<const InputSocket *> graph_outputs,
                 const Logger *logger,
-                const SideEffectProvider *side_effect_provider);
+                const SideEffectProvider *side_effect_provider,
+                const NodeExecuteWrapper *node_execute_wrapper);
 
   void *init_storage(LinearAllocator<> &allocator) const override;
   void destruct_storage(void *storage) const override;

source/blender/functions/intern/lazy_function_graph_executor.cc

@@ -1412,7 +1412,12 @@ inline void Executor::execute_node(const FunctionNode &node,
   };
 
   lazy_threading::HintReceiver blocking_hint_receiver{blocking_hint_fn};
-  fn.execute(node_params, fn_context);
+  if (self_.node_execute_wrapper_) {
+    self_.node_execute_wrapper_->execute_node(node, node_params, fn_context);
+  }
+  else {
+    fn.execute(node_params, fn_context);
+  }
 
   if (self_.logger_ != nullptr) {
     self_.logger_->log_after_node_execute(node, node_params, fn_context);
@@ -1423,12 +1428,14 @@ GraphExecutor::GraphExecutor(const Graph &graph,
                              const Span<const OutputSocket *> graph_inputs,
                              const Span<const InputSocket *> graph_outputs,
                              const Logger *logger,
-                             const SideEffectProvider *side_effect_provider)
+                             const SideEffectProvider *side_effect_provider,
+                             const NodeExecuteWrapper *node_execute_wrapper)
     : graph_(graph),
       graph_inputs_(graph_inputs),
       graph_outputs_(graph_outputs),
       logger_(logger),
-      side_effect_provider_(side_effect_provider)
+      side_effect_provider_(side_effect_provider),
+      node_execute_wrapper_(node_execute_wrapper)
 {
   /* The graph executor can handle partial execution when there are still missing inputs. */
   allow_missing_requested_inputs_ = true;

source/blender/functions/tests/FN_lazy_function_test.cc

@@ -108,7 +108,8 @@ TEST(lazy_function, SideEffects)
 
   SimpleSideEffectProvider side_effect_provider{{&store_node}};
 
-  GraphExecutor executor_fn{graph, {&input_node.output(0)}, {}, nullptr, &side_effect_provider};
+  GraphExecutor executor_fn{
+      graph, {&input_node.output(0)}, {}, nullptr, &side_effect_provider, nullptr};
   execute_lazy_function_eagerly(
       executor_fn, nullptr, nullptr, std::make_tuple(5), std::make_tuple());
 
@@ -172,7 +173,7 @@ TEST(lazy_function, GraphWithCycle)
   graph.update_node_indices();
 
   GraphExecutor executor_fn{
-      graph, {&input_node.output(0)}, {&output_node.input(0)}, nullptr, nullptr};
+      graph, {&input_node.output(0)}, {&output_node.input(0)}, nullptr, nullptr, nullptr};
   int result = 0;
   execute_lazy_function_eagerly(
       executor_fn, nullptr, nullptr, std::make_tuple(10), std::make_tuple(&result));

source/blender/nodes/intern/geometry_nodes_execute.cc

@@ -586,8 +586,12 @@ bke::GeometrySet execute_geometry_nodes_on_geometry(
   nodes::GeometryNodesLazyFunctionLogger lf_logger(lf_graph_info);
   nodes::GeometryNodesLazyFunctionSideEffectProvider lf_side_effect_provider;
 
-  lf::GraphExecutor graph_executor{
-      lf_graph_info.graph, graph_inputs, graph_outputs, &lf_logger, &lf_side_effect_provider};
+  lf::GraphExecutor graph_executor{lf_graph_info.graph,
+                                   graph_inputs,
+                                   graph_outputs,
+                                   &lf_logger,
+                                   &lf_side_effect_provider,
+                                   nullptr};
 
   nodes::GeoNodesLFUserData user_data;
   fill_user_data(user_data);

source/blender/nodes/intern/geometry_nodes_lazy_function.cc

@@ -1008,7 +1008,8 @@ class LazyFunctionForGroupNode : public LazyFunction {
                             std::move(graph_inputs),
                             std::move(graph_outputs),
                             &*lf_logger_,
-                            &*lf_side_effect_provider_);
+                            &*lf_side_effect_provider_,
+                            nullptr);
   }
 
   void execute_impl(lf::Params &params, const lf::Context &context) const override
@@ -2020,7 +2021,7 @@ struct GeometryNodesLazyFunctionGraphBuilder {
     auto &side_effect_provider = scope_.construct<GeometryNodesLazyFunctionSideEffectProvider>();
 
     const auto &lf_graph_fn = scope_.construct<lf::GraphExecutor>(
-        lf_graph, lf_zone_inputs, lf_zone_outputs, &logger, &side_effect_provider);
+        lf_graph, lf_zone_inputs, lf_zone_outputs, &logger, &side_effect_provider, nullptr);
     const auto &zone_function = scope_.construct<LazyFunctionForSimulationZone>(*zone.output_node,
                                                                                 lf_graph_fn);
     zone_info.lazy_function = &zone_function;
@@ -2158,7 +2159,7 @@ struct GeometryNodesLazyFunctionGraphBuilder {
     auto &logger = scope_.construct<GeometryNodesLazyFunctionLogger>(*lf_graph_info_);
     auto &side_effect_provider = scope_.construct<GeometryNodesLazyFunctionSideEffectProvider>();
     body_fn.function = &scope_.construct<lf::GraphExecutor>(
-        lf_body_graph, lf_body_inputs, lf_body_outputs, &logger, &side_effect_provider);
+        lf_body_graph, lf_body_inputs, lf_body_outputs, &logger, &side_effect_provider, nullptr);
 
     // std::cout << "\n\n" << lf_body_graph.to_dot() << "\n\n";