Functions: remove multi-function network

The multi-function network system was able to compose multiple
multi-functions into a new one and to evaluate that efficiently.
This functionality was heavily used by the particle nodes prototype
a year ago. However, since then we only used multi-functions
without the need to compose them in geometry nodes.

The upcoming "fields" in geometry nodes will need a way to
compose multi-functions again. Unfortunately, the code removed
in this commit was not ideal for this different kind of function
composition. I've been working on an alternative that will be added
separately when it becomes needed.

I've had to update all the function nodes, because their interface
depended on the multi-function network data structure a bit.
The actual multi-function implementations are still the same though.

27 files changed, 306 insertions, 1016 deletions

diff --git a/source/blender/nodes/CMakeLists.txt b/source/blender/nodes/CMakeLists.txt
index 46fb9f54bfe..8680fcee49a 100644
--- a/source/blender/nodes/CMakeLists.txt
+++ b/source/blender/nodes/CMakeLists.txt
@@ -345,11 +345,10 @@ set(SRC
   intern/node_common.c
   intern/node_exec.cc
   intern/node_geometry_exec.cc
+  intern/node_multi_function.cc
   intern/node_socket.cc
-  intern/node_tree_multi_function.cc
   intern/node_tree_ref.cc
   intern/node_util.c
-  intern/type_callbacks.cc
   intern/type_conversions.cc
 
   composite/node_composite_util.h
@@ -366,13 +365,12 @@ set(SRC
   NOD_geometry_exec.hh
   NOD_geometry_nodes_eval_log.hh
   NOD_math_functions.hh
-  NOD_node_tree_multi_function.hh
+  NOD_multi_function.hh
   NOD_node_tree_ref.hh
   NOD_shader.h
   NOD_socket.h
   NOD_static_types.h
   NOD_texture.h
-  NOD_type_callbacks.hh
   NOD_type_conversions.hh
   intern/node_common.h
   intern/node_exec.h
diff --git a/source/blender/nodes/NOD_multi_function.hh b/source/blender/nodes/NOD_multi_function.hh
new file mode 100644
index 00000000000..2f4b104fb4c
--- /dev/null
+++ b/source/blender/nodes/NOD_multi_function.hh
@@ -0,0 +1,130 @@
+/*
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ */
+
+#pragma once
+
+#include "FN_multi_function.hh"
+
+#include "DNA_node_types.h"
+
+#include "NOD_derived_node_tree.hh"
+
+namespace blender::nodes {
+
+using namespace fn::multi_function_types;
+
+class NodeMultiFunctions;
+
+/**
+ * Utility class to help nodes build a multi-function for themselves.
+ */
+class NodeMultiFunctionBuilder : NonCopyable, NonMovable {
+ private:
+  ResourceScope &resource_scope_;
+  bNode &node_;
+  bNodeTree &tree_;
+  const MultiFunction *built_fn_ = nullptr;
+
+  friend NodeMultiFunctions;
+
+ public:
+  NodeMultiFunctionBuilder(ResourceScope &resource_scope, bNode &node, bNodeTree &tree);
+
+  /**
+   * Assign a multi-function for the current node. The input and output parameters of the function
+   * have to match the available sockets in the node.
+   */
+  void set_matching_fn(const MultiFunction *fn);
+  void set_matching_fn(const MultiFunction &fn);
+
+  /**
+   * Utility method for creating and assigning a multi-function when it can't have a static
+   * lifetime.
+   */
+  template<typename T, typename... Args> void construct_and_set_matching_fn(Args &&...args);
+
+  bNode &node();
+  bNodeTree &tree();
+
+  ResourceScope &resource_scope();
+};
+
+/**
+ * Gives access to multi-functions for all nodes in a node tree that support them.
+ */
+class NodeMultiFunctions {
+ private:
+  Map<const bNode *, const MultiFunction *> map_;
+
+ public:
+  NodeMultiFunctions(const DerivedNodeTree &tree, ResourceScope &resource_scope);
+
+  const MultiFunction *try_get(const DNode &node) const;
+};
+
+/* --------------------------------------------------------------------
+ * NodeMultiFunctionBuilder inline methods.
+ */
+
+inline NodeMultiFunctionBuilder::NodeMultiFunctionBuilder(ResourceScope &resource_scope,
+                                                          bNode &node,
+                                                          bNodeTree &tree)
+    : resource_scope_(resource_scope), node_(node), tree_(tree)
+{
+}
+
+inline bNode &NodeMultiFunctionBuilder::node()
+{
+  return node_;
+}
+
+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;
+}
+
+inline void NodeMultiFunctionBuilder::set_matching_fn(const MultiFunction &fn)
+{
+  this->set_matching_fn(&fn);
+}
+
+template<typename T, typename... Args>
+inline void NodeMultiFunctionBuilder::construct_and_set_matching_fn(Args &&...args)
+{
+  const T &fn = resource_scope_.construct<T>(__func__, std::forward<Args>(args)...);
+  this->set_matching_fn(&fn);
+}
+
+/* --------------------------------------------------------------------
+ * NodeMultiFunctions inline methods.
+ */
+
+inline const MultiFunction *NodeMultiFunctions::try_get(const DNode &node) const
+{
+  return map_.lookup_default(node->bnode(), nullptr);
+}
+
+}  // namespace blender::nodes
diff --git a/source/blender/nodes/NOD_node_tree_multi_function.hh b/source/blender/nodes/NOD_node_tree_multi_function.hh
deleted file mode 100644
index 7eeeaef0b98..00000000000
--- a/source/blender/nodes/NOD_node_tree_multi_function.hh
+++ /dev/null
@@ -1,390 +0,0 @@
-/*
- * This program is free software; you can redistribute it and/or
- * modify it under the terms of the GNU General Public License
- * as published by the Free Software Foundation; either version 2
- * of the License, or (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program; if not, write to the Free Software Foundation,
- * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
- */
-
-#pragma once
-
-/** \file
- * \ingroup nodes
- *
- * This file allows you to generate a multi-function network from a user-generated node tree.
- */
-
-#include "FN_multi_function_builder.hh"
-#include "FN_multi_function_network.hh"
-
-#include "NOD_derived_node_tree.hh"
-#include "NOD_type_callbacks.hh"
-
-#include "BLI_multi_value_map.hh"
-#include "BLI_resource_scope.hh"
-
-namespace blender::nodes {
-
-/**
- * A MFNetworkTreeMap maps various components of a node tree to components of a fn::MFNetwork. This
- * is necessary for further processing of a multi-function network that has been generated from a
- * node tree.
- */
-class MFNetworkTreeMap {
- private:
-  /**
-   * Store by id instead of using a hash table to avoid unnecessary hash table lookups.
-   *
-   * Input sockets in a node tree can have multiple corresponding sockets in the generated
-   * MFNetwork. This is because nodes are allowed to expand into multiple multi-function nodes.
-   */
-  const DerivedNodeTree &tree_;
-  fn::MFNetwork &network_;
-  MultiValueMap<DSocket, fn::MFSocket *> sockets_by_dsocket_;
-
- public:
-  MFNetworkTreeMap(const DerivedNodeTree &tree, fn::MFNetwork &network)
-      : tree_(tree), network_(network)
-  {
-  }
-
-  const DerivedNodeTree &tree() const
-  {
-    return tree_;
-  }
-
-  const fn::MFNetwork &network() const
-  {
-    return network_;
-  }
-
-  fn::MFNetwork &network()
-  {
-    return network_;
-  }
-
-  void add(const DSocket &dsocket, fn::MFSocket &socket)
-  {
-    BLI_assert(dsocket->is_input() == socket.is_input());
-    BLI_assert(dsocket->is_input() || sockets_by_dsocket_.lookup(dsocket).is_empty());
-    sockets_by_dsocket_.add(dsocket, &socket);
-  }
-
-  void add(const DInputSocket &dsocket, fn::MFInputSocket &socket)
-  {
-    sockets_by_dsocket_.add(dsocket, &socket);
-  }
-
-  void add(const DOutputSocket &dsocket, fn::MFOutputSocket &socket)
-  {
-    /* There can be at most one matching output socket. */
-    BLI_assert(sockets_by_dsocket_.lookup(dsocket).is_empty());
-    sockets_by_dsocket_.add(dsocket, &socket);
-  }
-
-  void add(const DTreeContext &context,
-           Span<const InputSocketRef *> dsockets,
-           Span<fn::MFInputSocket *> sockets)
-  {
-    assert_same_size(dsockets, sockets);
-    for (int i : dsockets.index_range()) {
-      this->add(DInputSocket(&context, dsockets[i]), *sockets[i]);
-    }
-  }
-
-  void add(const DTreeContext &context,
-           Span<const OutputSocketRef *> dsockets,
-           Span<fn::MFOutputSocket *> sockets)
-  {
-    assert_same_size(dsockets, sockets);
-    for (int i : dsockets.index_range()) {
-      this->add(DOutputSocket(&context, dsockets[i]), *sockets[i]);
-    }
-  }
-
-  void add_try_match(const DNode &dnode, fn::MFNode &node)
-  {
-    this->add_try_match(*dnode.context(),
-                        dnode->inputs().cast<const SocketRef *>(),
-                        node.inputs().cast<fn::MFSocket *>());
-    this->add_try_match(*dnode.context(),
-                        dnode->outputs().cast<const SocketRef *>(),
-                        node.outputs().cast<fn::MFSocket *>());
-  }
-
-  void add_try_match(const DTreeContext &context,
-                     Span<const InputSocketRef *> dsockets,
-                     Span<fn::MFInputSocket *> sockets)
-  {
-    this->add_try_match(
-        context, dsockets.cast<const SocketRef *>(), sockets.cast<fn::MFSocket *>());
-  }
-
-  void add_try_match(const DTreeContext &context,
-                     Span<const OutputSocketRef *> dsockets,
-                     Span<fn::MFOutputSocket *> sockets)
-  {
-    this->add_try_match(
-        context, dsockets.cast<const SocketRef *>(), sockets.cast<fn::MFSocket *>());
-  }
-
-  void add_try_match(const DTreeContext &context,
-                     Span<const SocketRef *> dsockets,
-                     Span<fn::MFSocket *> sockets)
-  {
-    int used_sockets = 0;
-    for (const SocketRef *dsocket : dsockets) {
-      if (!dsocket->is_available()) {
-        continue;
-      }
-      if (!socket_is_mf_data_socket(*dsocket->typeinfo())) {
-        continue;
-      }
-      fn::MFSocket *socket = sockets[used_sockets];
-      this->add(DSocket(&context, dsocket), *socket);
-      used_sockets++;
-    }
-  }
-
-  fn::MFOutputSocket &lookup(const DOutputSocket &dsocket)
-  {
-    return sockets_by_dsocket_.lookup(dsocket)[0]->as_output();
-  }
-
-  Span<fn::MFInputSocket *> lookup(const DInputSocket &dsocket)
-  {
-    return sockets_by_dsocket_.lookup(dsocket).cast<fn::MFInputSocket *>();
-  }
-
-  fn::MFInputSocket &lookup_dummy(const DInputSocket &dsocket)
-  {
-    Span<fn::MFInputSocket *> sockets = this->lookup(dsocket);
-    BLI_assert(sockets.size() == 1);
-    fn::MFInputSocket &socket = *sockets[0];
-    BLI_assert(socket.node().is_dummy());
-    return socket;
-  }
-
-  fn::MFOutputSocket &lookup_dummy(const DOutputSocket &dsocket)
-  {
-    fn::MFOutputSocket &socket = this->lookup(dsocket);
-    BLI_assert(socket.node().is_dummy());
-    return socket;
-  }
-
-  bool is_mapped(const DSocket &dsocket) const
-  {
-    return !sockets_by_dsocket_.lookup(dsocket).is_empty();
-  }
-};
-
-/**
- * This data is necessary throughout the generation of a MFNetwork from a node tree.
- */
-struct CommonMFNetworkBuilderData {
-  ResourceScope &scope;
-  fn::MFNetwork &network;
-  MFNetworkTreeMap &network_map;
-  const DerivedNodeTree &tree;
-};
-
-class MFNetworkBuilderBase {
- protected:
-  CommonMFNetworkBuilderData &common_;
-
- public:
-  MFNetworkBuilderBase(CommonMFNetworkBuilderData &common) : common_(common)
-  {
-  }
-
-  /**
-   * Returns the network that is currently being built.
-   */
-  fn::MFNetwork &network()
-  {
-    return common_.network;
-  }
-
-  /**
-   * Returns the map between the node tree and the multi-function network that is being built.
-   */
-  MFNetworkTreeMap &network_map()
-  {
-    return common_.network_map;
-  }
-
-  /**
-   * Returns a resource collector that will only be destructed after the multi-function network is
-   * destructed.
-   */
-  ResourceScope &resource_scope()
-  {
-    return common_.scope;
-  }
-
-  /**
-   * Constructs a new function that will live at least as long as the MFNetwork.
-   */
-  template<typename T, typename... Args> T &construct_fn(Args &&...args)
-  {
-    BLI_STATIC_ASSERT((std::is_base_of_v<fn::MultiFunction, T>), "");
-    void *buffer = common_.scope.linear_allocator().allocate(sizeof(T), alignof(T));
-    T *fn = new (buffer) T(std::forward<Args>(args)...);
-    common_.scope.add(destruct_ptr<T>(fn), fn->name().c_str());
-    return *fn;
-  }
-};
-
-/**
- * This class is used by socket implementations to define how an unlinked input socket is handled
- * in a multi-function network.
- */
-class SocketMFNetworkBuilder : public MFNetworkBuilderBase {
- private:
-  bNodeSocket *bsocket_;
-  fn::MFOutputSocket *built_socket_ = nullptr;
-
- public:
-  SocketMFNetworkBuilder(CommonMFNetworkBuilderData &common, const DSocket &dsocket)
-      : MFNetworkBuilderBase(common), bsocket_(dsocket->bsocket())
-  {
-  }
-
-  /**
-   * Returns the socket that is currently being built.
-   */
-  bNodeSocket &bsocket()
-  {
-    return *bsocket_;
-  }
-
-  /**
-   * Utility method that returns bsocket->default_value for the current socket.
-   */
-  template<typename T> T *socket_default_value()
-  {
-    return static_cast<T *>(bsocket_->default_value);
-  }
-
-  /**
-   * Builds a function node for that socket that outputs the given constant value.
-   */
-  template<typename T> void set_constant_value(T value)
-  {
-    this->construct_generator_fn<fn::CustomMF_Constant<T>>(std::move(value));
-  }
-  void set_constant_value(const CPPType &type, const void *value)
-  {
-    /* The value has live as long as the generated mf network. */
-    this->construct_generator_fn<fn::CustomMF_GenericConstant>(type, value);
-  }
-
-  template<typename T, typename... Args> void construct_generator_fn(Args &&...args)
-  {
-    const fn::MultiFunction &fn = this->construct_fn<T>(std::forward<Args>(args)...);
-    this->set_generator_fn(fn);
-  }
-
-  /**
-   * Uses the first output of the given multi-function as value of the socket.
-   */
-  void set_generator_fn(const fn::MultiFunction &fn)
-  {
-    fn::MFFunctionNode &node = common_.network.add_function(fn);
-    this->set_socket(node.output(0));
-  }
-
-  /**
-   * Define a multi-function socket that outputs the value of the bsocket.
-   */
-  void set_socket(fn::MFOutputSocket &socket)
-  {
-    built_socket_ = &socket;
-  }
-
-  fn::MFOutputSocket *built_socket()
-  {
-    return built_socket_;
-  }
-};
-
-/**
- * This class is used by node implementations to define how a user-level node expands into
- * multi-function nodes internally.
- */
-class NodeMFNetworkBuilder : public MFNetworkBuilderBase {
- private:
-  DNode dnode_;
-
- public:
-  NodeMFNetworkBuilder(CommonMFNetworkBuilderData &common, DNode dnode)
-      : MFNetworkBuilderBase(common), dnode_(dnode)
-  {
-  }
-
-  /**
-   * Tells the builder to build a function that corresponds to the node that is being built. It
-   * will try to match up sockets.
-   */
-  template<typename T, typename... Args> T &construct_and_set_matching_fn(Args &&...args)
-  {
-    T &function = this->construct_fn<T>(std::forward<Args>(args)...);
-    this->set_matching_fn(function);
-    return function;
-  }
-
-  const fn::MultiFunction &get_not_implemented_fn()
-  {
-    return this->get_default_fn("Not Implemented (" + dnode_->name() + ")");
-  }
-
-  const fn::MultiFunction &get_default_fn(StringRef name);
-
-  const void set_not_implemented()
-  {
-    this->set_matching_fn(this->get_not_implemented_fn());
-  }
-
-  /**
-   * Tells the builder that the given function corresponds to the node that is being built. It will
-   * try to match up sockets. For that it skips unavailable and non-data sockets.
-   */
-  void set_matching_fn(const fn::MultiFunction &function)
-  {
-    fn::MFFunctionNode &node = common_.network.add_function(function);
-    common_.network_map.add_try_match(dnode_, node);
-  }
-
-  /**
-   * Returns the node that is currently being built.
-   */
-  bNode &bnode()
-  {
-    return *dnode_->bnode();
-  }
-
-  /**
-   * Returns the node that is currently being built.
-   */
-  const DNode &dnode() const
-  {
-    return dnode_;
-  }
-};
-
-MFNetworkTreeMap insert_node_tree_into_mf_network(fn::MFNetwork &network,
-                                                  const DerivedNodeTree &tree,
-                                                  ResourceScope &scope);
-
-using MultiFunctionByNode = Map<DNode, const fn::MultiFunction *>;
-MultiFunctionByNode get_multi_function_per_node(const DerivedNodeTree &tree, ResourceScope &scope);
-
-}  // namespace blender::nodes
diff --git a/source/blender/nodes/function/node_function_util.hh b/source/blender/nodes/function/node_function_util.hh
index 9fbd6712827..96a8f29c3e9 100644
--- a/source/blender/nodes/function/node_function_util.hh
+++ b/source/blender/nodes/function/node_function_util.hh
@@ -30,7 +30,7 @@
 #include "BLT_translation.h"
 
 #include "NOD_function.h"
-#include "NOD_node_tree_multi_function.hh"
+#include "NOD_multi_function.hh"
 
 #include "node_util.h"
 
diff --git a/source/blender/nodes/function/nodes/node_fn_boolean_math.cc b/source/blender/nodes/function/nodes/node_fn_boolean_math.cc
index 7a83ff8e016..0ba9080918c 100644
--- a/source/blender/nodes/function/nodes/node_fn_boolean_math.cc
+++ b/source/blender/nodes/function/nodes/node_fn_boolean_math.cc
@@ -58,7 +58,7 @@ static void node_boolean_math_label(bNodeTree *UNUSED(ntree), bNode *node, char
   BLI_strncpy(label, IFACE_(name), maxlen);
 }
 
-static const blender::fn::MultiFunction &get_multi_function(bNode &bnode)
+static const blender::fn::MultiFunction *get_multi_function(bNode &bnode)
 {
   static blender::fn::CustomMF_SI_SI_SO<bool, bool, bool> and_fn{
       "And", [](bool a, bool b) { return a && b; }};
@@ -68,20 +68,21 @@ static const blender::fn::MultiFunction &get_multi_function(bNode &bnode)
 
   switch (bnode.custom1) {
     case NODE_BOOLEAN_MATH_AND:
-      return and_fn;
+      return &and_fn;
     case NODE_BOOLEAN_MATH_OR:
-      return or_fn;
+      return &or_fn;
     case NODE_BOOLEAN_MATH_NOT:
-      return not_fn;
+      return &not_fn;
   }
 
-  BLI_assert(false);
-  return blender::fn::dummy_multi_function;
+  BLI_assert_unreachable();
+  return nullptr;
 }
 
-static void node_boolean_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void fn_node_boolean_math_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  const blender::fn::MultiFunction &fn = get_multi_function(builder.bnode());
+  const blender::fn::MultiFunction *fn = get_multi_function(builder.node());
   builder.set_matching_fn(fn);
 }
 
