Geometry Nodes: decouple multi-function lifetimes from modifier

Previously, some multi-functions were allocated in a resource scope.
This was fine as long as the multi-functions were only needed during
the current evaluation of the node tree. However, now cases arise
that require the multi-functions to be alive after the modifier is finished.
For example, we want to evaluate fields created with geometry nodes
outside of geometry nodes.

To make this work, `std::shared_ptr` has to be used in a few more places.
Realistically, this shouldn't have a noticable impact on performance.
If this does become a bottleneck in the future, we can think about ways
to make this work without using `shared_ptr` for multi-functions that
are only used once.

6 files changed, 42 insertions, 35 deletions

diff --git a/source/blender/functions/FN_field.hh b/source/blender/functions/FN_field.hh
index f65c4e443f2..2fca78fa6e7 100644
--- a/source/blender/functions/FN_field.hh
+++ b/source/blender/functions/FN_field.hh
@@ -204,14 +204,14 @@ class FieldOperation : public FieldNode {
    * The multi-function used by this node. It is optionally owned.
    * Multi-functions with mutable or vector parameters are not supported currently.
    */
-  std::unique_ptr<const MultiFunction> owned_function_;
+  std::shared_ptr<const MultiFunction> owned_function_;
   const MultiFunction *function_;
 
   /** Inputs to the operation. */
   blender::Vector<GField> inputs_;
 
  public:
-  FieldOperation(std::unique_ptr<const MultiFunction> function, Vector<GField> inputs = {});
+  FieldOperation(std::shared_ptr<const MultiFunction> function, Vector<GField> inputs = {});
   FieldOperation(const MultiFunction &function, Vector<GField> inputs = {});
 
   Span<GField> inputs() const;
diff --git a/source/blender/functions/intern/field.cc b/source/blender/functions/intern/field.cc
index 39688ef3daf..03af3f53065 100644
--- a/source/blender/functions/intern/field.cc
+++ b/source/blender/functions/intern/field.cc
@@ -534,7 +534,7 @@ const GVArray *IndexFieldInput::get_varray_for_context(const fn::FieldContext &U
  * FieldOperation.
  */
 
-FieldOperation::FieldOperation(std::unique_ptr<const MultiFunction> function,
+FieldOperation::FieldOperation(std::shared_ptr<const MultiFunction> function,
                                Vector<GField> inputs)
     : FieldOperation(*function, std::move(inputs))
 {
diff --git a/source/blender/modifiers/intern/MOD_nodes.cc b/source/blender/modifiers/intern/MOD_nodes.cc
index ebccbb3d2cd..14b40582516 100644
--- a/source/blender/modifiers/intern/MOD_nodes.cc
+++ b/source/blender/modifiers/intern/MOD_nodes.cc
@@ -924,7 +924,7 @@ static GeometrySet compute_geometry(const DerivedNodeTree &tree,
 {
   blender::ResourceScope scope;
   blender::LinearAllocator<> &allocator = scope.linear_allocator();
-  blender::nodes::NodeMultiFunctions mf_by_node{tree, scope};
+  blender::nodes::NodeMultiFunctions mf_by_node{tree};
 
   Map<DOutputSocket, GMutablePointer> group_inputs;
 
diff --git a/source/blender/modifiers/intern/MOD_nodes_evaluator.cc b/source/blender/modifiers/intern/MOD_nodes_evaluator.cc
index 20ee6127504..0e7f5fe155b 100644
--- a/source/blender/modifiers/intern/MOD_nodes_evaluator.cc
+++ b/source/blender/modifiers/intern/MOD_nodes_evaluator.cc
@@ -882,9 +882,9 @@ class GeometryNodesEvaluator {
     }
 
     /* Use the multi-function implementation if it exists. */
-    const MultiFunction *multi_function = params_.mf_by_node->try_get(node);
-    if (multi_function != nullptr) {
-      this->execute_multi_function_node(node, *multi_function, node_state);
+    const nodes::NodeMultiFunctions::Item &fn_item = params_.mf_by_node->try_get(node);
+    if (fn_item.fn != nullptr) {
+      this->execute_multi_function_node(node, fn_item, node_state);
       return;
     }
 
@@ -905,7 +905,7 @@ class GeometryNodesEvaluator {
   }
 
   void execute_multi_function_node(const DNode node,
-                                   const MultiFunction &fn,
+                                   const nodes::NodeMultiFunctions::Item &fn_item,
                                    NodeState &node_state)
   {
     if (node->idname().find("Legacy") != StringRef::not_found) {
@@ -933,7 +933,13 @@ class GeometryNodesEvaluator {
       input_fields.append(std::move(*(GField *)single_value.value));
     }
 
-    auto operation = std::make_shared<fn::FieldOperation>(fn, std::move(input_fields));
+    std::shared_ptr<fn::FieldOperation> operation;
+    if (fn_item.owned_fn) {
+      operation = std::make_shared<fn::FieldOperation>(fn_item.owned_fn, std::move(input_fields));
+    }
+    else {
+      operation = std::make_shared<fn::FieldOperation>(*fn_item.fn, std::move(input_fields));
+    }
 
     /* Forward outputs. */
     int output_index = 0;
diff --git a/source/blender/nodes/NOD_multi_function.hh b/source/blender/nodes/NOD_multi_function.hh
index c1952cf8014..367ceaab9f7 100644
--- a/source/blender/nodes/NOD_multi_function.hh
+++ b/source/blender/nodes/NOD_multi_function.hh
@@ -33,15 +33,15 @@ class NodeMultiFunctions;
  */
 class NodeMultiFunctionBuilder : NonCopyable, NonMovable {
  private:
-  ResourceScope &resource_scope_;
   bNode &node_;
   bNodeTree &tree_;
+  std::shared_ptr<MultiFunction> owned_built_fn_;
   const MultiFunction *built_fn_ = nullptr;
 
   friend NodeMultiFunctions;
 
  public:
-  NodeMultiFunctionBuilder(ResourceScope &resource_scope, bNode &node, bNodeTree &tree);
+  NodeMultiFunctionBuilder(bNode &node, bNodeTree &tree);
 
   /**
    * Assign a multi-function for the current node. The input and output parameters of the function
@@ -58,31 +58,33 @@ class NodeMultiFunctionBuilder : NonCopyable, NonMovable {
 
   bNode &node();
   bNodeTree &tree();
-
-  ResourceScope &resource_scope();
 };
 
 /**
  * Gives access to multi-functions for all nodes in a node tree that support them.
  */
 class NodeMultiFunctions {
+ public:
+  struct Item {
+    const MultiFunction *fn = nullptr;
+    std::shared_ptr<MultiFunction> owned_fn;
+  };
+
  private:
-  Map<const bNode *, const MultiFunction *> map_;
+  Map<const bNode *, Item> map_;
 
  public:
-  NodeMultiFunctions(const DerivedNodeTree &tree, ResourceScope &resource_scope);
+  NodeMultiFunctions(const DerivedNodeTree &tree);
 
-  const MultiFunction *try_get(const DNode &node) const;
+  const Item &try_get(const DNode &node) const;
 };
 
 /* -------------------------------------------------------------------- */
 /** \name #NodeMultiFunctionBuilder Inline Methods
  * \{ */
 
-inline NodeMultiFunctionBuilder::NodeMultiFunctionBuilder(ResourceScope &resource_scope,
-                                                          bNode &node,
-                                                          bNodeTree &tree)
-    : resource_scope_(resource_scope), node_(node), tree_(tree)
+inline NodeMultiFunctionBuilder::NodeMultiFunctionBuilder(bNode &node, bNodeTree &tree)
+    : node_(node), tree_(tree)
 {
 }
 
@@ -96,11 +98,6 @@ inline bNodeTree &NodeMultiFunctionBuilder::tree()
   return tree_;
 }
 
-inline ResourceScope &NodeMultiFunctionBuilder::resource_scope()
-{
-  return resource_scope_;
-}
-
 inline void NodeMultiFunctionBuilder::set_matching_fn(const MultiFunction *fn)
 {
   built_fn_ = fn;
@@ -108,14 +105,14 @@ inline void NodeMultiFunctionBuilder::set_matching_fn(const MultiFunction *fn)
 
 inline void NodeMultiFunctionBuilder::set_matching_fn(const MultiFunction &fn)
 {
-  this->set_matching_fn(&fn);
+  built_fn_ = &fn;
 }
 
 template<typename T, typename... Args>
 inline void NodeMultiFunctionBuilder::construct_and_set_matching_fn(Args &&...args)
 {
-  const T &fn = resource_scope_.construct<T>(std::forward<Args>(args)...);
-  this->set_matching_fn(&fn);
+  owned_built_fn_ = std::make_shared<T>(std::forward<Args>(args)...);
+  built_fn_ = &*owned_built_fn_;
 }
 
 /** \} */
@@ -124,9 +121,14 @@ inline void NodeMultiFunctionBuilder::construct_and_set_matching_fn(Args &&...ar
 /** \name #NodeMultiFunctions Inline Methods
  * \{ */
 
-inline const MultiFunction *NodeMultiFunctions::try_get(const DNode &node) const
+inline const NodeMultiFunctions::Item &NodeMultiFunctions::try_get(const DNode &node) const
 {
-  return map_.lookup_default(node->bnode(), nullptr);
+  static Item empty_item;
+  const Item *item = map_.lookup_ptr(node->bnode());
+  if (item == nullptr) {
+    return empty_item;
+  }
+  return *item;
 }
 
 /** \} */
diff --git a/source/blender/nodes/intern/node_multi_function.cc b/source/blender/nodes/intern/node_multi_function.cc
index c91899ed8c2..6d79ed839b2 100644
--- a/source/blender/nodes/intern/node_multi_function.cc
+++ b/source/blender/nodes/intern/node_multi_function.cc
@@ -18,7 +18,7 @@
 
 namespace blender::nodes {
 
-NodeMultiFunctions::NodeMultiFunctions(const DerivedNodeTree &tree, ResourceScope &resource_scope)
+NodeMultiFunctions::NodeMultiFunctions(const DerivedNodeTree &tree)
 {
   for (const NodeTreeRef *tree_ref : tree.used_node_tree_refs()) {
     bNodeTree *btree = tree_ref->btree();
@@ -27,11 +27,10 @@ NodeMultiFunctions::NodeMultiFunctions(const DerivedNodeTree &tree, ResourceScop
       if (bnode->typeinfo->build_multi_function == nullptr) {
         continue;
       }
-      NodeMultiFunctionBuilder builder{resource_scope, *bnode, *btree};
+      NodeMultiFunctionBuilder builder{*bnode, *btree};
       bnode->typeinfo->build_multi_function(builder);
-      const MultiFunction *fn = builder.built_fn_;
-      if (fn != nullptr) {
-        map_.add_new(bnode, fn);
+      if (builder.built_fn_ != nullptr) {
+        map_.add_new(bnode, {builder.built_fn_, std::move(builder.owned_built_fn_)});
       }
     }
   }