@@ -93,7 +94,7 @@ void register_node_type_fn_boolean_math()
   node_type_socket_templates(&ntype, fn_node_boolean_math_in, fn_node_boolean_math_out);
   node_type_label(&ntype, node_boolean_math_label);
   node_type_update(&ntype, node_boolean_math_update);
-  ntype.expand_in_mf_network = node_boolean_expand_in_mf_network;
+  ntype.build_multi_function = fn_node_boolean_math_build_multi_function;
   ntype.draw_buttons = fn_node_boolean_math_layout;
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/function/nodes/node_fn_float_compare.cc b/source/blender/nodes/function/nodes/node_fn_float_compare.cc
index 6c8df8f2ea0..16ffb761a15 100644
--- a/source/blender/nodes/function/nodes/node_fn_float_compare.cc
+++ b/source/blender/nodes/function/nodes/node_fn_float_compare.cc
@@ -64,7 +64,7 @@ static void node_float_compare_label(bNodeTree *UNUSED(ntree),
   BLI_strncpy(label, IFACE_(name), maxlen);
 }
 
-static const blender::fn::MultiFunction &get_multi_function(bNode &node)
+static const blender::fn::MultiFunction *get_multi_function(bNode &node)
 {
   static blender::fn::CustomMF_SI_SI_SO<float, float, bool> less_than_fn{
       "Less Than", [](float a, float b) { return a < b; }};
@@ -81,26 +81,27 @@ static const blender::fn::MultiFunction &get_multi_function(bNode &node)
 
   switch (node.custom1) {
     case NODE_FLOAT_COMPARE_LESS_THAN:
-      return less_than_fn;
+      return &less_than_fn;
     case NODE_FLOAT_COMPARE_LESS_EQUAL:
-      return less_equal_fn;
+      return &less_equal_fn;
     case NODE_FLOAT_COMPARE_GREATER_THAN:
-      return greater_than_fn;
+      return &greater_than_fn;
     case NODE_FLOAT_COMPARE_GREATER_EQUAL:
-      return greater_equal_fn;
+      return &greater_equal_fn;
     case NODE_FLOAT_COMPARE_EQUAL:
-      return equal_fn;
+      return &equal_fn;
     case NODE_FLOAT_COMPARE_NOT_EQUAL:
-      return not_equal_fn;
+      return &not_equal_fn;
   }
 
-  BLI_assert(false);
-  return blender::fn::dummy_multi_function;
+  BLI_assert_unreachable();
+  return nullptr;
 }
 
-static void node_float_compare_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void fn_node_float_compare_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  const blender::fn::MultiFunction &fn = get_multi_function(builder.bnode());
+  const blender::fn::MultiFunction *fn = get_multi_function(builder.node());
   builder.set_matching_fn(fn);
 }
 
@@ -112,7 +113,7 @@ void register_node_type_fn_float_compare()
   node_type_socket_templates(&ntype, fn_node_float_compare_in, fn_node_float_compare_out);
   node_type_label(&ntype, node_float_compare_label);
   node_type_update(&ntype, node_float_compare_update);
-  ntype.expand_in_mf_network = node_float_compare_expand_in_mf_network;
+  ntype.build_multi_function = fn_node_float_compare_build_multi_function;
   ntype.draw_buttons = geo_node_float_compare_layout;
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/function/nodes/node_fn_float_to_int.cc b/source/blender/nodes/function/nodes/node_fn_float_to_int.cc
index 26cde576400..52acfefe615 100644
--- a/source/blender/nodes/function/nodes/node_fn_float_to_int.cc
+++ b/source/blender/nodes/function/nodes/node_fn_float_to_int.cc
@@ -50,7 +50,7 @@ static void node_float_to_int_label(bNodeTree *UNUSED(ntree), bNode *node, char
   BLI_strncpy(label, IFACE_(name), maxlen);
 }
 
-static const blender::fn::MultiFunction &get_multi_function(bNode &bnode)
+static const blender::fn::MultiFunction *get_multi_function(bNode &bnode)
 {
   static blender::fn::CustomMF_SI_SO<float, int> round_fn{"Round",
                                                           [](float a) { return (int)round(a); }};
@@ -63,22 +63,23 @@ static const blender::fn::MultiFunction &get_multi_function(bNode &bnode)
 
   switch (static_cast<FloatToIntRoundingMode>(bnode.custom1)) {
     case FN_NODE_FLOAT_TO_INT_ROUND:
-      return round_fn;
+      return &round_fn;
     case FN_NODE_FLOAT_TO_INT_FLOOR:
-      return floor_fn;
+      return &floor_fn;
     case FN_NODE_FLOAT_TO_INT_CEIL:
-      return ceil_fn;
+      return &ceil_fn;
     case FN_NODE_FLOAT_TO_INT_TRUNCATE:
-      return trunc_fn;
+      return &trunc_fn;
   }
 
   BLI_assert_unreachable();
-  return blender::fn::dummy_multi_function;
+  return nullptr;
 }
 
-static void node_float_to_int_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void fn_node_float_to_int_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  const blender::fn::MultiFunction &fn = get_multi_function(builder.bnode());
+  const blender::fn::MultiFunction *fn = get_multi_function(builder.node());
   builder.set_matching_fn(fn);
 }
 
@@ -89,7 +90,7 @@ void register_node_type_fn_float_to_int()
   fn_node_type_base(&ntype, FN_NODE_FLOAT_TO_INT, "Float to Integer", NODE_CLASS_CONVERTOR, 0);
   node_type_socket_templates(&ntype, fn_node_float_to_int_in, fn_node_float_to_int_out);
   node_type_label(&ntype, node_float_to_int_label);
-  ntype.expand_in_mf_network = node_float_to_int_expand_in_mf_network;
+  ntype.build_multi_function = fn_node_float_to_int_build_multi_function;
   ntype.draw_buttons = fn_node_float_to_int_layout;
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/function/nodes/node_fn_input_string.cc b/source/blender/nodes/function/nodes/node_fn_input_string.cc
index f16bdef2f38..560ace57aba 100644
--- a/source/blender/nodes/function/nodes/node_fn_input_string.cc
+++ b/source/blender/nodes/function/nodes/node_fn_input_string.cc
@@ -29,14 +29,14 @@ static void fn_node_input_string_layout(uiLayout *layout, bContext *UNUSED(C), P
   uiItemR(layout, ptr, "string", 0, "", ICON_NONE);
 }
 
-static void fn_node_input_string_expand_in_mf_network(
-    blender::nodes::NodeMFNetworkBuilder &builder)
+static void fn_node_input_string_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  bNode &bnode = builder.bnode();
+  bNode &bnode = builder.node();
   NodeInputString *node_storage = static_cast<NodeInputString *>(bnode.storage);
   std::string string = std::string((node_storage->string) ? node_storage->string : "");
-
-  builder.construct_and_set_matching_fn<blender::fn::CustomMF_Constant<std::string>>(string);
+  builder.construct_and_set_matching_fn<blender::fn::CustomMF_Constant<std::string>>(
+      std::move(string));
 }
 
 static void fn_node_input_string_init(bNodeTree *UNUSED(ntree), bNode *node)
@@ -78,7 +78,7 @@ void register_node_type_fn_input_string()
   node_type_socket_templates(&ntype, nullptr, fn_node_input_string_out);
   node_type_init(&ntype, fn_node_input_string_init);
   node_type_storage(&ntype, "NodeInputString", fn_node_input_string_free, fn_node_string_copy);
-  ntype.expand_in_mf_network = fn_node_input_string_expand_in_mf_network;
+  ntype.build_multi_function = fn_node_input_string_build_multi_function;
   ntype.draw_buttons = fn_node_input_string_layout;
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/function/nodes/node_fn_input_vector.cc b/source/blender/nodes/function/nodes/node_fn_input_vector.cc
index 2cd4eb1d9df..244c045de9a 100644
--- a/source/blender/nodes/function/nodes/node_fn_input_vector.cc
+++ b/source/blender/nodes/function/nodes/node_fn_input_vector.cc
@@ -32,16 +32,14 @@ static void fn_node_input_vector_layout(uiLayout *layout, bContext *UNUSED(C), P
   uiItemR(col, ptr, "vector", UI_ITEM_R_EXPAND, "", ICON_NONE);
 }
 
-static void fn_node_vector_input_expand_in_mf_network(
-    blender::nodes::NodeMFNetworkBuilder &builder)
+static void fn_node_vector_input_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  bNode &bnode = builder.bnode();
+  bNode &bnode = builder.node();
   NodeInputVector *node_storage = static_cast<NodeInputVector *>(bnode.storage);
   blender::float3 vector(node_storage->vector);
-
   builder.construct_and_set_matching_fn<blender::fn::CustomMF_Constant<blender::float3>>(vector);
 }
-
 static void fn_node_input_vector_init(bNodeTree *UNUSED(ntree), bNode *node)
 {
   NodeInputVector *data = (NodeInputVector *)MEM_callocN(sizeof(NodeInputVector),
@@ -58,7 +56,7 @@ void register_node_type_fn_input_vector()
   node_type_init(&ntype, fn_node_input_vector_init);
   node_type_storage(
       &ntype, "NodeInputVector", node_free_standard_storage, node_copy_standard_storage);
-  ntype.expand_in_mf_network = fn_node_vector_input_expand_in_mf_network;
+  ntype.build_multi_function = fn_node_vector_input_build_multi_function;
   ntype.draw_buttons = fn_node_input_vector_layout;
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/function/nodes/node_fn_random_float.cc b/source/blender/nodes/function/nodes/node_fn_random_float.cc
index a3c9f44b6a1..47ec9adf6bd 100644
--- a/source/blender/nodes/function/nodes/node_fn_random_float.cc
+++ b/source/blender/nodes/function/nodes/node_fn_random_float.cc
@@ -67,10 +67,11 @@ class RandomFloatFunction : public blender::fn::MultiFunction {
   }
 };
 
-static void fn_node_random_float_expand_in_mf_network(
-    blender::nodes::NodeMFNetworkBuilder &builder)
+static void fn_node_random_float_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  builder.construct_and_set_matching_fn<RandomFloatFunction>();
+  static RandomFloatFunction fn;
+  builder.set_matching_fn(fn);
 }
 
 void register_node_type_fn_random_float()
@@ -79,6 +80,6 @@ void register_node_type_fn_random_float()
 
   fn_node_type_base(&ntype, FN_NODE_RANDOM_FLOAT, "Random Float", 0, 0);
   node_type_socket_templates(&ntype, fn_node_random_float_in, fn_node_random_float_out);
-  ntype.expand_in_mf_network = fn_node_random_float_expand_in_mf_network;
+  ntype.build_multi_function = fn_node_random_float_build_multi_function;
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/intern/node_geometry_exec.cc b/source/blender/nodes/intern/node_geometry_exec.cc
index ffa20579acc..a3bbca90731 100644
--- a/source/blender/nodes/intern/node_geometry_exec.cc
+++ b/source/blender/nodes/intern/node_geometry_exec.cc
@@ -19,7 +19,6 @@
 #include "DEG_depsgraph_query.h"
 
 #include "NOD_geometry_exec.hh"
-#include "NOD_type_callbacks.hh"
 #include "NOD_type_conversions.hh"
 
 #include "node_geometry_util.hh"
diff --git a/source/blender/nodes/NOD_type_callbacks.hh b/source/blender/nodes/intern/node_multi_function.cc
index 2be78f929db..c91899ed8c2 100644
--- a/source/blender/nodes/NOD_type_callbacks.hh
+++ b/source/blender/nodes/intern/node_multi_function.cc
@@ -14,21 +14,27 @@
  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  */
 
-#pragma once
-
-#include <optional>
-
-#include "BKE_node.h"
-
-#include "FN_multi_function_data_type.hh"
+#include "NOD_multi_function.hh"
 
 namespace blender::nodes {
 
-using fn::CPPType;
-using fn::MFDataType;
-
-std::optional<MFDataType> socket_mf_type_get(const bNodeSocketType &stype);
-bool socket_is_mf_data_socket(const bNodeSocketType &stype);
-void socket_expand_in_mf_network(SocketMFNetworkBuilder &builder);
+NodeMultiFunctions::NodeMultiFunctions(const DerivedNodeTree &tree, ResourceScope &resource_scope)
+{
+  for (const NodeTreeRef *tree_ref : tree.used_node_tree_refs()) {
+    bNodeTree *btree = tree_ref->btree();
+    for (const NodeRef *node : tree_ref->nodes()) {
+      bNode *bnode = node->bnode();
+      if (bnode->typeinfo->build_multi_function == nullptr) {
+        continue;
+      }
+      NodeMultiFunctionBuilder builder{resource_scope, *bnode, *btree};
+      bnode->typeinfo->build_multi_function(builder);
+      const MultiFunction *fn = builder.built_fn_;
+      if (fn != nullptr) {
+        map_.add_new(bnode, fn);
+      }
+    }
+  }
+}
 
 }  // namespace blender::nodes
diff --git a/source/blender/nodes/intern/node_socket.cc b/source/blender/nodes/intern/node_socket.cc
index 4be3fd2468b..528616eb23a 100644
--- a/source/blender/nodes/intern/node_socket.cc
+++ b/source/blender/nodes/intern/node_socket.cc
@@ -44,7 +44,6 @@
 
 #include "MEM_guardedalloc.h"
 
-#include "NOD_node_tree_multi_function.hh"
 #include "NOD_socket.h"
 
 #include "FN_cpp_type_make.hh"
diff --git a/source/blender/nodes/intern/node_tree_multi_function.cc b/source/blender/nodes/intern/node_tree_multi_function.cc
deleted file mode 100644
index 7ab6495f733..00000000000
--- a/source/blender/nodes/intern/node_tree_multi_function.cc
+++ /dev/null
@@ -1,409 +0,0 @@
-/*
- * This program is free software; you can redistribute it and/or
- * modify it under the terms of the GNU General Public License
- * as published by the Free Software Foundation; either version 2
- * of the License, or (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program; if not, write to the Free Software Foundation,
- * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
- */
-
-#include "NOD_node_tree_multi_function.hh"
-#include "NOD_type_conversions.hh"
-
-#include "FN_multi_function_network_evaluation.hh"
-
-#include "BLI_color.hh"
-#include "BLI_float2.hh"
-#include "BLI_float3.hh"
-
-namespace blender::nodes {
-
-const fn::MultiFunction &NodeMFNetworkBuilder::get_default_fn(StringRef name)
-{
-  Vector<fn::MFDataType, 10> input_types;
-  Vector<fn::MFDataType, 10> output_types;
-
-  for (const InputSocketRef *dsocket : dnode_->inputs()) {
-    if (dsocket->is_available()) {
-      std::optional<fn::MFDataType> data_type = socket_mf_type_get(*dsocket->typeinfo());
-      if (data_type.has_value()) {
-        input_types.append(*data_type);
-      }
-    }
-  }
-  for (const OutputSocketRef *dsocket : dnode_->outputs()) {
-    if (dsocket->is_available()) {
-      std::optional<fn::MFDataType> data_type = socket_mf_type_get(*dsocket->typeinfo());
-      if (data_type.has_value()) {
-        output_types.append(*data_type);
-      }
-    }
-  }
-
-  const fn::MultiFunction &fn = this->construct_fn<fn::CustomMF_DefaultOutput>(
-      name, input_types, output_types);
-  return fn;
-}
-
-static void insert_dummy_node(CommonMFNetworkBuilderData &common, const DNode &dnode)
-{
-  constexpr int stack_capacity = 10;
-
-  Vector<fn::MFDataType, stack_capacity> input_types;
-  Vector<StringRef, stack_capacity> input_names;
-  Vector<const InputSocketRef *, stack_capacity> input_dsockets;
-
-  for (const InputSocketRef *dsocket : dnode->inputs()) {
-    if (dsocket->is_available()) {
-      std::optional<fn::MFDataType> data_type = socket_mf_type_get(*dsocket->bsocket()->typeinfo);
-      if (data_type.has_value()) {
-        input_types.append(*data_type);
-        input_names.append(dsocket->name());
-        input_dsockets.append(dsocket);
-      }
-    }
-  }
-
-  Vector<fn::MFDataType, stack_capacity> output_types;
-  Vector<StringRef, stack_capacity> output_names;
-  Vector<const OutputSocketRef *, stack_capacity> output_dsockets;
-
-  for (const OutputSocketRef *dsocket : dnode->outputs()) {
-    if (dsocket->is_available()) {
-      std::optional<fn::MFDataType> data_type = socket_mf_type_get(*dsocket->bsocket()->typeinfo);
-      if (data_type.has_value()) {
-        output_types.append(*data_type);
-        output_names.append(dsocket->name());
-        output_dsockets.append(dsocket);
-      }
-    }
-  }
-
-  fn::MFDummyNode &dummy_node = common.network.add_dummy(
-      dnode->name(), input_types, output_types, input_names, output_names);
-
-  common.network_map.add(*dnode.context(), input_dsockets, dummy_node.inputs());
-  common.network_map.add(*dnode.context(), output_dsockets, dummy_node.outputs());
-}
-
-static bool has_data_sockets(const DNode &dnode)
-{
-  for (const InputSocketRef *socket : dnode->inputs()) {
-    if (socket_is_mf_data_socket(*socket->bsocket()->typeinfo)) {
-      return true;
-    }
-  }
-  for (const OutputSocketRef *socket : dnode->outputs()) {
-    if (socket_is_mf_data_socket(*socket->bsocket()->typeinfo)) {
-      return true;
-    }
-  }
-  return false;
-}
-
-static void foreach_node_to_insert(CommonMFNetworkBuilderData &common,
-                                   FunctionRef<void(DNode)> callback)
-{
-  common.tree.foreach_node([&](const DNode dnode) {
-    if (dnode->is_group_node()) {
-      return;
-    }
-    /* Don't insert non-root group input/output nodes, because they will be inlined. */
-    if (!dnode.context()->is_root()) {
-      if (dnode->is_group_input_node() || dnode->is_group_output_node()) {
-        return;
-      }
-    }
-    callback(dnode);
-  });
-}
-
-/**
- * Expands all function nodes in the multi-function network. Nodes that don't have an expand
- * function, but do have data sockets, will get corresponding dummy nodes.
- */
-static void insert_nodes(CommonMFNetworkBuilderData &common)
-{
-  foreach_node_to_insert(common, [&](const DNode dnode) {
-    const bNodeType *node_type = dnode->typeinfo();
-    if (node_type->expand_in_mf_network != nullptr) {
-      NodeMFNetworkBuilder builder{common, dnode};
-      node_type->expand_in_mf_network(builder);
-    }
-    else if (has_data_sockets(dnode)) {
-      insert_dummy_node(common, dnode);
-    }
-  });
-}
-
-static fn::MFOutputSocket &insert_default_value_for_type(CommonMFNetworkBuilderData &common,
-                                                         fn::MFDataType type)
-{
-  const fn::MultiFunction *default_fn;
-  if (type.is_single()) {
-    default_fn = &common.scope.construct<fn::CustomMF_GenericConstant>(
-        AT, type.single_type(), type.single_type().default_value());
-  }
-  else {
-    default_fn = &common.scope.construct<fn::CustomMF_GenericConstantArray>(
-        AT, fn::GSpan(type.vector_base_type()));
-  }
-
-  fn::MFNode &node = common.network.add_function(*default_fn);
-  return node.output(0);
-}
-
-static fn::MFOutputSocket *insert_unlinked_input(CommonMFNetworkBuilderData &common,
-                                                 const DInputSocket &dsocket)
-{
-  BLI_assert(socket_is_mf_data_socket(*dsocket->typeinfo()));
-
-  SocketMFNetworkBuilder builder{common, dsocket};
-  socket_expand_in_mf_network(builder);
-
-  fn::MFOutputSocket *built_socket = builder.built_socket();
-  BLI_assert(built_socket != nullptr);
-  return built_socket;
-}
-
-static void insert_links_and_unlinked_inputs(CommonMFNetworkBuilderData &common)
-{
-  foreach_node_to_insert(common, [&](const DNode dnode) {
-    for (const InputSocketRef *socket_ref : dnode->inputs()) {
-      const DInputSocket to_dsocket{dnode.context(), socket_ref};
-      if (!to_dsocket->is_available()) {
-        continue;
-      }
-      if (!socket_is_mf_data_socket(*to_dsocket->typeinfo())) {
-        continue;
-      }
-
-      Span<fn::MFInputSocket *> to_sockets = common.network_map.lookup(to_dsocket);
-      BLI_assert(to_sockets.size() >= 1);
-      const fn::MFDataType to_type = to_sockets[0]->data_type();
-
-      Vector<DSocket> from_dsockets;
-      to_dsocket.foreach_origin_socket([&](DSocket socket) { from_dsockets.append(socket); });
-      if (from_dsockets.size() > 1) {
-        fn::MFOutputSocket &from_socket = insert_default_value_for_type(common, to_type);
-        for (fn::MFInputSocket *to_socket : to_sockets) {
-          common.network.add_link(from_socket, *to_socket);
-        }
-        continue;
-      }
-      if (from_dsockets.is_empty()) {
-        /* The socket is not linked. Need to use the value of the socket itself. */
-        fn::MFOutputSocket *built_socket = insert_unlinked_input(common, to_dsocket);
-        for (fn::MFInputSocket *to_socket : to_sockets) {
-          common.network.add_link(*built_socket, *to_socket);
-        }
-        continue;
-      }
-      if (from_dsockets[0]->is_input()) {
-        DInputSocket from_dsocket{from_dsockets[0]};
-        fn::MFOutputSocket *built_socket = insert_unlinked_input(common, from_dsocket);
-        for (fn::MFInputSocket *to_socket : to_sockets) {
-          common.network.add_link(*built_socket, *to_socket);
-        }
-        continue;
-      }
-      DOutputSocket from_dsocket{from_dsockets[0]};
-      fn::MFOutputSocket *from_socket = &common.network_map.lookup(from_dsocket);
-      const fn::MFDataType from_type = from_socket->data_type();
-
-      if (from_type != to_type) {
-        const fn::MultiFunction *conversion_fn =
-            get_implicit_type_conversions().get_conversion_multi_function(from_type, to_type);
-        if (conversion_fn != nullptr) {
-          fn::MFNode &node = common.network.add_function(*conversion_fn);
-          common.network.add_link(*from_socket, node.input(0));
-          from_socket = &node.output(0);
-        }
-        else {
-          from_socket = &insert_default_value_for_type(common, to_type);
-        }
-      }
-
-      for (fn::MFInputSocket *to_socket : to_sockets) {
-        common.network.add_link(*from_socket, *to_socket);
-      }
-    }
-  });
-}
-
-/**
- * Expands all function nodes contained in the given node tree within the given multi-function
- * network.
- *
- * Returns a mapping between the original node tree and the generated nodes/sockets for further
- * processing.
- */
-MFNetworkTreeMap insert_node_tree_into_mf_network(fn::MFNetwork &network,
-                                                  const DerivedNodeTree &tree,
-                                                  ResourceScope &scope)
-{
-  MFNetworkTreeMap network_map{tree, network};
-
-  CommonMFNetworkBuilderData common{scope, network, network_map, tree};
-
-  insert_nodes(common);
-  insert_links_and_unlinked_inputs(common);
-
-  return network_map;
-}
-
-/**
- * A single node is allowed to expand into multiple nodes before evaluation. Depending on what
- * nodes it expands to, it belongs a different type of the ones below.
- */
-enum class NodeExpandType {
-  SingleFunctionNode,
-  MultipleFunctionNodes,
-  HasDummyNodes,
-};
-
-/**
- * Checks how the given node expanded in the multi-function network. If it is only a single
- * function node, the corresponding function is returned as well.
- */
-static NodeExpandType get_node_expand_type(MFNetworkTreeMap &network_map,
-                                           const DNode &dnode,
-                                           const fn::MultiFunction **r_single_function)
-{
-  const fn::MFFunctionNode *single_function_node = nullptr;
-  bool has_multiple_nodes = false;
-  bool has_dummy_nodes = false;
-
-  auto check_mf_node = [&](fn::MFNode &mf_node) {
-    if (mf_node.is_function()) {
-      if (single_function_node == nullptr) {
-        single_function_node = &mf_node.as_function();
-      }
-      if (&mf_node != single_function_node) {
-        has_multiple_nodes = true;
-      }
-    }
-    else {
-      BLI_assert(mf_node.is_dummy());
-      has_dummy_nodes = true;
-    }
-  };
-
-  for (const InputSocketRef *dsocket : dnode->inputs()) {
-    if (dsocket->is_available()) {
-      for (fn::MFInputSocket *mf_input :
-           network_map.lookup(DInputSocket(dnode.context(), dsocket))) {
-        check_mf_node(mf_input->node());
-      }
-    }
-  }
-  for (const OutputSocketRef *dsocket : dnode->outputs()) {
-    if (dsocket->is_available()) {
-      fn::MFOutputSocket &mf_output = network_map.lookup(DOutputSocket(dnode.context(), dsocket));
-      check_mf_node(mf_output.node());
-    }
-  }
-
-  if (has_dummy_nodes) {
-    return NodeExpandType::HasDummyNodes;
-  }
-  if (has_multiple_nodes) {
-    return NodeExpandType::MultipleFunctionNodes;
-  }
-  *r_single_function = &single_function_node->function();
-  return NodeExpandType::SingleFunctionNode;
-}
-
-static const fn::MultiFunction &create_function_for_node_that_expands_into_multiple(
-    const DNode &dnode,
-    fn::MFNetwork &network,
-    MFNetworkTreeMap &network_map,
-    ResourceScope &scope)
-{
-  Vector<const fn::MFOutputSocket *> dummy_fn_inputs;
-  for (const InputSocketRef *dsocket : dnode->inputs()) {
-    if (dsocket->is_available()) {
-      MFDataType data_type = *socket_mf_type_get(*dsocket->typeinfo());
-      fn::MFOutputSocket &fn_input = network.add_input(data_type.to_string(), data_type);
-      for (fn::MFInputSocket *mf_input :
-           network_map.lookup(DInputSocket(dnode.context(), dsocket))) {
-        network.add_link(fn_input, *mf_input);
-        dummy_fn_inputs.append(&fn_input);
-      }
-    }
-  }
-  Vector<const fn::MFInputSocket *> dummy_fn_outputs;
-  for (const OutputSocketRef *dsocket : dnode->outputs()) {
-    if (dsocket->is_available()) {
-      fn::MFOutputSocket &mf_output = network_map.lookup(DOutputSocket(dnode.context(), dsocket));
-      MFDataType data_type = mf_output.data_type();
-      fn::MFInputSocket &fn_output = network.add_output(data_type.to_string(), data_type);
-      network.add_link(mf_output, fn_output);
-      dummy_fn_outputs.append(&fn_output);
-    }
-  }
-
-  fn::MFNetworkEvaluator &fn_evaluator = scope.construct<fn::MFNetworkEvaluator>(
-      __func__, std::move(dummy_fn_inputs), std::move(dummy_fn_outputs));
-  return fn_evaluator;
-}
-
-/**
- * Returns a single multi-function for every node that supports it. This makes it easier to reuse
- * the multi-function implementation of nodes in different contexts.
- */
-MultiFunctionByNode get_multi_function_per_node(const DerivedNodeTree &tree, ResourceScope &scope)
-{
-  /* Build a network that nodes can insert themselves into. However, the individual nodes are not
-   * connected. */
-  fn::MFNetwork &network = scope.construct<fn::MFNetwork>(__func__);
-  MFNetworkTreeMap network_map{tree, network};
-  MultiFunctionByNode functions_by_node;
-
-  CommonMFNetworkBuilderData common{scope, network, network_map, tree};
-
-  tree.foreach_node([&](DNode dnode) {
-    const bNodeType *node_type = dnode->typeinfo();
-    if (node_type->expand_in_mf_network == nullptr) {
-      /* This node does not have a multi-function implementation. */
-      return;
-    }
-
-    NodeMFNetworkBuilder builder{common, dnode};
-    node_type->expand_in_mf_network(builder);
-
-    const fn::MultiFunction *single_function = nullptr;
-    const NodeExpandType expand_type = get_node_expand_type(network_map, dnode, &single_function);
-
-    switch (expand_type) {
-      case NodeExpandType::HasDummyNodes: {
-        /* Dummy nodes cannot be executed, so skip them. */
-        break;
-      }
-      case NodeExpandType::SingleFunctionNode: {
-        /* This is the common case. Most nodes just expand to a single function. */
-        functions_by_node.add_new(dnode, single_function);
-        break;
-      }
-      case NodeExpandType::MultipleFunctionNodes: {
-        /* If a node expanded into multiple functions, a new function has to be created that
-         * combines those. */
-        const fn::MultiFunction &fn = create_function_for_node_that_expands_into_multiple(
-            dnode, network, network_map, scope);
-        functions_by_node.add_new(dnode, &fn);
-        break;
-      }
-    }
-  });
-
-  return functions_by_node;
-}
-
-}  // namespace blender::nodes
diff --git a/source/blender/nodes/intern/type_callbacks.cc b/source/blender/nodes/intern/type_callbacks.cc
deleted file mode 100644
index 881a02c92e9..00000000000
--- a/source/blender/nodes/intern/type_callbacks.cc
+++ /dev/null
@@ -1,60 +0,0 @@
-/*
- * This program is free software; you can redistribute it and/or
- * modify it under the terms of the GNU General Public License
- * as published by the Free Software Foundation; either version 2
- * of the License, or (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program; if not, write to the Free Software Foundation,
- * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
- */
-
-#include "NOD_node_tree_multi_function.hh"
-#include "NOD_type_callbacks.hh"
-
-namespace blender::nodes {
-
-std::optional<MFDataType> socket_mf_type_get(const bNodeSocketType &stype)
-{
-  const CPPType *cpp_type = stype.get_base_cpp_type ? stype.get_base_cpp_type() : nullptr;
-  if (cpp_type != nullptr) {
-    return MFDataType::ForSingle(*cpp_type);
-  }
-  return {};
-}
-
-bool socket_is_mf_data_socket(const bNodeSocketType &stype)
-{
-  if (!socket_mf_type_get(stype).has_value()) {
-    return false;
-  }
-  if (stype.expand_in_mf_network == nullptr && stype.get_base_cpp_value == nullptr) {
-    return false;
-  }
-  return true;
-}
-
-void socket_expand_in_mf_network(SocketMFNetworkBuilder &builder)
-{
-  bNodeSocket &socket = builder.bsocket();
-  if (socket.typeinfo->expand_in_mf_network != nullptr) {
-    socket.typeinfo->expand_in_mf_network(builder);
-  }
-  else if (socket.typeinfo->get_base_cpp_value != nullptr) {
-    const CPPType &type = *socket.typeinfo->get_base_cpp_type();
-    void *buffer = builder.resource_scope().linear_allocator().allocate(type.size(),
-                                                                        type.alignment());
-    socket.typeinfo->get_base_cpp_value(socket, buffer);
-    builder.set_constant_value(type, buffer);
-  }
-  else {
-    BLI_assert_unreachable();
-  }
-}
-
-}  // namespace blender::nodes
diff --git a/source/blender/nodes/shader/node_shader_util.h b/source/blender/nodes/shader/node_shader_util.h
index dc44f0fa98f..a75354d3381 100644
--- a/source/blender/nodes/shader/node_shader_util.h
+++ b/source/blender/nodes/shader/node_shader_util.h
@@ -72,7 +72,7 @@
 #ifdef __cplusplus
 #  include "FN_multi_function_builder.hh"
 
-#  include "NOD_node_tree_multi_function.hh"
+#  include "NOD_multi_function.hh"
 
 #  include "BLI_color.hh"
 #  include "BLI_float3.hh"
diff --git a/source/blender/nodes/shader/nodes/node_shader_clamp.cc b/source/blender/nodes/shader/nodes/node_shader_clamp.cc
index 4f77421cfe0..f105f8bcaf9 100644
--- a/source/blender/nodes/shader/nodes/node_shader_clamp.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_clamp.cc
@@ -51,7 +51,7 @@ static int gpu_shader_clamp(GPUMaterial *mat,
              GPU_stack_link(mat, node, "clamp_range", in, out);
 }
 
-static void sh_node_clamp_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_clamp_build_multi_function(blender::nodes::NodeMultiFunctionBuilder &builder)
 {
   static blender::fn::CustomMF_SI_SI_SI_SO<float, float, float, float> minmax_fn{
       "Clamp (Min Max)",
@@ -65,7 +65,7 @@ static void sh_node_clamp_expand_in_mf_network(blender::nodes::NodeMFNetworkBuil
         return clamp_f(value, b, a);
       }};
 
-  int clamp_type = builder.bnode().custom1;
+  int clamp_type = builder.node().custom1;
   if (clamp_type == NODE_CLAMP_MINMAX) {
     builder.set_matching_fn(minmax_fn);
   }
@@ -82,7 +82,7 @@ void register_node_type_sh_clamp(void)
   node_type_socket_templates(&ntype, sh_node_clamp_in, sh_node_clamp_out);
   node_type_init(&ntype, node_shader_init_clamp);
   node_type_gpu(&ntype, gpu_shader_clamp);
-  ntype.expand_in_mf_network = sh_node_clamp_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_clamp_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_curves.cc b/source/blender/nodes/shader/nodes/node_shader_curves.cc
index f1d5040a292..df075d6e973 100644
--- a/source/blender/nodes/shader/nodes/node_shader_curves.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_curves.cc
@@ -143,9 +143,10 @@ class CurveVecFunction : public blender::fn::MultiFunction {
   }
 };
 
-static void sh_node_curve_vec_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_curve_vec_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  bNode &bnode = builder.bnode();
+  bNode &bnode = builder.node();
   CurveMapping *cumap = (CurveMapping *)bnode.storage;
   BKE_curvemapping_init(cumap);
   builder.construct_and_set_matching_fn<CurveVecFunction>(*cumap);
@@ -162,7 +163,7 @@ void register_node_type_sh_curve_vec(void)
   node_type_storage(&ntype, "CurveMapping", node_free_curves, node_copy_curves);
   node_type_exec(&ntype, node_initexec_curves, nullptr, node_shader_exec_curve_vec);
   node_type_gpu(&ntype, gpu_shader_curve_vec);
-  ntype.expand_in_mf_network = sh_node_curve_vec_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_curve_vec_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
@@ -317,9 +318,10 @@ class CurveRGBFunction : public blender::fn::MultiFunction {
   }
 };
 
-static void sh_node_curve_rgb_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_curve_rgb_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  bNode &bnode = builder.bnode();
+  bNode &bnode = builder.node();
   CurveMapping *cumap = (CurveMapping *)bnode.storage;
   BKE_curvemapping_init(cumap);
   builder.construct_and_set_matching_fn<CurveRGBFunction>(*cumap);
@@ -336,7 +338,7 @@ void register_node_type_sh_curve_rgb(void)
   node_type_storage(&ntype, "CurveMapping", node_free_curves, node_copy_curves);
   node_type_exec(&ntype, node_initexec_curves, nullptr, node_shader_exec_curve_rgb);
   node_type_gpu(&ntype, gpu_shader_curve_rgb);
-  ntype.expand_in_mf_network = sh_node_curve_rgb_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_curve_rgb_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_map_range.cc b/source/blender/nodes/shader/nodes/node_shader_map_range.cc
index ad7abd9d491..e4739e2864d 100644
--- a/source/blender/nodes/shader/nodes/node_shader_map_range.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_map_range.cc
@@ -261,9 +261,10 @@ class MapRangeSmootherstepFunction : public blender::fn::MultiFunction {
   }
 };
 
-static void sh_node_map_range_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_map_range_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  bNode &bnode = builder.bnode();
+  bNode &bnode = builder.node();
   bool clamp = bnode.custom1 != 0;
   int interpolation_type = bnode.custom2;
 
@@ -301,7 +302,6 @@ static void sh_node_map_range_expand_in_mf_network(blender::nodes::NodeMFNetwork
       break;
     }
     default:
-      builder.set_not_implemented();
       break;
   }
 }
@@ -315,7 +315,7 @@ void register_node_type_sh_map_range(void)
   node_type_init(&ntype, node_shader_init_map_range);
   node_type_update(&ntype, node_shader_update_map_range);
   node_type_gpu(&ntype, gpu_shader_map_range);
-  ntype.expand_in_mf_network = sh_node_map_range_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_map_range_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_math.cc b/source/blender/nodes/shader/nodes/node_shader_math.cc
index 7a846031456..c30f2948ab1 100644
--- a/source/blender/nodes/shader/nodes/node_shader_math.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_math.cc
@@ -69,11 +69,9 @@ static int gpu_shader_math(GPUMaterial *mat,
   return 0;
 }
 
-static const blender::fn::MultiFunction &get_base_multi_function(
-    blender::nodes::NodeMFNetworkBuilder &builder)
+static const blender::fn::MultiFunction *get_base_multi_function(bNode &node)
 {
-  const int mode = builder.bnode().custom1;
-
+  const int mode = node.custom1;
   const blender::fn::MultiFunction *base_fn = nullptr;
 
   blender::nodes::try_dispatch_float_math_fl_to_fl(
@@ -82,7 +80,7 @@ static const blender::fn::MultiFunction &get_base_multi_function(
         base_fn = &fn;
       });
   if (base_fn != nullptr) {
-    return *base_fn;
+    return base_fn;
   }
 
   blender::nodes::try_dispatch_float_math_fl_fl_to_fl(
@@ -92,7 +90,7 @@ static const blender::fn::MultiFunction &get_base_multi_function(
         base_fn = &fn;
       });
   if (base_fn != nullptr) {
-    return *base_fn;
+    return base_fn;
   }
 
   blender::nodes::try_dispatch_float_math_fl_fl_fl_to_fl(
@@ -102,36 +100,51 @@ static const blender::fn::MultiFunction &get_base_multi_function(
         base_fn = &fn;
       });
   if (base_fn != nullptr) {
-    return *base_fn;
+    return base_fn;
   }
 
-  return builder.get_not_implemented_fn();
+  return nullptr;
 }
 
-static void sh_node_math_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
-{
-  const blender::fn::MultiFunction &base_function = get_base_multi_function(builder);
+class ClampWrapperFunction : public blender::fn::MultiFunction {
+ private:
+  const blender::fn::MultiFunction &fn_;
 
-  const blender::nodes::DNode &dnode = builder.dnode();
-  blender::fn::MFNetwork &network = builder.network();
-  blender::fn::MFFunctionNode &base_node = network.add_function(base_function);
+ public:
+  ClampWrapperFunction(const blender::fn::MultiFunction &fn) : fn_(fn)
+  {
+    this->set_signature(&fn.signature());
+  }
 
-  builder.network_map().add_try_match(*dnode.context(), dnode->inputs(), base_node.inputs());
+  void call(blender::IndexMask mask,
+            blender::fn::MFParams params,
+            blender::fn::MFContext context) const override
+  {
+    fn_.call(mask, params, context);
+
+    /* Assumes the output parameter is the last one. */
+    const int output_param_index = this->param_amount() - 1;
+    /* This has actually been initialized in the call above. */
+    blender::MutableSpan<float> results = params.uninitialized_single_output<float>(
+        output_param_index);
+
+    for (const int i : mask) {
+      float &value = results[i];
+      CLAMP(value, 0.0f, 1.0f);
+    }
+  }
+};
+
+static void sh_node_math_build_multi_function(blender::nodes::NodeMultiFunctionBuilder &builder)
+{
+  const blender::fn::MultiFunction *base_function = get_base_multi_function(builder.node());
 
-  const bool clamp_output = builder.bnode().custom2 != 0;
+  const bool clamp_output = builder.node().custom2 != 0;
   if (clamp_output) {
-    static blender::fn::CustomMF_SI_SO<float, float> clamp_fn{"Clamp", [](float value) {
-                                                                CLAMP(value, 0.0f, 1.0f);
-                                                                return value;
-                                                              }};
-    blender::fn::MFFunctionNode &clamp_node = network.add_function(clamp_fn);
-    network.add_link(base_node.output(0), clamp_node.input(0));
-    builder.network_map().add(blender::nodes::DOutputSocket(dnode.context(), &dnode->output(0)),
-                              clamp_node.output(0));
+    builder.construct_and_set_matching_fn<ClampWrapperFunction>(*base_function);
   }
   else {
-    builder.network_map().add(blender::nodes::DOutputSocket(dnode.context(), &dnode->output(0)),
-                              base_node.output(0));
+    builder.set_matching_fn(base_function);
   }
 }
 
@@ -144,7 +157,7 @@ void register_node_type_sh_math(void)
   node_type_label(&ntype, node_math_label);
   node_type_gpu(&ntype, gpu_shader_math);
   node_type_update(&ntype, node_math_update);
-  ntype.expand_in_mf_network = sh_node_math_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_math_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_mixRgb.cc b/source/blender/nodes/shader/nodes/node_shader_mixRgb.cc
index 47011caeeb6..ade35a40366 100644
--- a/source/blender/nodes/shader/nodes/node_shader_mixRgb.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_mixRgb.cc
@@ -174,9 +174,9 @@ class MixRGBFunction : public blender::fn::MultiFunction {
   }
 };
 
-static void sh_node_mix_rgb_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_mix_rgb_build_multi_function(blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  bNode &node = builder.bnode();
+  bNode &node = builder.node();
   bool clamp = node.custom2 & SHD_MIXRGB_CLAMP;
   int mix_type = node.custom1;
   builder.construct_and_set_matching_fn<MixRGBFunction>(clamp, mix_type);
@@ -191,7 +191,7 @@ void register_node_type_sh_mix_rgb(void)
   node_type_label(&ntype, node_blend_label);
   node_type_exec(&ntype, nullptr, nullptr, node_shader_exec_mix_rgb);
   node_type_gpu(&ntype, gpu_shader_mix_rgb);
-  ntype.expand_in_mf_network = sh_node_mix_rgb_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_mix_rgb_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_sepcombRGB.cc b/source/blender/nodes/shader/nodes/node_shader_sepcombRGB.cc
index a7239154633..2779fc6bf68 100644
--- a/source/blender/nodes/shader/nodes/node_shader_sepcombRGB.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_sepcombRGB.cc
@@ -96,7 +96,7 @@ class SeparateRGBFunction : public blender::fn::MultiFunction {
   }
 };
 
-static void sh_node_seprgb_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_seprgb_build_multi_function(blender::nodes::NodeMultiFunctionBuilder &builder)
 {
   static SeparateRGBFunction fn;
   builder.set_matching_fn(fn);
@@ -110,7 +110,7 @@ void register_node_type_sh_seprgb(void)
   node_type_socket_templates(&ntype, sh_node_seprgb_in, sh_node_seprgb_out);
   node_type_exec(&ntype, nullptr, nullptr, node_shader_exec_seprgb);
   node_type_gpu(&ntype, gpu_shader_seprgb);
-  ntype.expand_in_mf_network = sh_node_seprgb_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_seprgb_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
@@ -153,7 +153,7 @@ static int gpu_shader_combrgb(GPUMaterial *mat,
   return GPU_stack_link(mat, node, "combine_rgb", in, out);
 }
 
-static void sh_node_combrgb_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_combrgb_build_multi_function(blender::nodes::NodeMultiFunctionBuilder &builder)
 {
   static blender::fn::CustomMF_SI_SI_SI_SO<float, float, float, blender::ColorGeometry4f> fn{
       "Combine RGB",
@@ -169,7 +169,7 @@ void register_node_type_sh_combrgb(void)
   node_type_socket_templates(&ntype, sh_node_combrgb_in, sh_node_combrgb_out);
   node_type_exec(&ntype, nullptr, nullptr, node_shader_exec_combrgb);
   node_type_gpu(&ntype, gpu_shader_combrgb);
-  ntype.expand_in_mf_network = sh_node_combrgb_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_combrgb_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_sepcombXYZ.cc b/source/blender/nodes/shader/nodes/node_shader_sepcombXYZ.cc
index efa9581c414..1fd794cdd0a 100644
--- a/source/blender/nodes/shader/nodes/node_shader_sepcombXYZ.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_sepcombXYZ.cc
@@ -81,7 +81,7 @@ class MF_SeparateXYZ : public blender::fn::MultiFunction {
   }
 };
 
-static void sh_node_sepxyz_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_sepxyz_build_multi_function(blender::nodes::NodeMultiFunctionBuilder &builder)
 {
   static MF_SeparateXYZ separate_fn;
   builder.set_matching_fn(separate_fn);
@@ -94,7 +94,7 @@ void register_node_type_sh_sepxyz(void)
   sh_fn_node_type_base(&ntype, SH_NODE_SEPXYZ, "Separate XYZ", NODE_CLASS_CONVERTOR, 0);
   node_type_socket_templates(&ntype, sh_node_sepxyz_in, sh_node_sepxyz_out);
   node_type_gpu(&ntype, gpu_shader_sepxyz);
-  ntype.expand_in_mf_network = sh_node_sepxyz_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_sepxyz_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
@@ -120,7 +120,7 @@ static int gpu_shader_combxyz(GPUMaterial *mat,
   return GPU_stack_link(mat, node, "combine_xyz", in, out);
 }
 
-static void sh_node_combxyz_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_combxyz_build_multi_function(blender::nodes::NodeMultiFunctionBuilder &builder)
 {
   static blender::fn::CustomMF_SI_SI_SI_SO<float, float, float, blender::float3> fn{
       "Combine Vector", [](float x, float y, float z) { return blender::float3(x, y, z); }};
@@ -134,7 +134,7 @@ void register_node_type_sh_combxyz(void)
   sh_fn_node_type_base(&ntype, SH_NODE_COMBXYZ, "Combine XYZ", NODE_CLASS_CONVERTOR, 0);
   node_type_socket_templates(&ntype, sh_node_combxyz_in, sh_node_combxyz_out);
   node_type_gpu(&ntype, gpu_shader_combxyz);
-  ntype.expand_in_mf_network = sh_node_combxyz_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_combxyz_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_valToRgb.cc b/source/blender/nodes/shader/nodes/node_shader_valToRgb.cc
index 5b2eb300aac..1bc42ab0cc6 100644
--- a/source/blender/nodes/shader/nodes/node_shader_valToRgb.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_valToRgb.cc
@@ -163,9 +163,10 @@ class ColorBandFunction : public blender::fn::MultiFunction {
   }
 };
 
-static void sh_node_valtorgb_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_valtorgb_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  bNode &bnode = builder.bnode();
+  bNode &bnode = builder.node();
   const ColorBand *color_band = (const ColorBand *)bnode.storage;
   builder.construct_and_set_matching_fn<ColorBandFunction>(*color_band);
 }
@@ -181,7 +182,7 @@ void register_node_type_sh_valtorgb(void)
   node_type_storage(&ntype, "ColorBand", node_free_standard_storage, node_copy_standard_storage);
   node_type_exec(&ntype, nullptr, nullptr, node_shader_exec_valtorgb);
   node_type_gpu(&ntype, gpu_shader_valtorgb);
-  ntype.expand_in_mf_network = sh_node_valtorgb_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_valtorgb_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_value.cc b/source/blender/nodes/shader/nodes/node_shader_value.cc
index 495c8d12824..602d5a1cf56 100644
--- a/source/blender/nodes/shader/nodes/node_shader_value.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_value.cc
@@ -39,9 +39,9 @@ static int gpu_shader_value(GPUMaterial *mat,
   return GPU_stack_link(mat, node, "set_value", in, out, link);
 }
 
-static void sh_node_value_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_value_build_multi_function(blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  const bNodeSocket *bsocket = builder.dnode()->output(0).bsocket();
+  const bNodeSocket *bsocket = (bNodeSocket *)builder.node().outputs.first;
   const bNodeSocketValueFloat *value = (const bNodeSocketValueFloat *)bsocket->default_value;
   builder.construct_and_set_matching_fn<blender::fn::CustomMF_Constant<float>>(value->value);
 }
@@ -53,7 +53,7 @@ void register_node_type_sh_value(void)
   sh_fn_node_type_base(&ntype, SH_NODE_VALUE, "Value", NODE_CLASS_INPUT, 0);
   node_type_socket_templates(&ntype, nullptr, sh_node_value_out);
   node_type_gpu(&ntype, gpu_shader_value);
-  ntype.expand_in_mf_network = sh_node_value_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_value_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_vector_math.cc b/source/blender/nodes/shader/nodes/node_shader_vector_math.cc
index 419a11201aa..4424db6aed1 100644
--- a/source/blender/nodes/shader/nodes/node_shader_vector_math.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_vector_math.cc
@@ -183,12 +183,11 @@ static void node_shader_update_vector_math(bNodeTree *UNUSED(ntree), bNode *node
   }
 }
 
-static const blender::fn::MultiFunction &get_multi_function(
-    blender::nodes::NodeMFNetworkBuilder &builder)
+static const blender::fn::MultiFunction *get_multi_function(bNode &node)
 {
   using blender::float3;
 
-  NodeVectorMathOperation operation = NodeVectorMathOperation(builder.bnode().custom1);
+  NodeVectorMathOperation operation = NodeVectorMathOperation(node.custom1);
 
   const blender::fn::MultiFunction *multi_fn = nullptr;
 
@@ -199,7 +198,7 @@ static const blender::fn::MultiFunction &get_multi_function(
         multi_fn = &fn;
       });
   if (multi_fn != nullptr) {
-    return *multi_fn;
+    return multi_fn;
   }
 
   blender::nodes::try_dispatch_float_math_fl3_fl3_fl3_to_fl3(
@@ -209,7 +208,7 @@ static const blender::fn::MultiFunction &get_multi_function(
         multi_fn = &fn;
       });
   if (multi_fn != nullptr) {
-    return *multi_fn;
+    return multi_fn;
   }
 
   blender::nodes::try_dispatch_float_math_fl3_fl3_fl_to_fl3(
@@ -219,7 +218,7 @@ static const blender::fn::MultiFunction &get_multi_function(
         multi_fn = &fn;
       });
   if (multi_fn != nullptr) {
-    return *multi_fn;
+    return multi_fn;
   }
 
   blender::nodes::try_dispatch_float_math_fl3_fl3_to_fl(
@@ -229,7 +228,7 @@ static const blender::fn::MultiFunction &get_multi_function(
         multi_fn = &fn;
       });
   if (multi_fn != nullptr) {
-    return *multi_fn;
+    return multi_fn;
   }
 
   blender::nodes::try_dispatch_float_math_fl3_fl_to_fl3(
@@ -239,7 +238,7 @@ static const blender::fn::MultiFunction &get_multi_function(
         multi_fn = &fn;
       });
   if (multi_fn != nullptr) {
-    return *multi_fn;
+    return multi_fn;
   }
 
   blender::nodes::try_dispatch_float_math_fl3_to_fl3(
@@ -248,7 +247,7 @@ static const blender::fn::MultiFunction &get_multi_function(
         multi_fn = &fn;
       });
   if (multi_fn != nullptr) {
-    return *multi_fn;
+    return multi_fn;
   }
 
   blender::nodes::try_dispatch_float_math_fl3_to_fl(
@@ -257,15 +256,16 @@ static const blender::fn::MultiFunction &get_multi_function(
         multi_fn = &fn;
       });
   if (multi_fn != nullptr) {
-    return *multi_fn;
+    return multi_fn;
   }
 
-  return builder.get_not_implemented_fn();
+  return nullptr;
 }
 
-static void sh_node_vector_math_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_vector_math_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  const blender::fn::MultiFunction &fn = get_multi_function(builder);
+  const blender::fn::MultiFunction *fn = get_multi_function(builder.node());
   builder.set_matching_fn(fn);
 }
 
@@ -278,7 +278,7 @@ void register_node_type_sh_vect_math(void)
   node_type_label(&ntype, node_vector_math_label);
   node_type_gpu(&ntype, gpu_shader_vector_math);
   node_type_update(&ntype, node_shader_update_vector_math);
-  ntype.expand_in_mf_network = sh_node_vector_math_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_vector_math_build_multi_function;
 
   nodeRegisterType(&ntype);
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_vector_rotate.cc b/source/blender/nodes/shader/nodes/node_shader_vector_rotate.cc
index 3b2c2fa5a03..bc51b7e29ea 100644
--- a/source/blender/nodes/shader/nodes/node_shader_vector_rotate.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_vector_rotate.cc
@@ -100,11 +100,10 @@ static float3 sh_node_vector_rotate_euler(const float3 vector,
   return result + center;
 }
 
-static const blender::fn::MultiFunction &get_multi_function(
-    blender::nodes::NodeMFNetworkBuilder &builder)
+static const blender::fn::MultiFunction *get_multi_function(bNode &node)
 {
-  bool invert = builder.bnode().custom2;
-  const int mode = builder.bnode().custom1;
+  bool invert = node.custom2;
+  const int mode = node.custom1;
 
   switch (mode) {
     case NODE_VECTOR_ROTATE_TYPE_AXIS: {
@@ -113,13 +112,13 @@ static const blender::fn::MultiFunction &get_multi_function(
             "Rotate Axis", [](float3 in, float3 center, float3 axis, float angle) {
               return sh_node_vector_rotate_around_axis(in, center, axis, -angle);
             }};
-        return fn;
+        return &fn;
       }
       static blender::fn::CustomMF_SI_SI_SI_SI_SO<float3, float3, float3, float, float3> fn{
           "Rotate Axis", [](float3 in, float3 center, float3 axis, float angle) {
             return sh_node_vector_rotate_around_axis(in, center, axis, angle);
           }};
-      return fn;
+      return &fn;
     }
     case NODE_VECTOR_ROTATE_TYPE_AXIS_X: {
       float3 axis = float3(1.0f, 0.0f, 0.0f);
@@ -128,13 +127,13 @@ static const blender::fn::MultiFunction &get_multi_function(
             "Rotate X-Axis", [=](float3 in, float3 center, float angle) {
               return sh_node_vector_rotate_around_axis(in, center, axis, -angle);
             }};
-        return fn;
+        return &fn;
       }
       static blender::fn::CustomMF_SI_SI_SI_SO<float3, float3, float, float3> fn{
           "Rotate X-Axis", [=](float3 in, float3 center, float angle) {
             return sh_node_vector_rotate_around_axis(in, center, axis, angle);
           }};
-      return fn;
+      return &fn;
     }
     case NODE_VECTOR_ROTATE_TYPE_AXIS_Y: {
       float3 axis = float3(0.0f, 1.0f, 0.0f);
@@ -143,13 +142,13 @@ static const blender::fn::MultiFunction &get_multi_function(
             "Rotate Y-Axis", [=](float3 in, float3 center, float angle) {
               return sh_node_vector_rotate_around_axis(in, center, axis, -angle);
             }};
-        return fn;
+        return &fn;
       }
       static blender::fn::CustomMF_SI_SI_SI_SO<float3, float3, float, float3> fn{
           "Rotate Y-Axis", [=](float3 in, float3 center, float angle) {
             return sh_node_vector_rotate_around_axis(in, center, axis, angle);
           }};
-      return fn;
+      return &fn;
     }
     case NODE_VECTOR_ROTATE_TYPE_AXIS_Z: {
       float3 axis = float3(0.0f, 0.0f, 1.0f);
@@ -158,13 +157,13 @@ static const blender::fn::MultiFunction &get_multi_function(
             "Rotate Z-Axis", [=](float3 in, float3 center, float angle) {
               return sh_node_vector_rotate_around_axis(in, center, axis, -angle);
             }};
-        return fn;
+        return &fn;
       }
       static blender::fn::CustomMF_SI_SI_SI_SO<float3, float3, float, float3> fn{
           "Rotate Z-Axis", [=](float3 in, float3 center, float angle) {
             return sh_node_vector_rotate_around_axis(in, center, axis, angle);
           }};
-      return fn;
+      return &fn;
     }
     case NODE_VECTOR_ROTATE_TYPE_EULER_XYZ: {
       if (invert) {
@@ -172,24 +171,24 @@ static const blender::fn::MultiFunction &get_multi_function(
             "Rotate Euler", [](float3 in, float3 center, float3 rotation) {
               return sh_node_vector_rotate_euler(in, center, rotation, true);
             }};
-        return fn;
+        return &fn;
       }
       static blender::fn::CustomMF_SI_SI_SI_SO<float3, float3, float3, float3> fn{
           "Rotate Euler", [](float3 in, float3 center, float3 rotation) {
             return sh_node_vector_rotate_euler(in, center, rotation, false);
           }};
-      return fn;
+      return &fn;
     }
     default:
       BLI_assert_unreachable();
-      return builder.get_not_implemented_fn();
+      return nullptr;
   }
 }
 
-static void sh_node_vector_rotate_expand_in_mf_network(
-    blender::nodes::NodeMFNetworkBuilder &builder)
+static void sh_node_vector_rotate_build_multi_function(
+    blender::nodes::NodeMultiFunctionBuilder &builder)
 {
-  const blender::fn::MultiFunction &fn = get_multi_function(builder);
+  const blender::fn::MultiFunction *fn = get_multi_function(builder.node());
   builder.set_matching_fn(fn);
 }
 
@@ -211,7 +210,7 @@ void register_node_type_sh_vector_rotate(void)
   node_type_socket_templates(&ntype, sh_node_vector_rotate_in, sh_node_vector_rotate_out);
   node_type_gpu(&ntype, gpu_shader_vector_rotate);
   node_type_update(&ntype, node_shader_update_vector_rotate);
-  ntype.expand_in_mf_network = sh_node_vector_rotate_expand_in_mf_network;
+  ntype.build_multi_function = sh_node_vector_rotate_build_multi_function;
 
   nodeRegisterType(&ntype);
 }