Nodes: refactor derived node tree

This is a complete rewrite of the derived node tree data structure.
It is a much thinner abstraction about `NodeTreeRef` than before.
This gives the user of the derived node tree more control and allows
for greater introspection capabilities (e.g. before muted nodes were
completely abstracted away; this was convenient, but came with
limitations).

Another nice benefit of the new structure is that it is much cheaper
to build, because it does not inline all nodes and sockets in nested
node groups.

Differential Revision: https://developer.blender.org/D10620

11 files changed, 703 insertions, 1169 deletions

diff --git a/source/blender/modifiers/intern/MOD_nodes.cc b/source/blender/modifiers/intern/MOD_nodes.cc
index b47f5806c9c..4d1823b8951 100644
--- a/source/blender/modifiers/intern/MOD_nodes.cc
+++ b/source/blender/modifiers/intern/MOD_nodes.cc
@@ -77,6 +77,7 @@
 #include "NOD_type_callbacks.hh"
 
 using blender::float3;
+using blender::FunctionRef;
 using blender::IndexRange;
 using blender::Map;
 using blender::Set;
@@ -90,8 +91,8 @@ using blender::bke::PersistentObjectHandle;
 using blender::fn::GMutablePointer;
 using blender::fn::GValueMap;
 using blender::nodes::GeoNodeExecParams;
-using namespace blender::nodes::derived_node_tree_types;
 using namespace blender::fn::multi_function_types;
+using namespace blender::nodes::derived_node_tree_types;
 
 static void initData(ModifierData *md)
 {
@@ -254,8 +255,8 @@ static bool isDisabled(const struct Scene *UNUSED(scene),
 class GeometryNodesEvaluator {
  private:
   blender::LinearAllocator<> allocator_;
-  Map<std::pair<const DInputSocket *, const DOutputSocket *>, GMutablePointer> value_by_input_;
-  Vector<const DInputSocket *> group_outputs_;
+  Map<std::pair<DInputSocket, DOutputSocket>, GMutablePointer> value_by_input_;
+  Vector<DInputSocket> group_outputs_;
   blender::nodes::MultiFunctionByNode &mf_by_node_;
   const blender::nodes::DataTypeConversions &conversions_;
   const PersistentDataHandleMap &handle_map_;
@@ -264,8 +265,8 @@ class GeometryNodesEvaluator {
   Depsgraph *depsgraph_;
 
  public:
-  GeometryNodesEvaluator(const Map<const DOutputSocket *, GMutablePointer> &group_input_data,
-                         Vector<const DInputSocket *> group_outputs,
+  GeometryNodesEvaluator(const Map<DOutputSocket, GMutablePointer> &group_input_data,
+                         Vector<DInputSocket> group_outputs,
                          blender::nodes::MultiFunctionByNode &mf_by_node,
                          const PersistentDataHandleMap &handle_map,
                          const Object *self_object,
@@ -280,15 +281,15 @@ class GeometryNodesEvaluator {
         depsgraph_(depsgraph)
   {
     for (auto item : group_input_data.items()) {
-      this->forward_to_inputs(*item.key, item.value);
+      this->forward_to_inputs(item.key, item.value);
     }
   }
 
   Vector<GMutablePointer> execute()
   {
     Vector<GMutablePointer> results;
-    for (const DInputSocket *group_output : group_outputs_) {
-      Vector<GMutablePointer> result = this->get_input_values(*group_output);
+    for (const DInputSocket &group_output : group_outputs_) {
+      Vector<GMutablePointer> result = this->get_input_values(group_output);
       results.append(result[0]);
     }
     for (GMutablePointer value : value_by_input_.values()) {
@@ -298,62 +299,63 @@ class GeometryNodesEvaluator {
   }
 
  private:
-  Vector<GMutablePointer> get_input_values(const DInputSocket &socket_to_compute)
+  Vector<GMutablePointer> get_input_values(const DInputSocket socket_to_compute)
   {
+    Vector<DSocket> from_sockets;
+    socket_to_compute.foreach_origin_socket([&](DSocket socket) { from_sockets.append(socket); });
 
-    Span<const DOutputSocket *> from_sockets = socket_to_compute.linked_sockets();
-    Span<const DGroupInput *> from_group_inputs = socket_to_compute.linked_group_inputs();
-    const int total_inputs = from_sockets.size() + from_group_inputs.size();
+    /* Multi-input sockets contain a vector of inputs. */
+    if (socket_to_compute->is_multi_input_socket()) {
+      Vector<GMutablePointer> values;
+      for (const DSocket from_socket : from_sockets) {
+        GMutablePointer value = get_input_from_incoming_link(socket_to_compute, from_socket);
+        values.append(value);
+      }
+      return values;
+    }
 
-    if (total_inputs == 0) {
+    if (from_sockets.is_empty()) {
       /* The input is not connected, use the value from the socket itself. */
-      return {get_unlinked_input_value(socket_to_compute)};
+      const CPPType &type = *blender::nodes::socket_cpp_type_get(*socket_to_compute->typeinfo());
+      return {get_unlinked_input_value(socket_to_compute, type)};
     }
 
-    if (from_group_inputs.size() == 1) {
-      return {get_unlinked_input_value(socket_to_compute)};
-    }
+    const DSocket from_socket = from_sockets[0];
+    GMutablePointer value = this->get_input_from_incoming_link(socket_to_compute, from_socket);
+    return {value};
+  }
 
-    /* Multi-input sockets contain a vector of inputs. */
-    if (socket_to_compute.is_multi_input_socket()) {
-      Vector<GMutablePointer> values;
-      for (const DOutputSocket *from_socket : from_sockets) {
-        const std::pair<const DInputSocket *, const DOutputSocket *> key = std::make_pair(
-            &socket_to_compute, from_socket);
-        std::optional<GMutablePointer> value = value_by_input_.pop_try(key);
-        if (value.has_value()) {
-          values.append(*value);
-        }
-        else {
-          this->compute_output_and_forward(*from_socket);
-          GMutablePointer value = value_by_input_.pop(key);
-          values.append(value);
-        }
+  GMutablePointer get_input_from_incoming_link(const DInputSocket socket_to_compute,
+                                               const DSocket from_socket)
+  {
+    if (from_socket->is_output()) {
+      const DOutputSocket from_output_socket{from_socket};
+      const std::pair<DInputSocket, DOutputSocket> key = std::make_pair(socket_to_compute,
+                                                                        from_output_socket);
+      std::optional<GMutablePointer> value = value_by_input_.pop_try(key);
+      if (value.has_value()) {
+        /* This input has been computed before, return it directly. */
+        return {*value};
       }
-      return values;
-    }
 
-    const DOutputSocket &from_socket = *from_sockets[0];
-    const std::pair<const DInputSocket *, const DOutputSocket *> key = std::make_pair(
-        &socket_to_compute, &from_socket);
-    std::optional<GMutablePointer> value = value_by_input_.pop_try(key);
-    if (value.has_value()) {
-      /* This input has been computed before, return it directly. */
-      return {*value};
+      /* Compute the socket now. */
+      this->compute_output_and_forward(from_output_socket);
+      return {value_by_input_.pop(key)};
     }
 
-    /* Compute the socket now. */
-    this->compute_output_and_forward(from_socket);
-    return {value_by_input_.pop(key)};
+    /* Get value from an unlinked input socket. */
+    const CPPType &type = *blender::nodes::socket_cpp_type_get(*socket_to_compute->typeinfo());
+    const DInputSocket from_input_socket{from_socket};
+    return {get_unlinked_input_value(from_input_socket, type)};
   }
 
-  void compute_output_and_forward(const DOutputSocket &socket_to_compute)
+  void compute_output_and_forward(const DOutputSocket socket_to_compute)
   {
-    const DNode &node = socket_to_compute.node();
+    const DNode node{socket_to_compute.context(), &socket_to_compute->node()};
 
-    if (!socket_to_compute.is_available()) {
+    if (!socket_to_compute->is_available()) {
       /* If the output is not available, use a default value. */
-      const CPPType &type = *blender::nodes::socket_cpp_type_get(*socket_to_compute.typeinfo());
+      const CPPType &type = *blender::nodes::socket_cpp_type_get(*socket_to_compute->typeinfo());
       void *buffer = allocator_.allocate(type.size(), type.alignment());
       type.copy_to_uninitialized(type.default_value(), buffer);
       this->forward_to_inputs(socket_to_compute, {type, buffer});
@@ -362,9 +364,9 @@ class GeometryNodesEvaluator {
 
     /* Prepare inputs required to execute the node. */
     GValueMap<StringRef> node_inputs_map{allocator_};
-    for (const DInputSocket *input_socket : node.inputs()) {
+    for (const InputSocketRef *input_socket : node->inputs()) {
       if (input_socket->is_available()) {
-        Vector<GMutablePointer> values = this->get_input_values(*input_socket);
+        Vector<GMutablePointer> values = this->get_input_values({node.context(), input_socket});
         for (int i = 0; i < values.size(); ++i) {
           /* Values from Multi Input Sockets are stored in input map with the format
            * <identifier>[<index>]. */
@@ -382,15 +384,15 @@ class GeometryNodesEvaluator {
     this->execute_node(node, params);
 
     /* Forward computed outputs to linked input sockets. */
-    for (const DOutputSocket *output_socket : node.outputs()) {
+    for (const OutputSocketRef *output_socket : node->outputs()) {
       if (output_socket->is_available()) {
         GMutablePointer value = node_outputs_map.extract(output_socket->identifier());
-        this->forward_to_inputs(*output_socket, value);
+        this->forward_to_inputs({node.context(), output_socket}, value);
       }
     }
   }
 
-  void execute_node(const DNode &node, GeoNodeExecParams params)
+  void execute_node(const DNode node, GeoNodeExecParams params)
   {
     const bNode &bnode = params.node();
 
@@ -403,7 +405,7 @@ class GeometryNodesEvaluator {
     }
 
     /* Use the multi-function implementation if it exists. */
-    const MultiFunction *multi_function = mf_by_node_.lookup_default(&node, nullptr);
+    const MultiFunction *multi_function = mf_by_node_.lookup_default(node, nullptr);
     if (multi_function != nullptr) {
       this->execute_multi_function_node(node, params, *multi_function);
       return;
@@ -413,51 +415,52 @@ class GeometryNodesEvaluator {
     this->execute_unknown_node(node, params);
   }
 
-  void store_ui_hints(const DNode &node, GeoNodeExecParams params) const
+  void store_ui_hints(const DNode node, GeoNodeExecParams params) const
   {
-    for (const DInputSocket *dsocket : node.inputs()) {
-      if (!dsocket->is_available()) {
+    for (const InputSocketRef *socket_ref : node->inputs()) {
+      if (!socket_ref->is_available()) {
         continue;
       }
-      if (dsocket->bsocket()->type != SOCK_GEOMETRY) {
+      if (socket_ref->bsocket()->type != SOCK_GEOMETRY) {
         continue;
       }
 
-      bNodeTree *btree_cow = node.node_ref().tree().btree();
+      bNodeTree *btree_cow = node->btree();
       bNodeTree *btree_original = (bNodeTree *)DEG_get_original_id((ID *)btree_cow);
       const NodeTreeEvaluationContext context(*self_object_, *modifier_);
 
-      const GeometrySet &geometry_set = params.get_input<GeometrySet>(dsocket->identifier());
+      const GeometrySet &geometry_set = params.get_input<GeometrySet>(socket_ref->identifier());
       const Vector<const GeometryComponent *> components = geometry_set.get_components_for_read();
 
       for (const GeometryComponent *component : components) {
-        component->attribute_foreach([&](StringRefNull attribute_name,
-                                         const AttributeMetaData &UNUSED(meta_data)) {
-          BKE_nodetree_attribute_hint_add(*btree_original, context, *node.bnode(), attribute_name);
-          return true;
-        });
+        component->attribute_foreach(
+            [&](StringRefNull attribute_name, const AttributeMetaData &UNUSED(meta_data)) {
+              BKE_nodetree_attribute_hint_add(
+                  *btree_original, context, *node->bnode(), attribute_name);
+              return true;
+            });
       }
     }
   }
 
-  void execute_multi_function_node(const DNode &node,
+  void execute_multi_function_node(const DNode node,
                                    GeoNodeExecParams params,
                                    const MultiFunction &fn)
   {
     MFContextBuilder fn_context;
     MFParamsBuilder fn_params{fn, 1};
     Vector<GMutablePointer> input_data;
-    for (const DInputSocket *dsocket : node.inputs()) {
-      if (dsocket->is_available()) {
-        GMutablePointer data = params.extract_input(dsocket->identifier());
+    for (const InputSocketRef *socket_ref : node->inputs()) {
+      if (socket_ref->is_available()) {
+        GMutablePointer data = params.extract_input(socket_ref->identifier());
         fn_params.add_readonly_single_input(GSpan(*data.type(), data.get(), 1));
         input_data.append(data);
       }
     }
     Vector<GMutablePointer> output_data;
-    for (const DOutputSocket *dsocket : node.outputs()) {
-      if (dsocket->is_available()) {
-        const CPPType &type = *blender::nodes::socket_cpp_type_get(*dsocket->typeinfo());
+    for (const OutputSocketRef *socket_ref : node->outputs()) {
+      if (socket_ref->is_available()) {
+        const CPPType &type = *blender::nodes::socket_cpp_type_get(*socket_ref->typeinfo());
         void *buffer = allocator_.allocate(type.size(), type.alignment());
         fn_params.add_uninitialized_single_output(GMutableSpan(type, buffer, 1));
         output_data.append(GMutablePointer(type, buffer));
@@ -468,19 +471,19 @@ class GeometryNodesEvaluator {
       value.destruct();
     }
     int output_index = 0;
-    for (const int i : node.outputs().index_range()) {
-      if (node.output(i).is_available()) {
+    for (const int i : node->outputs().index_range()) {
+      if (node->output(i).is_available()) {
         GMutablePointer value = output_data[output_index];
-        params.set_output_by_move(node.output(i).identifier(), value);
+        params.set_output_by_move(node->output(i).identifier(), value);
         value.destruct();
         output_index++;
       }
     }
   }
 
-  void execute_unknown_node(const DNode &node, GeoNodeExecParams params)
+  void execute_unknown_node(const DNode node, GeoNodeExecParams params)
   {
-    for (const DOutputSocket *socket : node.outputs()) {
+    for (const OutputSocketRef *socket : node->outputs()) {
       if (socket->is_available()) {
         const CPPType &type = *blender::nodes::socket_cpp_type_get(*socket->typeinfo());
         params.set_output_by_copy(socket->identifier(), {type, type.default_value()});
@@ -488,17 +491,18 @@ class GeometryNodesEvaluator {
     }
   }
 
-  void forward_to_inputs(const DOutputSocket &from_socket, GMutablePointer value_to_forward)
+  void forward_to_inputs(const DOutputSocket from_socket, GMutablePointer value_to_forward)
   {
     /* For all sockets that are linked with the from_socket push the value to their node. */
-    Span<const DInputSocket *> to_sockets_all = from_socket.linked_sockets();
+    Vector<DInputSocket> to_sockets_all;
+    from_socket.foreach_target_socket(
+        [&](DInputSocket to_socket) { to_sockets_all.append(to_socket); });
 
     const CPPType &from_type = *value_to_forward.type();
-    Vector<const DInputSocket *> to_sockets_same_type;
-    for (const DInputSocket *to_socket : to_sockets_all) {
+    Vector<DInputSocket> to_sockets_same_type;
+    for (const DInputSocket &to_socket : to_sockets_all) {
       const CPPType &to_type = *blender::nodes::socket_cpp_type_get(*to_socket->typeinfo());
-      const std::pair<const DInputSocket *, const DOutputSocket *> key = std::make_pair(
-          to_socket, &from_socket);
+      const std::pair<DInputSocket, DOutputSocket> key = std::make_pair(to_socket, from_socket);
       if (from_type == to_type) {
         to_sockets_same_type.append(to_socket);
       }
@@ -520,23 +524,21 @@ class GeometryNodesEvaluator {
     }
     else if (to_sockets_same_type.size() == 1) {
       /* This value is only used on one input socket, no need to copy it. */
-      const DInputSocket *to_socket = to_sockets_same_type[0];
-      const std::pair<const DInputSocket *, const DOutputSocket *> key = std::make_pair(
-          to_socket, &from_socket);
+      const DInputSocket to_socket = to_sockets_same_type[0];
+      const std::pair<DInputSocket, DOutputSocket> key = std::make_pair(to_socket, from_socket);
 
       add_value_to_input_socket(key, value_to_forward);
     }
     else {
       /* Multiple inputs use the value, make a copy for every input except for one. */
-      const DInputSocket *first_to_socket = to_sockets_same_type[0];
-      Span<const DInputSocket *> other_to_sockets = to_sockets_same_type.as_span().drop_front(1);
+      const DInputSocket first_to_socket = to_sockets_same_type[0];
+      Span<DInputSocket> other_to_sockets = to_sockets_same_type.as_span().drop_front(1);
       const CPPType &type = *value_to_forward.type();
-      const std::pair<const DInputSocket *, const DOutputSocket *> first_key = std::make_pair(
-          first_to_socket, &from_socket);
+      const std::pair<DInputSocket, DOutputSocket> first_key = std::make_pair(first_to_socket,
+                                                                              from_socket);
       add_value_to_input_socket(first_key, value_to_forward);
-      for (const DInputSocket *to_socket : other_to_sockets) {
-        const std::pair<const DInputSocket *, const DOutputSocket *> key = std::make_pair(
-            to_socket, &from_socket);
+      for (const DInputSocket &to_socket : other_to_sockets) {
+        const std::pair<DInputSocket, DOutputSocket> key = std::make_pair(to_socket, from_socket);
         void *buffer = allocator_.allocate(type.size(), type.alignment());
         type.copy_to_uninitialized(value_to_forward.get(), buffer);
         add_value_to_input_socket(key, GMutablePointer{type, buffer});
@@ -544,22 +546,17 @@ class GeometryNodesEvaluator {
     }
   }
 
-  void add_value_to_input_socket(const std::pair<const DInputSocket *, const DOutputSocket *> key,
+  void add_value_to_input_socket(const std::pair<DInputSocket, DOutputSocket> key,
                                  GMutablePointer value)
   {
     value_by_input_.add_new(key, value);
   }
 
-  GMutablePointer get_unlinked_input_value(const DInputSocket &socket)
+  GMutablePointer get_unlinked_input_value(const DInputSocket &socket,
+                                           const CPPType &required_type)
   {
-    bNodeSocket *bsocket;
-    if (socket.linked_group_inputs().size() == 0) {
-      bsocket = socket.bsocket();
-    }
-    else {
-      bsocket = socket.linked_group_inputs()[0]->bsocket();
-    }
-    const CPPType &type = *blender::nodes::socket_cpp_type_get(*socket.typeinfo());
+    bNodeSocket *bsocket = socket->bsocket();
+    const CPPType &type = *blender::nodes::socket_cpp_type_get(*socket->typeinfo());
     void *buffer = allocator_.allocate(type.size(), type.alignment());
 
     if (bsocket->type == SOCK_OBJECT) {
@@ -576,7 +573,19 @@ class GeometryNodesEvaluator {
       blender::nodes::socket_cpp_value_get(*bsocket, buffer);
     }
 
-    return {type, buffer};
+    if (type == required_type) {
+      return {type, buffer};
+    }
+    if (conversions_.is_convertible(type, required_type)) {
+      void *converted_buffer = allocator_.allocate(required_type.size(),
+                                                   required_type.alignment());
+      conversions_.convert(type, required_type, buffer, converted_buffer);
+      type.destruct(buffer);
+      return {required_type, converted_buffer};
+    }
+    void *default_buffer = allocator_.allocate(required_type.size(), required_type.alignment());
+    type.copy_to_uninitialized(type.default_value(), default_buffer);
+    return {required_type, default_buffer};
   }
 };
 
@@ -985,7 +994,7 @@ static void fill_data_handle_map(const NodesModifierSettings &settings,
 {
   Set<ID *> used_ids;
   find_used_ids_from_settings(settings, used_ids);
-  find_used_ids_from_nodes(*tree.btree(), used_ids);
+  find_used_ids_from_nodes(*tree.root_context().tree().btree(), used_ids);
 
   int current_handle = 0;
   for (ID *id : used_ids) {
@@ -1013,8 +1022,8 @@ static void reset_tree_ui_storage(Span<const blender::nodes::NodeTreeRef *> tree
  * often than necessary. It's going to be replaced soon.
  */
 static GeometrySet compute_geometry(const DerivedNodeTree &tree,
-                                    Span<const DOutputSocket *> group_input_sockets,
-                                    const DInputSocket &socket_to_compute,
+                                    Span<const OutputSocketRef *> group_input_sockets,
+                                    const InputSocketRef &socket_to_compute,
                                     GeometrySet input_geometry_set,
                                     NodesModifierData *nmd,
                                     const ModifierEvalContext *ctx)
@@ -1026,32 +1035,33 @@ static GeometrySet compute_geometry(const DerivedNodeTree &tree,
   PersistentDataHandleMap handle_map;
   fill_data_handle_map(nmd->settings, tree, handle_map);
 
-  Map<const DOutputSocket *, GMutablePointer> group_inputs;
+  Map<DOutputSocket, GMutablePointer> group_inputs;
 
+  const DTreeContext *root_context = &tree.root_context();
   if (group_input_sockets.size() > 0) {
-    Span<const DOutputSocket *> remaining_input_sockets = group_input_sockets;
+    Span<const OutputSocketRef *> remaining_input_sockets = group_input_sockets;
 
     /* If the group expects a geometry as first input, use the geometry that has been passed to
      * modifier. */
-    const DOutputSocket *first_input_socket = group_input_sockets[0];
+    const OutputSocketRef *first_input_socket = group_input_sockets[0];
     if (first_input_socket->bsocket()->type == SOCK_GEOMETRY) {
       GeometrySet *geometry_set_in = allocator.construct<GeometrySet>(
           std::move(input_geometry_set));
-      group_inputs.add_new(first_input_socket, geometry_set_in);
+      group_inputs.add_new({root_context, first_input_socket}, geometry_set_in);
       remaining_input_sockets = remaining_input_sockets.drop_front(1);
     }
 
     /* Initialize remaining group inputs. */
-    for (const DOutputSocket *socket : remaining_input_sockets) {
+    for (const OutputSocketRef *socket : remaining_input_sockets) {
       const CPPType &cpp_type = *blender::nodes::socket_cpp_type_get(*socket->typeinfo());
       void *value_in = allocator.allocate(cpp_type.size(), cpp_type.alignment());
       initialize_group_input(*nmd, handle_map, *socket->bsocket(), cpp_type, value_in);
-      group_inputs.add_new(socket, {cpp_type, value_in});
+      group_inputs.add_new({root_context, socket}, {cpp_type, value_in});
     }
   }
 
-  Vector<const DInputSocket *> group_outputs;
-  group_outputs.append(&socket_to_compute);
+  Vector<DInputSocket> group_outputs;
+  group_outputs.append({root_context, &socket_to_compute});
 
   GeometryNodesEvaluator evaluator{group_inputs,
                                    group_outputs,
@@ -1126,16 +1136,17 @@ static void modifyGeometry(ModifierData *md,
 
   check_property_socket_sync(ctx->object, md);
 
-  blender::nodes::NodeTreeRefMap tree_refs;
-  DerivedNodeTree tree{nmd->node_group, tree_refs};
+  NodeTreeRefMap tree_refs;
+  DerivedNodeTree tree{*nmd->node_group, tree_refs};
 
   if (tree.has_link_cycles()) {
     BKE_modifier_set_error(ctx->object, md, "Node group has cycles");
     return;
   }
 
-  Span<const DNode *> input_nodes = tree.nodes_by_type("NodeGroupInput");
-  Span<const DNode *> output_nodes = tree.nodes_by_type("NodeGroupOutput");
+  const NodeTreeRef &root_tree_ref = tree.root_context().tree();
+  Span<const NodeRef *> input_nodes = root_tree_ref.nodes_by_type("NodeGroupInput");
+  Span<const NodeRef *> output_nodes = root_tree_ref.nodes_by_type("NodeGroupOutput");
 
   if (input_nodes.size() > 1) {
     return;
@@ -1144,16 +1155,18 @@ static void modifyGeometry(ModifierData *md,
     return;
   }
 
-  Span<const DOutputSocket *> group_inputs = (input_nodes.size() == 1) ?
-                                                 input_nodes[0]->outputs().drop_back(1) :
-                                                 Span<const DOutputSocket *>{};
-  Span<const DInputSocket *> group_outputs = output_nodes[0]->inputs().drop_back(1);
+  Span<const OutputSocketRef *> group_inputs;
+  if (input_nodes.size() == 1) {
+    group_inputs = input_nodes[0]->outputs().drop_back(1);
+  }
+
+  Span<const InputSocketRef *> group_outputs = output_nodes[0]->inputs().drop_back(1);
 
   if (group_outputs.size() == 0) {
     return;
   }
 
-  const DInputSocket *group_output = group_outputs[0];
+  const InputSocketRef *group_output = group_outputs[0];
   if (group_output->idname() != "NodeSocketGeometry") {
     return;
   }
diff --git a/source/blender/nodes/NOD_derived_node_tree.hh b/source/blender/nodes/NOD_derived_node_tree.hh
index ea67f23eade..a254f908bfd 100644
--- a/source/blender/nodes/NOD_derived_node_tree.hh
+++ b/source/blender/nodes/NOD_derived_node_tree.hh
@@ -19,544 +19,360 @@
 /** \file
  * \ingroup nodes
  *
- * DerivedNodeTree provides a flattened view on a bNodeTree, i.e. node groups are inlined. It
- * builds on top of NodeTreeRef and supports similar queries efficiently.
- *
- * Every inlined node remembers its path to the parent ("call stack").
- *
- * Unlinked group node inputs are handled separately from other sockets.
- *
- * There is a dot graph exporter for debugging purposes.
+ * DerivedNodeTree builds on top of NodeTreeRef and makes working with (nested) node groups more
+ * convenient and safe. It does so by pairing nodes and sockets with a context. The context
+ * contains information about the current "instance" of the node or socket. A node might be
+ * "instanced" multiple times when it is in a node group that is used multiple times.
  */
 
-#include "NOD_node_tree_ref.hh"
-
+#include "BLI_function_ref.hh"
 #include "BLI_vector_set.hh"
 
+#include "NOD_node_tree_ref.hh"
+
 namespace blender::nodes {
 
-class DSocket;
-class DInputSocket;
-class DOutputSocket;
-class DNode;
-class DParentNode;
-class DGroupInput;
+class DTreeContext;
 class DerivedNodeTree;
 
-class DSocket : NonCopyable, NonMovable {
- protected:
-  DNode *node_;
-  const SocketRef *socket_ref_;
-  int id_;
-
-  friend DerivedNodeTree;
-
- public:
-  const DNode &node() const;
-
-  int id() const;
-  int index() const;
-
-  bool is_input() const;
-  bool is_output() const;
-
-  const DSocket &as_base() const;
-  const DInputSocket &as_input() const;
-  const DOutputSocket &as_output() const;
-
-  PointerRNA *rna() const;
-  StringRefNull idname() const;
-  StringRefNull name() const;
-  StringRefNull identifier() const;
-  bNodeSocketType *typeinfo() const;
-
-  const SocketRef &socket_ref() const;
-  bNodeSocket *bsocket() const;
-
-  bool is_available() const;
-};
+struct DNode;
+struct DSocket;
+struct DInputSocket;
+struct DOutputSocket;
 
-class DInputSocket : public DSocket {
+/**
+ * The context attached to every node or socket in a derived node tree. It can be used to determine
+ * the place of a node in a hierarchy of node groups.
+ *
+ * Contexts are organized in a tree data structure to avoid having to store the entire path to the
+ * root node group for every node/socket.
+ */
+class DTreeContext {
  private:
-  Vector<DOutputSocket *> linked_sockets_;
-  Vector<DGroupInput *> linked_group_inputs_;
-  bool is_multi_input_socket_;
+  /* Null when this context is for the root node group. Otherwise it points to the context one
+   * level up. */
+  DTreeContext *parent_context_;
+  /* Null when this context is for the root node group. Otherwise it points to the group node in
+   * the parent node group that contains this context. */
+  const NodeRef *parent_node_;
+  /* The current node tree. */
+  const NodeTreeRef *tree_;
+  /* All the children contexts of this context. */
+  Map<const NodeRef *, DTreeContext *> children_;
 
   friend DerivedNodeTree;
 
  public:
-  const InputSocketRef &socket_ref() const;
-
-  Span<const DOutputSocket *> linked_sockets() const;
-  Span<const DGroupInput *> linked_group_inputs() const;
-
-  bool is_linked() const;
-  bool is_multi_input_socket() const;
+  const NodeTreeRef &tree() const;
+  const DTreeContext *parent_context() const;
+  const NodeRef *parent_node() const;
+  const DTreeContext *child_context(const NodeRef &node) const;
+  bool is_root() const;
 };
 
-class DOutputSocket : public DSocket {
+/* A (nullable) reference to a node and the context it is in. It is unique within an entire nested
+ * node group hierarchy. This type is small and can be passed around by value. */
+class DNode {
  private:
-  Vector<DInputSocket *> linked_sockets_;
-
-  friend DerivedNodeTree;
+  const DTreeContext *context_ = nullptr;
+  const NodeRef *node_ref_ = nullptr;
 
  public:
-  const OutputSocketRef &socket_ref() const;
-  Span<const DInputSocket *> linked_sockets() const;
-};
+  DNode() = default;
+  DNode(const DTreeContext *context, const NodeRef *node);
 
-class DGroupInput : NonCopyable, NonMovable {
- private:
-  const InputSocketRef *socket_ref_;
-  DParentNode *parent_;
-  Vector<DInputSocket *> linked_sockets_;
-  int id_;
+  const DTreeContext *context() const;
+  const NodeRef *node_ref() const;
+  const NodeRef *operator->() const;
 
-  friend DerivedNodeTree;
+  friend bool operator==(const DNode &a, const DNode &b);
+  friend bool operator!=(const DNode &a, const DNode &b);
+  operator bool() const;
 
- public:
-  const InputSocketRef &socket_ref() const;
-  bNodeSocket *bsocket() const;
-  const DParentNode *parent() const;
-  Span<const DInputSocket *> linked_sockets() const;
-  int id() const;
-  StringRefNull name() const;
+  uint64_t hash() const;
 };
 
-class DNode : NonCopyable, NonMovable {
- private:
-  const NodeRef *node_ref_;
-  DParentNode *parent_;
-
-  Span<DInputSocket *> inputs_;
-  Span<DOutputSocket *> outputs_;
-
-  int id_;
-
-  friend DerivedNodeTree;
+/* A (nullable) reference to a socket and the context it is in. It is unique within an entire
+ * nested node group hierarchy. This type is small and can be passed around by value.
+ *
+ * A #DSocket can represent an input or an output socket. If the type of a socket is known at
+ * compile time is is preferable to use #DInputSocket or #DOutputSocket instead. */
+class DSocket {
+ protected:
+  const DTreeContext *context_ = nullptr;
+  const SocketRef *socket_ref_ = nullptr;
 
  public:
-  const NodeRef &node_ref() const;
-  const DParentNode *parent() const;
+  DSocket() = default;
+  DSocket(const DTreeContext *context, const SocketRef *socket);
+  DSocket(const DInputSocket &input_socket);
+  DSocket(const DOutputSocket &output_socket);
 
-  Span<const DInputSocket *> inputs() const;
-  Span<const DOutputSocket *> outputs() const;
+  const DTreeContext *context() const;
+  const SocketRef *socket_ref() const;
+  const SocketRef *operator->() const;
 
-  const DInputSocket &input(int index) const;
-  const DOutputSocket &output(int index) const;
+  friend bool operator==(const DSocket &a, const DSocket &b);
+  friend bool operator!=(const DSocket &a, const DSocket &b);
+  operator bool() const;
 
-  const DInputSocket &input(int index, StringRef expected_name) const;
-  const DOutputSocket &output(int index, StringRef expected_name) const;
+  uint64_t hash() const;
+};
 
-  int id() const;
+/* A (nullable) reference to an input socket and the context it is in. */
+class DInputSocket : public DSocket {
+ public:
+  DInputSocket() = default;
+  DInputSocket(const DTreeContext *context, const InputSocketRef *socket);
+  explicit DInputSocket(const DSocket &base_socket);
 
-  PointerRNA *rna() const;
-  StringRefNull idname() const;
-  StringRefNull name() const;
-  bNode *bnode() const;
-  bNodeType *typeinfo() const;
+  const InputSocketRef *socket_ref() const;
+  const InputSocketRef *operator->() const;
 
- private:
-  void destruct_with_sockets();
+  DOutputSocket get_corresponding_group_node_output() const;
+  Vector<DOutputSocket, 4> get_corresponding_group_input_sockets() const;
+
+  void foreach_origin_socket(FunctionRef<void(DSocket)> callback) const;
 };
 
-class DParentNode : NonCopyable, NonMovable {
- private:
-  const NodeRef *node_ref_;
-  DParentNode *parent_;
-  int id_;
+/* A (nullable) reference to an output socket and the context it is in. */
+class DOutputSocket : public DSocket {
+ public:
+  DOutputSocket() = default;
+  DOutputSocket(const DTreeContext *context, const OutputSocketRef *socket);
+  explicit DOutputSocket(const DSocket &base_socket);
 
-  friend DerivedNodeTree;
+  const OutputSocketRef *socket_ref() const;
+  const OutputSocketRef *operator->() const;
 
- public:
-  const DParentNode *parent() const;
-  const NodeRef &node_ref() const;
-  int id() const;
+  DInputSocket get_corresponding_group_node_input() const;
+  DInputSocket get_active_corresponding_group_output_socket() const;
+
+  void foreach_target_socket(FunctionRef<void(DInputSocket)> callback) const;
 };
 
-class DerivedNodeTree : NonCopyable, NonMovable {
+class DerivedNodeTree {
  private:
   LinearAllocator<> allocator_;
-  Vector<DNode *> nodes_by_id_;
-  Vector<DGroupInput *> group_inputs_;
-  Vector<DParentNode *> parent_nodes_;
-
-  Vector<DSocket *> sockets_by_id_;
-  Vector<DInputSocket *> input_sockets_;
-  Vector<DOutputSocket *> output_sockets_;
-
-  MultiValueMap<const bNodeType *, DNode *> nodes_by_type_;
+  DTreeContext *root_context_;
   VectorSet<const NodeTreeRef *> used_node_tree_refs_;
-  bNodeTree *btree_;
 
  public:
-  DerivedNodeTree(bNodeTree *btree, NodeTreeRefMap &node_tree_refs);
+  DerivedNodeTree(bNodeTree &btree, NodeTreeRefMap &node_tree_refs);
   ~DerivedNodeTree();
 
-  bNodeTree *btree() const;
-
-  Span<const DNode *> nodes() const;
-  Span<const DNode *> nodes_by_type(StringRefNull idname) const;
-  Span<const DNode *> nodes_by_type(const bNodeType *nodetype) const;
-
-  Span<const DSocket *> sockets() const;
-  Span<const DInputSocket *> input_sockets() const;
-  Span<const DOutputSocket *> output_sockets() const;
-
-  Span<const DGroupInput *> group_inputs() const;
-
+  const DTreeContext &root_context() const;
   Span<const NodeTreeRef *> used_node_tree_refs() const;
 
   bool has_link_cycles() const;
-
-  std::string to_dot() const;
+  void foreach_node(FunctionRef<void(DNode)> callback) const;
 
  private:
-  /* Utility functions used during construction. */
-  void insert_nodes_and_links_in_id_order(const NodeTreeRef &tree_ref,
-                                          DParentNode *parent,
-                                          Vector<DNode *> &all_nodes);
-  DNode &create_node(const NodeRef &node_ref,
-                     DParentNode *parent,
-                     MutableSpan<DSocket *> r_sockets_map);
-  void expand_groups(Vector<DNode *> &all_nodes,
-                     Vector<DGroupInput *> &all_group_inputs,
-                     Vector<DParentNode *> &all_parent_nodes,
-                     NodeTreeRefMap &node_tree_refs);
-  void expand_group_node(DNode &group_node,
-                         Vector<DNode *> &all_nodes,
-                         Vector<DGroupInput *> &all_group_inputs,
-                         Vector<DParentNode *> &all_parent_nodes,
-                         NodeTreeRefMap &node_tree_refs);
-  void create_group_inputs_for_unlinked_inputs(DNode &node,
-                                               Vector<DGroupInput *> &all_group_inputs);
-  void relink_group_inputs(const NodeTreeRef &group_ref,
-                           Span<DNode *> nodes_by_id,
-                           DNode &group_node);
-  void relink_group_outputs(const NodeTreeRef &group_ref,
-                            Span<DNode *> nodes_by_id,
-                            DNode &group_node);
-  void remove_expanded_group_interfaces(Vector<DNode *> &all_nodes);
-  void remove_unused_group_inputs(Vector<DGroupInput *> &all_group_inputs);
-  void relink_and_remove_muted_nodes(Vector<DNode *> &all_nodes);
-  void relink_muted_node(DNode &muted_node);
-  void store_in_this_and_init_ids(Vector<DNode *> &&all_nodes,
-                                  Vector<DGroupInput *> &&all_group_inputs,
-                                  Vector<DParentNode *> &&all_parent_nodes);
+  DTreeContext &construct_context_recursively(DTreeContext *parent_context,
+                                              const NodeRef *parent_node,
+                                              bNodeTree &btree,
+                                              NodeTreeRefMap &node_tree_refs);
+  void destruct_context_recursively(DTreeContext *context);
+
+  void foreach_node_in_context_recursive(const DTreeContext &context,
+                                         FunctionRef<void(DNode)> callback) const;
 };
 
 namespace derived_node_tree_types {
+using namespace node_tree_ref_types;
 using nodes::DerivedNodeTree;
-using nodes::DGroupInput;
 using nodes::DInputSocket;
 using nodes::DNode;
 using nodes::DOutputSocket;
-using nodes::DParentNode;
-};  // namespace derived_node_tree_types
+using nodes::DSocket;
+using nodes::DTreeContext;
+}  // namespace derived_node_tree_types
 
 /* --------------------------------------------------------------------
- * DSocket inline methods.
+ * DTreeContext inline methods.
  */
 
-inline const DNode &DSocket::node() const
-{
-  return *node_;
-}
-
-inline int DSocket::id() const
-{
-  return id_;
-}
-
-inline int DSocket::index() const
+inline const NodeTreeRef &DTreeContext::tree() const
 {
-  return socket_ref_->index();
+  return *tree_;
 }
 
-inline bool DSocket::is_input() const
+inline const DTreeContext *DTreeContext::parent_context() const
 {
-  return socket_ref_->is_input();
+  return parent_context_;
 }
 
-inline bool DSocket::is_output() const
+inline const NodeRef *DTreeContext::parent_node() const
 {
-  return socket_ref_->is_output();
+  return parent_node_;
 }
 
-inline const DSocket &DSocket::as_base() const
+inline const DTreeContext *DTreeContext::child_context(const NodeRef &node) const
 {
-  return *this;
+  return children_.lookup_default(&node, nullptr);
 }
 
-inline const DInputSocket &DSocket::as_input() const
+inline bool DTreeContext::is_root() const
 {
-  return static_cast<const DInputSocket &>(*this);
+  return parent_context_ == nullptr;
 }
 
-inline const DOutputSocket &DSocket::as_output() const
-{
-  return static_cast<const DOutputSocket &>(*this);
-}
+/* --------------------------------------------------------------------
+ * DNode inline methods.
+ */
 
-inline PointerRNA *DSocket::rna() const
+inline DNode::DNode(const DTreeContext *context, const NodeRef *node_ref)
+    : context_(context), node_ref_(node_ref)
 {
-  return socket_ref_->rna();
+  BLI_assert(node_ref == nullptr || &node_ref->tree() == &context->tree());
 }
 
-inline StringRefNull DSocket::idname() const
+inline const DTreeContext *DNode::context() const
 {
-  return socket_ref_->idname();
+  return context_;
 }
 
-inline StringRefNull DSocket::name() const
+inline const NodeRef *DNode::node_ref() const
 {
-  return socket_ref_->name();
+  return node_ref_;
 }
 
-inline StringRefNull DSocket::identifier() const
+inline bool operator==(const DNode &a, const DNode &b)
 {
-  return socket_ref_->identifier();
+  return a.context_ == b.context_ && a.node_ref_ == b.node_ref_;
 }
 
-inline bNodeSocketType *DSocket::typeinfo() const
+inline bool operator!=(const DNode &a, const DNode &b)
 {
-  return socket_ref_->bsocket()->typeinfo;
+  return !(a == b);
 }
 
-inline const SocketRef &DSocket::socket_ref() const
+inline DNode::operator bool() const
 {
-  return *socket_ref_;
+  return node_ref_ != nullptr;
 }
 
-inline bNodeSocket *DSocket::bsocket() const
+inline const NodeRef *DNode::operator->() const
 {
-  return socket_ref_->bsocket();
+  return node_ref_;
 }
 
-inline bool DSocket::is_available() const
+inline uint64_t DNode::hash() const
 {
-  return (socket_ref_->bsocket()->flag & SOCK_UNAVAIL) == 0;
+  return DefaultHash<const DTreeContext *>{}(context_) ^ DefaultHash<const NodeRef *>{}(node_ref_);
 }
 
 /* --------------------------------------------------------------------
- * DInputSocket inline methods.
+ * DSocket inline methods.
  */
 
-inline const InputSocketRef &DInputSocket::socket_ref() const
+inline DSocket::DSocket(const DTreeContext *context, const SocketRef *socket_ref)
+    : context_(context), socket_ref_(socket_ref)
 {
-  return socket_ref_->as_input();
+  BLI_assert(socket_ref == nullptr || &socket_ref->tree() == &context->tree());
 }
 
-inline Span<const DOutputSocket *> DInputSocket::linked_sockets() const
+inline DSocket::DSocket(const DInputSocket &input_socket)
+    : DSocket(input_socket.context_, input_socket.socket_ref_)
 {
-  return linked_sockets_;
 }
 
-inline Span<const DGroupInput *> DInputSocket::linked_group_inputs() const
+inline DSocket::DSocket(const DOutputSocket &output_socket)
+    : DSocket(output_socket.context_, output_socket.socket_ref_)
 {
-  return linked_group_inputs_;
 }
 
-inline bool DInputSocket::is_linked() const
+inline const DTreeContext *DSocket::context() const
 {
-  return linked_sockets_.size() > 0 || linked_group_inputs_.size() > 0;
+  return context_;
 }
 
-inline bool DInputSocket::is_multi_input_socket() const
+inline const SocketRef *DSocket::socket_ref() const
 {
-  return is_multi_input_socket_;
+  return socket_ref_;
 }
 
-/* --------------------------------------------------------------------
- * DOutputSocket inline methods.
- */
-
-inline const OutputSocketRef &DOutputSocket::socket_ref() const
-{
-  return socket_ref_->as_output();
-}
-
-inline Span<const DInputSocket *> DOutputSocket::linked_sockets() const
-{
-  return linked_sockets_;
-}
-
-/* --------------------------------------------------------------------
- * DGroupInput inline methods.
- */
-
-inline const InputSocketRef &DGroupInput::socket_ref() const
+inline bool operator==(const DSocket &a, const DSocket &b)
 {
-  return *socket_ref_;
+  return a.context_ == b.context_ && a.socket_ref_ == b.socket_ref_;
 }
 
-inline bNodeSocket *DGroupInput::bsocket() const
+inline bool operator!=(const DSocket &a, const DSocket &b)
 {
-  return socket_ref_->bsocket();
+  return !(a == b);
 }
 
-inline const DParentNode *DGroupInput::parent() const
+inline DSocket::operator bool() const
 {
-  return parent_;
+  return socket_ref_ != nullptr;
 }
 
-inline Span<const DInputSocket *> DGroupInput::linked_sockets() const
+inline const SocketRef *DSocket::operator->() const
 {
-  return linked_sockets_;
+  return socket_ref_;
 }
 
-inline int DGroupInput::id() const
+inline uint64_t DSocket::hash() const
 {
-  return id_;
-}
-
-inline StringRefNull DGroupInput::name() const
-{
-  return socket_ref_->name();
+  return DefaultHash<const DTreeContext *>{}(context_) ^
+         DefaultHash<const SocketRef *>{}(socket_ref_);
 }
 
 /* --------------------------------------------------------------------
- * DNode inline methods.
+ * DInputSocket inline methods.
  */
 
-inline const NodeRef &DNode::node_ref() const
-{
-  return *node_ref_;
-}
-
-inline const DParentNode *DNode::parent() const
-{
-  return parent_;
-}
-
-inline Span<const DInputSocket *> DNode::inputs() const
-{
-  return inputs_;
-}
-
-inline Span<const DOutputSocket *> DNode::outputs() const
-{
-  return outputs_;
-}
-
-inline const DInputSocket &DNode::input(int index) const
-{
-  return *inputs_[index];
-}
-
-inline const DOutputSocket &DNode::output(int index) const
-{
-  return *outputs_[index];
-}
-
-inline const DInputSocket &DNode::input(int index, StringRef expected_name) const
+inline DInputSocket::DInputSocket(const DTreeContext *context, const InputSocketRef *socket_ref)
+    : DSocket(context, socket_ref)
 {
-  const DInputSocket &socket = *inputs_[index];
-  BLI_assert(socket.name() == expected_name);
-  UNUSED_VARS_NDEBUG(expected_name);
-  return socket;
 }
 
-inline const DOutputSocket &DNode::output(int index, StringRef expected_name) const
+inline DInputSocket::DInputSocket(const DSocket &base_socket) : DSocket(base_socket)
 {
-  const DOutputSocket &socket = *outputs_[index];
-  BLI_assert(socket.name() == expected_name);
-  UNUSED_VARS_NDEBUG(expected_name);
-  return socket;
+  BLI_assert(base_socket->is_input());
 }
 
-inline int DNode::id() const
+inline const InputSocketRef *DInputSocket::socket_ref() const
 {
-  return id_;
+  return (const InputSocketRef *)socket_ref_;
 }
 
-inline PointerRNA *DNode::rna() const
+inline const InputSocketRef *DInputSocket::operator->() const
 {
-  return node_ref_->rna();
+  return (const InputSocketRef *)socket_ref_;
 }
 
-inline StringRefNull DNode::idname() const
-{
-  return node_ref_->idname();
-}
-
-inline StringRefNull DNode::name() const
-{
-  return node_ref_->name();
-}
-
-inline bNode *DNode::bnode() const
-{
-  return node_ref_->bnode();
-}
+/* --------------------------------------------------------------------
+ * DOutputSocket inline methods.
+ */
 
-inline bNodeType *DNode::typeinfo() const
+inline DOutputSocket::DOutputSocket(const DTreeContext *context, const OutputSocketRef *socket_ref)
+    : DSocket(context, socket_ref)
 {
-  return node_ref_->bnode()->typeinfo;
 }
 
-/* --------------------------------------------------------------------
- * DParentNode inline methods.
- */
-
-inline const DParentNode *DParentNode::parent() const
+inline DOutputSocket::DOutputSocket(const DSocket &base_socket) : DSocket(base_socket)
 {
-  return parent_;
+  BLI_assert(base_socket->is_output());
 }
 
-inline const NodeRef &DParentNode::node_ref() const
+inline const OutputSocketRef *DOutputSocket::socket_ref() const
 {
-  return *node_ref_;
+  return (const OutputSocketRef *)socket_ref_;
 }
 
-inline int DParentNode::id() const
+inline const OutputSocketRef *DOutputSocket::operator->() const
 {
-  return id_;
+  return (const OutputSocketRef *)socket_ref_;
 }
 
 /* --------------------------------------------------------------------
  * DerivedNodeTree inline methods.
  */
 
-inline bNodeTree *DerivedNodeTree::btree() const
-{
-  return btree_;
-}
-
-inline Span<const DNode *> DerivedNodeTree::nodes() const
-{
-  return nodes_by_id_;
-}
-
-inline Span<const DNode *> DerivedNodeTree::nodes_by_type(StringRefNull idname) const
-{
-  const bNodeType *nodetype = nodeTypeFind(idname.c_str());
-  return this->nodes_by_type(nodetype);
-}
-
-inline Span<const DNode *> DerivedNodeTree::nodes_by_type(const bNodeType *nodetype) const
-{
-  return nodes_by_type_.lookup(nodetype);
-}
-
-inline Span<const DSocket *> DerivedNodeTree::sockets() const
-{
-  return sockets_by_id_;
-}
-
-inline Span<const DInputSocket *> DerivedNodeTree::input_sockets() const
-{
-  return input_sockets_;
-}
-
-inline Span<const DOutputSocket *> DerivedNodeTree::output_sockets() const
-{
-  return output_sockets_;
-}
-
-inline Span<const DGroupInput *> DerivedNodeTree::group_inputs() const
+inline const DTreeContext &DerivedNodeTree::root_context() const
 {
-  return group_inputs_;
+  return *root_context_;
 }
 
 inline Span<const NodeTreeRef *> DerivedNodeTree::used_node_tree_refs() const
diff --git a/source/blender/nodes/NOD_geometry_exec.hh b/source/blender/nodes/NOD_geometry_exec.hh
index e648d77337b..1772f92c4b6 100644
--- a/source/blender/nodes/NOD_geometry_exec.hh
+++ b/source/blender/nodes/NOD_geometry_exec.hh
@@ -59,7 +59,7 @@ using fn::GValueMap;
 
 class GeoNodeExecParams {
  private:
-  const DNode &node_;
+  const DNode node_;
   GValueMap<StringRef> &input_values_;
   GValueMap<StringRef> &output_values_;
   const PersistentDataHandleMap &handle_map_;
@@ -68,7 +68,7 @@ class GeoNodeExecParams {
   Depsgraph *depsgraph_;
 
  public:
-  GeoNodeExecParams(const DNode &node,
+  GeoNodeExecParams(const DNode node,
                     GValueMap<StringRef> &input_values,
                     GValueMap<StringRef> &output_values,
                     const PersistentDataHandleMap &handle_map,
@@ -182,7 +182,7 @@ class GeoNodeExecParams {
    */
   const bNode &node() const
   {
-    return *node_.bnode();
+    return *node_->bnode();
   }
 
   const PersistentDataHandleMap &handle_map() const
diff --git a/source/blender/nodes/NOD_node_tree_multi_function.hh b/source/blender/nodes/NOD_node_tree_multi_function.hh
index 552ef5509fa..df31ee18369 100644
--- a/source/blender/nodes/NOD_node_tree_multi_function.hh
+++ b/source/blender/nodes/NOD_node_tree_multi_function.hh
@@ -28,14 +28,15 @@
 #include "NOD_derived_node_tree.hh"
 #include "NOD_type_callbacks.hh"
 
+#include "BLI_multi_value_map.hh"
 #include "BLI_resource_collector.hh"
 
 namespace blender::nodes {
 
 /**
- * A MFNetworkTreeMap maps various components of a DerivedNodeTree 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.
+ * 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:
@@ -47,15 +48,11 @@ class MFNetworkTreeMap {
    */
   const DerivedNodeTree &tree_;
   fn::MFNetwork &network_;
-  Array<Vector<fn::MFSocket *, 1>> sockets_by_dsocket_id_;
-  Array<fn::MFOutputSocket *> socket_by_group_input_id_;
+  MultiValueMap<DSocket, fn::MFSocket *> sockets_by_dsocket_;
 
  public:
   MFNetworkTreeMap(const DerivedNodeTree &tree, fn::MFNetwork &network)
-      : tree_(tree),
-        network_(network),
-        sockets_by_dsocket_id_(tree.sockets().size()),
-        socket_by_group_input_id_(tree.group_inputs().size(), nullptr)
+      : tree_(tree), network_(network)
   {
   }
 
@@ -76,96 +73,95 @@ class MFNetworkTreeMap {
 
   void add(const DSocket &dsocket, fn::MFSocket &socket)
   {
-    BLI_assert(dsocket.is_input() == socket.is_input());
-    BLI_assert(dsocket.is_input() || sockets_by_dsocket_id_[dsocket.id()].size() == 0);
-    sockets_by_dsocket_id_[dsocket.id()].append(&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_id_[dsocket.id()].append(&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_id_[dsocket.id()].size() == 0);
-    sockets_by_dsocket_id_[dsocket.id()].append(&socket);
+    BLI_assert(sockets_by_dsocket_.lookup(dsocket).is_empty());
+    sockets_by_dsocket_.add(dsocket, &socket);
   }
 
-  void add(Span<const DInputSocket *> dsockets, Span<fn::MFInputSocket *> sockets)
+  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(*dsockets[i], *sockets[i]);
+      this->add(DInputSocket(&context, dsockets[i]), *sockets[i]);
     }
   }
 
-  void add(Span<const DOutputSocket *> dsockets, Span<fn::MFOutputSocket *> sockets)
+  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(*dsockets[i], *sockets[i]);
+      this->add(DOutputSocket(&context, dsockets[i]), *sockets[i]);
     }
   }
 
-  void add(const DGroupInput &group_input, fn::MFOutputSocket &socket)
-  {
-    BLI_assert(socket_by_group_input_id_[group_input.id()] == nullptr);
-    socket_by_group_input_id_[group_input.id()] = &socket;
-  }
-
   void add_try_match(const DNode &dnode, fn::MFNode &node)
   {
-    this->add_try_match(dnode.inputs().cast<const DSocket *>(),
+    this->add_try_match(*dnode.context(),
+                        dnode->inputs().cast<const SocketRef *>(),
                         node.inputs().cast<fn::MFSocket *>());
-    this->add_try_match(dnode.outputs().cast<const DSocket *>(),
+    this->add_try_match(*dnode.context(),
+                        dnode->outputs().cast<const SocketRef *>(),
                         node.outputs().cast<fn::MFSocket *>());
   }
 
-  void add_try_match(Span<const DInputSocket *> dsockets, Span<fn::MFInputSocket *> sockets)
+  void add_try_match(const DTreeContext &context,
+                     Span<const InputSocketRef *> dsockets,
+                     Span<fn::MFInputSocket *> sockets)
   {
-    this->add_try_match(dsockets.cast<const DSocket *>(), sockets.cast<fn::MFSocket *>());
+    this->add_try_match(
+        context, dsockets.cast<const SocketRef *>(), sockets.cast<fn::MFSocket *>());
   }
 
-  void add_try_match(Span<const DOutputSocket *> dsockets, Span<fn::MFOutputSocket *> sockets)
+  void add_try_match(const DTreeContext &context,
+                     Span<const OutputSocketRef *> dsockets,
+                     Span<fn::MFOutputSocket *> sockets)
   {
-    this->add_try_match(dsockets.cast<const DSocket *>(), sockets.cast<fn::MFSocket *>());
+    this->add_try_match(
+        context, dsockets.cast<const SocketRef *>(), sockets.cast<fn::MFSocket *>());
   }
 
-  void add_try_match(Span<const DSocket *> dsockets, Span<fn::MFSocket *> sockets)
+  void add_try_match(const DTreeContext &context,
+                     Span<const SocketRef *> dsockets,
+                     Span<fn::MFSocket *> sockets)
   {
     int used_sockets = 0;
-    for (const DSocket *dsocket : dsockets) {
+    for (const SocketRef *dsocket : dsockets) {
       if (!dsocket->is_available()) {
         continue;
       }
-      if (!socket_is_mf_data_socket(*dsocket->bsocket()->typeinfo)) {
+      if (!socket_is_mf_data_socket(*dsocket->typeinfo())) {
         continue;
       }
       fn::MFSocket *socket = sockets[used_sockets];
-      this->add(*dsocket, *socket);
+      this->add(DSocket(&context, dsocket), *socket);
       used_sockets++;
     }
   }
 
-  fn::MFOutputSocket &lookup(const DGroupInput &group_input)
-  {
-    fn::MFOutputSocket *socket = socket_by_group_input_id_[group_input.id()];
-    BLI_assert(socket != nullptr);
-    return *socket;
-  }
-
   fn::MFOutputSocket &lookup(const DOutputSocket &dsocket)
   {
-    auto &sockets = sockets_by_dsocket_id_[dsocket.id()];
-    BLI_assert(sockets.size() == 1);
-    return sockets[0]->as_output();
+    return sockets_by_dsocket_.lookup(dsocket)[0]->as_output();
   }
 
   Span<fn::MFInputSocket *> lookup(const DInputSocket &dsocket)
   {
-    return sockets_by_dsocket_id_[dsocket.id()].as_span().cast<fn::MFInputSocket *>();
+    return sockets_by_dsocket_.lookup(dsocket).cast<fn::MFInputSocket *>();
   }
 
   fn::MFInputSocket &lookup_dummy(const DInputSocket &dsocket)
@@ -186,7 +182,7 @@ class MFNetworkTreeMap {
 
   bool is_mapped(const DSocket &dsocket) const
   {
-    return sockets_by_dsocket_id_[dsocket.id()].size() >= 1;
+    return !sockets_by_dsocket_.lookup(dsocket).is_empty();
   }
 };
 
@@ -258,12 +254,7 @@ class SocketMFNetworkBuilder : public MFNetworkBuilderBase {
 
  public:
   SocketMFNetworkBuilder(CommonMFNetworkBuilderData &common, const DSocket &dsocket)
-      : MFNetworkBuilderBase(common), bsocket_(dsocket.bsocket())
-  {
-  }
-
-  SocketMFNetworkBuilder(CommonMFNetworkBuilderData &common, const DGroupInput &group_input)
-      : MFNetworkBuilderBase(common), bsocket_(group_input.bsocket())
+      : MFNetworkBuilderBase(common), bsocket_(dsocket->bsocket())
   {
   }
 
@@ -331,10 +322,10 @@ class SocketMFNetworkBuilder : public MFNetworkBuilderBase {
  */
 class NodeMFNetworkBuilder : public MFNetworkBuilderBase {
  private:
-  const DNode &dnode_;
+  DNode dnode_;
 
  public:
-  NodeMFNetworkBuilder(CommonMFNetworkBuilderData &common, const DNode &dnode)
+  NodeMFNetworkBuilder(CommonMFNetworkBuilderData &common, DNode dnode)
       : MFNetworkBuilderBase(common), dnode_(dnode)
   {
   }
@@ -352,7 +343,7 @@ class NodeMFNetworkBuilder : public MFNetworkBuilderBase {
 
   const fn::MultiFunction &get_not_implemented_fn()
   {
-    return this->get_default_fn("Not Implemented (" + dnode_.name() + ")");
+    return this->get_default_fn("Not Implemented (" + dnode_->name() + ")");
   }
 
   const fn::MultiFunction &get_default_fn(StringRef name);
@@ -377,7 +368,7 @@ class NodeMFNetworkBuilder : public MFNetworkBuilderBase {
    */
   bNode &bnode()
   {
-    return *dnode_.node_ref().bnode();
+    return *dnode_->bnode();
   }
 
   /**
@@ -393,7 +384,7 @@ MFNetworkTreeMap insert_node_tree_into_mf_network(fn::MFNetwork &network,
                                                   const DerivedNodeTree &tree,
                                                   ResourceCollector &resources);
 
-using MultiFunctionByNode = Map<const DNode *, const fn::MultiFunction *>;
+using MultiFunctionByNode = Map<DNode, const fn::MultiFunction *>;
 MultiFunctionByNode get_multi_function_per_node(const DerivedNodeTree &tree,
                                                 ResourceCollector &resources);
 
diff --git a/source/blender/nodes/function/nodes/node_fn_group_instance_id.cc b/source/blender/nodes/function/nodes/node_fn_group_instance_id.cc
index 1e22cde721d..3d0ea201239 100644
--- a/source/blender/nodes/function/nodes/node_fn_group_instance_id.cc
+++ b/source/blender/nodes/function/nodes/node_fn_group_instance_id.cc
@@ -26,9 +26,10 @@ static void fn_node_group_instance_id_expand_in_mf_network(
 {
   const blender::nodes::DNode &node = builder.dnode();
   std::string id = "/";
-  for (const blender::nodes::DParentNode *parent = node.parent(); parent;
-       parent = parent->parent()) {
-    id = "/" + parent->node_ref().name() + id;
+  for (const blender::nodes::DTreeContext *context = node.context();
+       context->parent_node() != nullptr;
+       context = context->parent_context()) {
+    id = "/" + context->parent_node()->name() + id;
   }
   builder.construct_and_set_matching_fn<blender::fn::CustomMF_Constant<std::string>>(
       std::move(id));
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 f3401a2c00d..d0ecb5592da 100644
--- a/source/blender/nodes/function/nodes/node_fn_random_float.cc
+++ b/source/blender/nodes/function/nodes/node_fn_random_float.cc
@@ -67,12 +67,13 @@ static void fn_node_random_float_expand_in_mf_network(
     blender::nodes::NodeMFNetworkBuilder &builder)
 {
   uint32_t function_seed = 1746872341u;
-  const blender::nodes::DNode &node = builder.dnode();
+  blender::nodes::DNode node = builder.dnode();
   const blender::DefaultHash<blender::StringRefNull> hasher;
-  function_seed = 33 * function_seed + hasher(node.name());
-  for (const blender::nodes::DParentNode *parent = node.parent(); parent != nullptr;
-       parent = parent->parent()) {
-    function_seed = 33 * function_seed + hasher(parent->node_ref().name());
+  function_seed = 33 * function_seed + hasher(node->name());
+  for (const blender::nodes::DTreeContext *context = node.context();
+       context->parent_node() != nullptr;
+       context = context->parent_context()) {
+    function_seed = 33 * function_seed + hasher(context->parent_node()->name());
   }
 
   builder.construct_and_set_matching_fn<RandomFloatFunction>(function_seed);
diff --git a/source/blender/nodes/intern/derived_node_tree.cc b/source/blender/nodes/intern/derived_node_tree.cc
index a65641bca2a..c20d6fa943e 100644
--- a/source/blender/nodes/intern/derived_node_tree.cc
+++ b/source/blender/nodes/intern/derived_node_tree.cc
@@ -16,518 +16,256 @@
 
 #include "NOD_derived_node_tree.hh"
 
-#include "BLI_dot_export.hh"
-
-#define UNINITIALIZED_ID UINT32_MAX
-
 namespace blender::nodes {
 
-DerivedNodeTree::DerivedNodeTree(bNodeTree *btree, NodeTreeRefMap &node_tree_refs) : btree_(btree)
+/* Construct a new derived node tree for a given root node tree. The generated derived node tree
+ * does not own the used node tree refs (so that those can be used by others as well). The caller
+ * has to make sure that the node tree refs added to #node_tree_refs live at least as long as the
+ * derived node tree. */
+DerivedNodeTree::DerivedNodeTree(bNodeTree &btree, NodeTreeRefMap &node_tree_refs)
 {
-  BLI_assert(btree != nullptr);
-
-  const NodeTreeRef &main_tree_ref = get_tree_ref_from_map(node_tree_refs, *btree);
-  used_node_tree_refs_.add_new(&main_tree_ref);
-
-  Vector<DNode *> all_nodes;
-  Vector<DGroupInput *> all_group_inputs;
-  Vector<DParentNode *> all_parent_nodes;
-
-  this->insert_nodes_and_links_in_id_order(main_tree_ref, nullptr, all_nodes);
-  this->expand_groups(all_nodes, all_group_inputs, all_parent_nodes, node_tree_refs);
-  this->relink_and_remove_muted_nodes(all_nodes);
-  this->remove_expanded_group_interfaces(all_nodes);
-  this->remove_unused_group_inputs(all_group_inputs);
-  this->store_in_this_and_init_ids(
-      std::move(all_nodes), std::move(all_group_inputs), std::move(all_parent_nodes));
+  /* Construct all possible contexts immediately. This is significantly cheaper than inlining all
+   * node groups. If it still becomes a performance issue in the future, contexts could be
+   * constructed lazily when they are needed. */
+  root_context_ = &this->construct_context_recursively(nullptr, nullptr, btree, node_tree_refs);
 }
 
-BLI_NOINLINE void DerivedNodeTree::insert_nodes_and_links_in_id_order(const NodeTreeRef &tree_ref,
-                                                                      DParentNode *parent,
-                                                                      Vector<DNode *> &all_nodes)
+DTreeContext &DerivedNodeTree::construct_context_recursively(DTreeContext *parent_context,
+                                                             const NodeRef *parent_node,
+                                                             bNodeTree &btree,
+                                                             NodeTreeRefMap &node_tree_refs)
 {
-  Array<DSocket *, 64> sockets_map(tree_ref.sockets().size());
-
-  /* Insert nodes. */
-  for (const NodeRef *node_ref : tree_ref.nodes()) {
-    DNode &node = this->create_node(*node_ref, parent, sockets_map);
-    all_nodes.append(&node);
-  }
-
-  /* Insert links. */
-  for (const NodeRef *node_ref : tree_ref.nodes()) {
-    for (const InputSocketRef *to_socket_ref : node_ref->inputs()) {
-      DInputSocket *to_socket = static_cast<DInputSocket *>(sockets_map[to_socket_ref->id()]);
-      for (const OutputSocketRef *from_socket_ref : to_socket_ref->linked_sockets()) {
-        DOutputSocket *from_socket = static_cast<DOutputSocket *>(
-            sockets_map[from_socket_ref->id()]);
-        to_socket->linked_sockets_.append(from_socket);
-        from_socket->linked_sockets_.append(to_socket);
+  DTreeContext &context = *allocator_.construct<DTreeContext>();
+  context.parent_context_ = parent_context;
+  context.parent_node_ = parent_node;
+  context.tree_ = &get_tree_ref_from_map(node_tree_refs, btree);
+  used_node_tree_refs_.add(context.tree_);
+
+  for (const NodeRef *node : context.tree_->nodes()) {
+    if (node->is_group_node()) {
+      bNode *bnode = node->bnode();
+      bNodeTree *child_btree = reinterpret_cast<bNodeTree *>(bnode->id);
+      if (child_btree != nullptr) {
+        DTreeContext &child = this->construct_context_recursively(
+            &context, node, *child_btree, node_tree_refs);
+        context.children_.add_new(node, &child);
       }
     }
   }
+
+  return context;
 }
 
-DNode &DerivedNodeTree::create_node(const NodeRef &node_ref,
-                                    DParentNode *parent,
-                                    MutableSpan<DSocket *> r_sockets_map)
+DerivedNodeTree::~DerivedNodeTree()
 {
-  DNode &node = *allocator_.construct<DNode>();
-  node.node_ref_ = &node_ref;
-  node.parent_ = parent;
-  node.id_ = UNINITIALIZED_ID;
-
-  node.inputs_ = allocator_.construct_elements_and_pointer_array<DInputSocket>(
-      node_ref.inputs().size());
-  node.outputs_ = allocator_.construct_elements_and_pointer_array<DOutputSocket>(
-      node_ref.outputs().size());
-
-  for (int i : node.inputs_.index_range()) {
-    const InputSocketRef &socket_ref = node_ref.input(i);
-    DInputSocket &socket = *node.inputs_[i];
-    socket.is_multi_input_socket_ = socket_ref.bsocket()->flag & SOCK_MULTI_INPUT;
-    socket.id_ = UNINITIALIZED_ID;
-    socket.node_ = &node;
-    socket.socket_ref_ = &socket_ref;
-
-    r_sockets_map[socket_ref.id()] = &socket;
-  }
-
-  for (int i : node.outputs_.index_range()) {
-    const OutputSocketRef &socket_ref = node_ref.output(i);
-    DOutputSocket &socket = *node.outputs_[i];
-
-    socket.id_ = UNINITIALIZED_ID;
-    socket.node_ = &node;
-    socket.socket_ref_ = &socket_ref;
-
-    r_sockets_map[socket_ref.id()] = &socket;
-  }
-
-  return node;
+  /* Has to be destructed manually, because the context info is allocated in a linear allocator. */
+  this->destruct_context_recursively(root_context_);
 }
 
-BLI_NOINLINE void DerivedNodeTree::expand_groups(Vector<DNode *> &all_nodes,
-                                                 Vector<DGroupInput *> &all_group_inputs,
-                                                 Vector<DParentNode *> &all_parent_nodes,
-                                                 NodeTreeRefMap &node_tree_refs)
+void DerivedNodeTree::destruct_context_recursively(DTreeContext *context)
 {
-  for (int i = 0; i < all_nodes.size(); i++) {
-    DNode &node = *all_nodes[i];
-    if (node.node_ref_->is_group_node()) {
-      /* Muted nodes are relinked in a separate step. */
-      if (!node.node_ref_->is_muted()) {
-        this->expand_group_node(
-            node, all_nodes, all_group_inputs, all_parent_nodes, node_tree_refs);
-      }
-    }
+  for (DTreeContext *child : context->children_.values()) {
+    this->destruct_context_recursively(child);
   }
+  context->~DTreeContext();
 }
 
-BLI_NOINLINE void DerivedNodeTree::expand_group_node(DNode &group_node,
-                                                     Vector<DNode *> &all_nodes,
-                                                     Vector<DGroupInput *> &all_group_inputs,
-                                                     Vector<DParentNode *> &all_parent_nodes,
-                                                     NodeTreeRefMap &node_tree_refs)
+/* Returns true if there are any cycles in the node tree. */
+bool DerivedNodeTree::has_link_cycles() const
 {
-  const NodeRef &group_node_ref = *group_node.node_ref_;
-  BLI_assert(group_node_ref.is_group_node());
-
-  bNodeTree *btree = reinterpret_cast<bNodeTree *>(group_node_ref.bnode()->id);
-  if (btree == nullptr) {
-    return;
+  for (const NodeTreeRef *tree_ref : used_node_tree_refs_) {
+    if (tree_ref->has_link_cycles()) {
+      return true;
+    }
   }
-
-  const NodeTreeRef &group_ref = get_tree_ref_from_map(node_tree_refs, *btree);
-  used_node_tree_refs_.add(&group_ref);
-
-  DParentNode &parent = *allocator_.construct<DParentNode>();
-  parent.id_ = all_parent_nodes.append_and_get_index(&parent);
-  parent.parent_ = group_node.parent_;
-  parent.node_ref_ = &group_node_ref;
-
-  this->insert_nodes_and_links_in_id_order(group_ref, &parent, all_nodes);
-  Span<DNode *> new_nodes_by_id = all_nodes.as_span().take_back(group_ref.nodes().size());
-
-  this->create_group_inputs_for_unlinked_inputs(group_node, all_group_inputs);
-  this->relink_group_inputs(group_ref, new_nodes_by_id, group_node);
-  this->relink_group_outputs(group_ref, new_nodes_by_id, group_node);
+  return false;
 }
 
-BLI_NOINLINE void DerivedNodeTree::create_group_inputs_for_unlinked_inputs(
-    DNode &node, Vector<DGroupInput *> &all_group_inputs)
+/* Calls the given callback on all nodes in the (possibly nested) derived node tree. */
+void DerivedNodeTree::foreach_node(FunctionRef<void(DNode)> callback) const
 {
-  for (DInputSocket *input_socket : node.inputs_) {
-    if (input_socket->is_linked()) {
-      continue;
-    }
-
-    DGroupInput &group_input = *allocator_.construct<DGroupInput>();
-    group_input.id_ = UNINITIALIZED_ID;
-    group_input.socket_ref_ = &input_socket->socket_ref();
-    group_input.parent_ = node.parent_;
-
-    group_input.linked_sockets_.append(input_socket);
-    input_socket->linked_group_inputs_.append(&group_input);
-    all_group_inputs.append(&group_input);
-  }
+  this->foreach_node_in_context_recursive(*root_context_, callback);
 }
 
-BLI_NOINLINE void DerivedNodeTree::relink_group_inputs(const NodeTreeRef &group_ref,
-                                                       Span<DNode *> nodes_by_id,
-                                                       DNode &group_node)
+void DerivedNodeTree::foreach_node_in_context_recursive(const DTreeContext &context,
+                                                        FunctionRef<void(DNode)> callback) const
 {
-  Span<const NodeRef *> node_refs = group_ref.nodes_by_type("NodeGroupInput");
-  if (node_refs.size() == 0) {
-    return;
+  for (const NodeRef *node_ref : context.tree_->nodes()) {
+    callback(DNode(&context, node_ref));
   }
-
-  int input_amount = group_node.inputs().size();
-
-  for (int input_index : IndexRange(input_amount)) {
-    DInputSocket *outside_group = group_node.inputs_[input_index];
-
-    for (DOutputSocket *outside_connected : outside_group->linked_sockets_) {
-      outside_connected->linked_sockets_.remove_first_occurrence_and_reorder(outside_group);
-    }
-
-    for (DGroupInput *outside_connected : outside_group->linked_group_inputs_) {
-      outside_connected->linked_sockets_.remove_first_occurrence_and_reorder(outside_group);
-    }
-
-    for (const NodeRef *input_node_ref : node_refs) {
-      DNode &input_node = *nodes_by_id[input_node_ref->id()];
-      DOutputSocket *inside_group = input_node.outputs_[input_index];
-
-      for (DInputSocket *inside_connected : inside_group->linked_sockets_) {
-        inside_connected->linked_sockets_.remove_first_occurrence_and_reorder(inside_group);
-
-        for (DOutputSocket *outside_connected : outside_group->linked_sockets_) {
-          inside_connected->linked_sockets_.append(outside_connected);
-          outside_connected->linked_sockets_.append(inside_connected);
-        }
-
-        for (DGroupInput *outside_connected : outside_group->linked_group_inputs_) {
-          inside_connected->linked_group_inputs_.append(outside_connected);
-          outside_connected->linked_sockets_.append(inside_connected);
-        }
-      }
-
-      inside_group->linked_sockets_.clear();
-    }
-
-    outside_group->linked_sockets_.clear();
-    outside_group->linked_group_inputs_.clear();
+  for (const DTreeContext *child_context : context.children_.values()) {
+    this->foreach_node_in_context_recursive(*child_context, callback);
   }
 }
 
-BLI_NOINLINE void DerivedNodeTree::relink_group_outputs(const NodeTreeRef &group_ref,
-                                                        Span<DNode *> nodes_by_id,
-                                                        DNode &group_node)
+DOutputSocket DInputSocket::get_corresponding_group_node_output() const
 {
-  Span<const NodeRef *> node_refs = group_ref.nodes_by_type("NodeGroupOutput");
-  if (node_refs.size() == 0) {
-    return;
-  }
-  /* TODO: Pick correct group output node if there are more than one. */
-  const NodeRef &output_node_ref = *node_refs[0];
-  DNode &output_node = *nodes_by_id[output_node_ref.id()];
-
-  int output_amount = group_node.outputs().size();
-  BLI_assert(output_amount == output_node_ref.inputs().size() - 1);
+  BLI_assert(*this);
+  BLI_assert(socket_ref_->node().is_group_output_node());
+  BLI_assert(socket_ref_->index() < socket_ref_->node().inputs().size() - 1);
 
-  for (int output_index : IndexRange(output_amount)) {
-    DOutputSocket *outside_group = group_node.outputs_[output_index];
-    DInputSocket *inside_group = output_node.inputs_[output_index];
+  const DTreeContext *parent_context = context_->parent_context();
+  const NodeRef *parent_node = context_->parent_node();
+  BLI_assert(parent_context != nullptr);
+  BLI_assert(parent_node != nullptr);
 
-    for (DInputSocket *outside_connected : outside_group->linked_sockets_) {
-      outside_connected->linked_sockets_.remove_first_occurrence_and_reorder(outside_group);
-    }
-
-    for (DOutputSocket *inside_connected : inside_group->linked_sockets_) {
-      inside_connected->linked_sockets_.remove_first_occurrence_and_reorder(inside_group);
-
-      for (DInputSocket *outside_connected : outside_group->linked_sockets_) {
-        inside_connected->linked_sockets_.append(outside_connected);
-        outside_connected->linked_sockets_.append(inside_connected);
-      }
-    }
-
-    for (DGroupInput *inside_connected : inside_group->linked_group_inputs_) {
-      inside_connected->linked_sockets_.remove_first_occurrence_and_reorder(inside_group);
-
-      for (DInputSocket *outside_connected : outside_group->linked_sockets_) {
-        inside_connected->linked_sockets_.append(outside_connected);
-        outside_connected->linked_group_inputs_.append(inside_connected);
-      }
-    }
-
-    outside_group->linked_sockets_.clear();
-    inside_group->linked_sockets_.clear();
-  }
+  const int socket_index = socket_ref_->index();
+  return {parent_context, &parent_node->output(socket_index)};
 }
 
-BLI_NOINLINE void DerivedNodeTree::remove_expanded_group_interfaces(Vector<DNode *> &all_nodes)
+Vector<DOutputSocket> DInputSocket::get_corresponding_group_input_sockets() const
 {
-  int index = 0;
-  while (index < all_nodes.size()) {
-    DNode &node = *all_nodes[index];
-    const NodeRef &node_ref = *node.node_ref_;
-    if (node_ref.is_group_node() ||
-        (node.parent_ != nullptr &&
-         (node_ref.is_group_input_node() || node_ref.is_group_output_node()))) {
-      all_nodes.remove_and_reorder(index);
-      node.destruct_with_sockets();
-    }
-    else {
-      index++;
-    }
+  BLI_assert(*this);
+  BLI_assert(socket_ref_->node().is_group_node());
+
+  const DTreeContext *child_context = context_->child_context(socket_ref_->node());
+  BLI_assert(child_context != nullptr);
+
+  const NodeTreeRef &child_tree = child_context->tree();
+  Span<const NodeRef *> group_input_nodes = child_tree.nodes_by_type("NodeGroupInput");
+  const int socket_index = socket_ref_->index();
+  Vector<DOutputSocket> sockets;
+  for (const NodeRef *group_input_node : group_input_nodes) {
+    sockets.append(DOutputSocket(child_context, &group_input_node->output(socket_index)));
   }
+  return sockets;
 }
 
-BLI_NOINLINE void DerivedNodeTree::remove_unused_group_inputs(
-    Vector<DGroupInput *> &all_group_inputs)
+DInputSocket DOutputSocket::get_corresponding_group_node_input() const
 {
-  int index = 0;
-  while (index < all_group_inputs.size()) {
-    DGroupInput &group_input = *all_group_inputs[index];
-    if (group_input.linked_sockets_.is_empty()) {
-      all_group_inputs.remove_and_reorder(index);
-      group_input.~DGroupInput();
-    }
-    else {
-      index++;
-    }
-  }
+  BLI_assert(*this);
+  BLI_assert(socket_ref_->node().is_group_input_node());
+  BLI_assert(socket_ref_->index() < socket_ref_->node().outputs().size() - 1);
+
+  const DTreeContext *parent_context = context_->parent_context();
+  const NodeRef *parent_node = context_->parent_node();
+  BLI_assert(parent_context != nullptr);
+  BLI_assert(parent_node != nullptr);
+
+  const int socket_index = socket_ref_->index();
+  return {parent_context, &parent_node->input(socket_index)};
 }
 
-BLI_NOINLINE void DerivedNodeTree::relink_and_remove_muted_nodes(Vector<DNode *> &all_nodes)
+DInputSocket DOutputSocket::get_active_corresponding_group_output_socket() const
 {
-  int index = 0;
-  while (index < all_nodes.size()) {
-    DNode &node = *all_nodes[index];
-    const NodeRef &node_ref = *node.node_ref_;
-    if (node_ref.is_muted()) {
-      this->relink_muted_node(node);
-      all_nodes.remove_and_reorder(index);
-      node.destruct_with_sockets();
-    }
-    else {
-      index++;
+  BLI_assert(*this);
+  BLI_assert(socket_ref_->node().is_group_node());
+
+  const DTreeContext *child_context = context_->child_context(socket_ref_->node());
+  BLI_assert(child_context != nullptr);
+
+  const NodeTreeRef &child_tree = child_context->tree();
+  Span<const NodeRef *> group_output_nodes = child_tree.nodes_by_type("NodeGroupOutput");
+  const int socket_index = socket_ref_->index();
+  for (const NodeRef *group_output_node : group_output_nodes) {
+    if (group_output_node->bnode()->flag & NODE_DO_OUTPUT || group_output_nodes.size() == 1) {
+      return {child_context, &group_output_node->input(socket_index)};
     }
   }
+  return {};
 }
 
-BLI_NOINLINE void DerivedNodeTree::relink_muted_node(DNode &node)
+/* Call the given callback for every "real" origin socket. "Real" means that reroutes, muted nodes
+ * and node groups are handled by this function. Origin sockets are ones where a node gets its
+ * inputs from. */
+void DInputSocket::foreach_origin_socket(FunctionRef<void(DSocket)> callback) const
 {
-  const bNode &bnode = *node.bnode();
-  LISTBASE_FOREACH (const bNodeLink *, internal_link, &bnode.internal_links) {
-    BLI_assert(internal_link->fromnode == &bnode);
-    BLI_assert(internal_link->tonode == &bnode);
-    bNodeSocket *input_bsocket = internal_link->fromsock;
-    bNodeSocket *output_bsocket = internal_link->tosock;
-
-    /* Find internally linked sockets. */
-    DInputSocket *input_socket = nullptr;
-    DOutputSocket *output_socket = nullptr;
-    for (DInputSocket *socket : node.inputs_) {
-      if (socket->bsocket() == input_bsocket) {
-        input_socket = socket;
-        break;
-      }
-    }
-    for (DOutputSocket *socket : node.outputs_) {
-      if (socket->bsocket() == output_bsocket) {
-        output_socket = socket;
-        break;
+  BLI_assert(*this);
+  for (const OutputSocketRef *linked_socket : socket_ref_->as_input().linked_sockets()) {
+    const NodeRef &linked_node = linked_socket->node();
+    DOutputSocket linked_dsocket{context_, linked_socket};
+
+    if (linked_node.is_muted()) {
+      /* If the node is muted, follow the internal links of the node. */
+      for (const InternalLinkRef *internal_link : linked_node.internal_links()) {
+        if (&internal_link->to() == linked_socket) {
+          DInputSocket input_of_muted_node{context_, &internal_link->from()};
+          input_of_muted_node.foreach_origin_socket(callback);
+        }
       }
     }
-    BLI_assert(input_socket != nullptr);
-    BLI_assert(output_socket != nullptr);
-
-    /* Link sockets connected to the input to sockets that are connected to the internally linked
-     * output. */
-    for (DInputSocket *to_socket : output_socket->linked_sockets_) {
-      for (DOutputSocket *from_socket : input_socket->linked_sockets_) {
-        from_socket->linked_sockets_.append_non_duplicates(to_socket);
-        to_socket->linked_sockets_.append_non_duplicates(from_socket);
+    else if (linked_node.is_group_input_node()) {
+      if (context_->is_root()) {
+        /* This is a group input in the root node group. */
+        callback(linked_dsocket);
       }
-      for (DGroupInput *group_input : input_socket->linked_group_inputs_) {
-        group_input->linked_sockets_.append_non_duplicates(to_socket);
-        to_socket->linked_group_inputs_.append_non_duplicates(group_input);
+      else {
+        DInputSocket socket_in_parent_group = linked_dsocket.get_corresponding_group_node_input();
+        if (socket_in_parent_group->is_linked()) {
+          /* Follow the links coming into the corresponding socket on the parent group node. */
+          socket_in_parent_group.foreach_origin_socket(callback);
+        }
+        else {
+          /* The corresponding input on the parent group node is not connected. Therefore, we use
+           * the value of that input socket directly. */
+          callback(socket_in_parent_group);
+        }
       }
     }
-  }
-
-  /* Remove remaining links from muted node. */
-  for (DInputSocket *to_socket : node.inputs_) {
-    for (DOutputSocket *from_socket : to_socket->linked_sockets_) {
-      from_socket->linked_sockets_.remove_first_occurrence_and_reorder(to_socket);
-    }
-    for (DGroupInput *from_group_input : to_socket->linked_group_inputs_) {
-      from_group_input->linked_sockets_.remove_first_occurrence_and_reorder(to_socket);
-    }
-    to_socket->linked_sockets_.clear();
-    to_socket->linked_group_inputs_.clear();
-  }
-  for (DOutputSocket *from_socket : node.outputs_) {
-    for (DInputSocket *to_socket : from_socket->linked_sockets_) {
-      to_socket->linked_sockets_.remove_first_occurrence_and_reorder(from_socket);
-    }
-    from_socket->linked_sockets_.clear();
-  }
-}
-
-void DNode::destruct_with_sockets()
-{
-  for (DInputSocket *socket : inputs_) {
-    socket->~DInputSocket();
-  }
-  for (DOutputSocket *socket : outputs_) {
-    socket->~DOutputSocket();
-  }
-  this->~DNode();
-}
-
-BLI_NOINLINE void DerivedNodeTree::store_in_this_and_init_ids(
-    Vector<DNode *> &&all_nodes,
-    Vector<DGroupInput *> &&all_group_inputs,
-    Vector<DParentNode *> &&all_parent_nodes)
-{
-  nodes_by_id_ = std::move(all_nodes);
-  group_inputs_ = std::move(all_group_inputs);
-  parent_nodes_ = std::move(all_parent_nodes);
-
-  for (int node_index : nodes_by_id_.index_range()) {
-    DNode *node = nodes_by_id_[node_index];
-    node->id_ = node_index;
-
-    const bNodeType *nodetype = node->node_ref_->bnode()->typeinfo;
-    nodes_by_type_.add(nodetype, node);
-
-    for (DInputSocket *socket : node->inputs_) {
-      socket->id_ = sockets_by_id_.append_and_get_index(socket);
-      input_sockets_.append(socket);
-    }
-    for (DOutputSocket *socket : node->outputs_) {
-      socket->id_ = sockets_by_id_.append_and_get_index(socket);
-      output_sockets_.append(socket);
+    else if (linked_node.is_group_node()) {
+      DInputSocket socket_in_group = linked_dsocket.get_active_corresponding_group_output_socket();
+      if (socket_in_group) {
+        if (socket_in_group->is_linked()) {
+          /* Follow the links coming into the group output node of the child node group. */
+          socket_in_group.foreach_origin_socket(callback);
+        }
+        else {
+          /* The output of the child node group is not connected, so we have to get the value from
+           * that socket. */
+          callback(socket_in_group);
+        }
+      }
     }
-  }
-
-  for (int i : group_inputs_.index_range()) {
-    group_inputs_[i]->id_ = i;
-  }
-}
-
-DerivedNodeTree::~DerivedNodeTree()
-{
-  for (DInputSocket *socket : input_sockets_) {
-    socket->~DInputSocket();
-  }
-  for (DOutputSocket *socket : output_sockets_) {
-    socket->~DOutputSocket();
-  }
-  for (DNode *node : nodes_by_id_) {
-    node->~DNode();
-  }
-  for (DGroupInput *group_input : group_inputs_) {
-    group_input->~DGroupInput();
-  }
-  for (DParentNode *parent : parent_nodes_) {
-    parent->~DParentNode();
-  }
-}
-
-bool DerivedNodeTree::has_link_cycles() const
-{
-  for (const NodeTreeRef *tree : used_node_tree_refs_) {
-    if (tree->has_link_cycles()) {
-      return true;
+    else {
+      /* The normal case: just use the value of a linked output socket. */
+      callback(linked_dsocket);
     }
   }
-  return false;
-}
-
-static dot::Cluster *get_cluster_for_parent(dot::DirectedGraph &graph,
-                                            Map<const DParentNode *, dot::Cluster *> &clusters,
-                                            const DParentNode *parent)
-{
-  if (parent == nullptr) {
-    return nullptr;
-  }
-  return clusters.lookup_or_add_cb(parent, [&]() {
-    dot::Cluster *parent_cluster = get_cluster_for_parent(graph, clusters, parent->parent());
-    bNodeTree *btree = reinterpret_cast<bNodeTree *>(parent->node_ref().bnode()->id);
-    dot::Cluster *new_cluster = &graph.new_cluster(parent->node_ref().name() + " / " +
-                                                   StringRef(btree->id.name + 2));
-    new_cluster->set_parent_cluster(parent_cluster);
-    return new_cluster;
-  });
 }
 
-std::string DerivedNodeTree::to_dot() const
+/* Calls the given callback for every "real" target socket. "Real" means that reroutes, muted nodes
+ * and node groups are handled by this function. Target sockets are on the nodes that use the value
+ * from this socket.   */
+void DOutputSocket::foreach_target_socket(FunctionRef<void(DInputSocket)> callback) const
 {
-  dot::DirectedGraph digraph;
-  digraph.set_rankdir(dot::Attr_rankdir::LeftToRight);
-
-  Map<const DNode *, dot::NodeWithSocketsRef> dot_nodes;
-  Map<const DGroupInput *, dot::NodeWithSocketsRef> dot_group_inputs;
-  Map<const DParentNode *, dot::Cluster *> dot_clusters;
-
-  for (const DNode *node : nodes_by_id_) {
-    dot::Node &dot_node = digraph.new_node("");
-    dot_node.set_background_color("white");
-
-    Vector<std::string> input_names;
-    for (const DInputSocket *socket : node->inputs()) {
-      input_names.append(socket->name());
-    }
-    Vector<std::string> output_names;
-    for (const DOutputSocket *socket : node->outputs()) {
-      output_names.append(socket->name());
+  for (const InputSocketRef *linked_socket : socket_ref_->as_output().linked_sockets()) {
+    const NodeRef &linked_node = linked_socket->node();
+    DInputSocket linked_dsocket{context_, linked_socket};
+
+    if (linked_node.is_muted()) {
+      /* If the target node is muted, follow its internal links. */
+      for (const InternalLinkRef *internal_link : linked_node.internal_links()) {
+        if (&internal_link->from() == linked_socket) {
+          DOutputSocket output_of_muted_node{context_, &internal_link->to()};
+          output_of_muted_node.foreach_target_socket(callback);
+        }
+      }
     }
-
-    dot_nodes.add_new(node,
-                      dot::NodeWithSocketsRef(dot_node, node->name(), input_names, output_names));
-
-    dot::Cluster *cluster = get_cluster_for_parent(digraph, dot_clusters, node->parent());
-    dot_node.set_parent_cluster(cluster);
-  }
-
-  for (const DGroupInput *group_input : group_inputs_) {
-    dot::Node &dot_node = digraph.new_node("");
-    dot_node.set_background_color("white");
-
-    std::string group_input_name = group_input->name();
-    dot_group_inputs.add_new(
-        group_input, dot::NodeWithSocketsRef(dot_node, "Group Input", {}, {group_input_name}));
-
-    dot::Cluster *cluster = get_cluster_for_parent(digraph, dot_clusters, group_input->parent());
-    dot_node.set_parent_cluster(cluster);
-  }
-
-  for (const DNode *to_node : nodes_by_id_) {
-    dot::NodeWithSocketsRef &to_dot_node = dot_nodes.lookup(to_node);
-
-    for (const DInputSocket *to_socket : to_node->inputs()) {
-      for (const DOutputSocket *from_socket : to_socket->linked_sockets()) {
-        const DNode *from_node = &from_socket->node();
-        dot::NodeWithSocketsRef &from_dot_node = dot_nodes.lookup(from_node);
-
-        digraph.new_edge(from_dot_node.output(from_socket->index()),
-                         to_dot_node.input(to_socket->index()));
+    else if (linked_node.is_group_output_node()) {
+      if (context_->is_root()) {
+        /* This is a group output in the root node group. */
+        callback(linked_dsocket);
       }
-      for (const DGroupInput *group_input : to_socket->linked_group_inputs()) {
-        dot::NodeWithSocketsRef &from_dot_node = dot_group_inputs.lookup(group_input);
-
-        digraph.new_edge(from_dot_node.output(0), to_dot_node.input(to_socket->index()));
+      else {
+        /* Follow the links going out of the group node in the parent node group. */
+        DOutputSocket socket_in_parent_group =
+            linked_dsocket.get_corresponding_group_node_output();
+        socket_in_parent_group.foreach_target_socket(callback);
       }
     }
+    else if (linked_node.is_group_node()) {
+      /* Follow the links within the nested node group. */
+      Vector<DOutputSocket> sockets_in_group =
+          linked_dsocket.get_corresponding_group_input_sockets();
+      for (DOutputSocket socket_in_group : sockets_in_group) {
+        socket_in_group.foreach_target_socket(callback);
+      }
+    }
+    else {
+      /* The normal case: just use the linked input socket as target. */
+      callback(linked_dsocket);
+    }
   }
-
-  digraph.set_random_cluster_bgcolors();
-  return digraph.to_dot_string();
 }
 
 }  // namespace blender::nodes
diff --git a/source/blender/nodes/intern/node_geometry_exec.cc b/source/blender/nodes/intern/node_geometry_exec.cc
index 9e62b7d7312..a4fb99a988e 100644
--- a/source/blender/nodes/intern/node_geometry_exec.cc
+++ b/source/blender/nodes/intern/node_geometry_exec.cc
@@ -20,7 +20,6 @@
 
 #include "DEG_depsgraph_query.h"
 
-#include "NOD_derived_node_tree.hh"
 #include "NOD_geometry_exec.hh"
 #include "NOD_type_callbacks.hh"
 
@@ -30,7 +29,7 @@ namespace blender::nodes {
 
 void GeoNodeExecParams::error_message_add(const NodeWarningType type, std::string message) const
 {
-  bNodeTree *btree_cow = node_.node_ref().tree().btree();
+  bNodeTree *btree_cow = node_->btree();
   BLI_assert(btree_cow != nullptr);
   if (btree_cow == nullptr) {
     return;
@@ -40,12 +39,12 @@ void GeoNodeExecParams::error_message_add(const NodeWarningType type, std::strin
   const NodeTreeEvaluationContext context(*self_object_, *modifier_);
 
   BKE_nodetree_error_message_add(
-      *btree_original, context, *node_.bnode(), type, std::move(message));
+      *btree_original, context, *node_->bnode(), type, std::move(message));
 }
 
 const bNodeSocket *GeoNodeExecParams::find_available_socket(const StringRef name) const
 {
-  for (const DSocket *socket : node_.inputs()) {
+  for (const InputSocketRef *socket : node_->inputs()) {
     if (socket->is_available() && socket->name() == name) {
       return socket->bsocket();
     }
@@ -176,7 +175,7 @@ void GeoNodeExecParams::check_extract_input(StringRef identifier,
                                             const CPPType *requested_type) const
 {
   bNodeSocket *found_socket = nullptr;
-  for (const DSocket *socket : node_.inputs()) {
+  for (const InputSocketRef *socket : node_->inputs()) {
     if (socket->identifier() == identifier) {
       found_socket = socket->bsocket();
       break;
@@ -186,7 +185,7 @@ void GeoNodeExecParams::check_extract_input(StringRef identifier,
   if (found_socket == nullptr) {
     std::cout << "Did not find an input socket with the identifier '" << identifier << "'.\n";
     std::cout << "Possible identifiers are: ";
-    for (const DSocket *socket : node_.inputs()) {
+    for (const InputSocketRef *socket : node_->inputs()) {
       if (socket->is_available()) {
         std::cout << "'" << socket->identifier() << "', ";
       }
@@ -218,7 +217,7 @@ void GeoNodeExecParams::check_extract_input(StringRef identifier,
 void GeoNodeExecParams::check_set_output(StringRef identifier, const CPPType &value_type) const
 {
   bNodeSocket *found_socket = nullptr;
-  for (const DSocket *socket : node_.outputs()) {
+  for (const OutputSocketRef *socket : node_->outputs()) {
     if (socket->identifier() == identifier) {
       found_socket = socket->bsocket();
       break;
@@ -228,7 +227,7 @@ void GeoNodeExecParams::check_set_output(StringRef identifier, const CPPType &va
   if (found_socket == nullptr) {
     std::cout << "Did not find an output socket with the identifier '" << identifier << "'.\n";
     std::cout << "Possible identifiers are: ";
-    for (const DSocket *socket : node_.outputs()) {
+    for (const OutputSocketRef *socket : node_->outputs()) {
       if (socket->is_available()) {
         std::cout << "'" << socket->identifier() << "', ";
       }
diff --git a/source/blender/nodes/intern/node_tree_multi_function.cc b/source/blender/nodes/intern/node_tree_multi_function.cc
index c2391667e86..bb1367573f8 100644
--- a/source/blender/nodes/intern/node_tree_multi_function.cc
+++ b/source/blender/nodes/intern/node_tree_multi_function.cc
@@ -29,17 +29,17 @@ const fn::MultiFunction &NodeMFNetworkBuilder::get_default_fn(StringRef name)
   Vector<fn::MFDataType, 10> input_types;
   Vector<fn::MFDataType, 10> output_types;
 
-  for (const DInputSocket *dsocket : dnode_.inputs()) {
+  for (const InputSocketRef *dsocket : dnode_->inputs()) {
     if (dsocket->is_available()) {
-      std::optional<fn::MFDataType> data_type = socket_mf_type_get(*dsocket->bsocket()->typeinfo);
+      std::optional<fn::MFDataType> data_type = socket_mf_type_get(*dsocket->typeinfo());
       if (data_type.has_value()) {
         input_types.append(*data_type);
       }
     }
   }
-  for (const DOutputSocket *dsocket : dnode_.outputs()) {
+  for (const OutputSocketRef *dsocket : dnode_->outputs()) {
     if (dsocket->is_available()) {
-      std::optional<fn::MFDataType> data_type = socket_mf_type_get(*dsocket->bsocket()->typeinfo);
+      std::optional<fn::MFDataType> data_type = socket_mf_type_get(*dsocket->typeinfo());
       if (data_type.has_value()) {
         output_types.append(*data_type);
       }
@@ -57,9 +57,9 @@ static void insert_dummy_node(CommonMFNetworkBuilderData &common, const DNode &d
 
   Vector<fn::MFDataType, stack_capacity> input_types;
   Vector<StringRef, stack_capacity> input_names;
-  Vector<const DInputSocket *, stack_capacity> input_dsockets;
+  Vector<const InputSocketRef *, stack_capacity> input_dsockets;
 
-  for (const DInputSocket *dsocket : dnode.inputs()) {
+  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()) {
@@ -72,9 +72,9 @@ static void insert_dummy_node(CommonMFNetworkBuilderData &common, const DNode &d
 
   Vector<fn::MFDataType, stack_capacity> output_types;
   Vector<StringRef, stack_capacity> output_names;
-  Vector<const DOutputSocket *, stack_capacity> output_dsockets;
+  Vector<const OutputSocketRef *, stack_capacity> output_dsockets;
 
-  for (const DOutputSocket *dsocket : dnode.outputs()) {
+  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()) {
@@ -86,20 +86,20 @@ static void insert_dummy_node(CommonMFNetworkBuilderData &common, const DNode &d
   }
 
   fn::MFDummyNode &dummy_node = common.network.add_dummy(
-      dnode.name(), input_types, output_types, input_names, output_names);
+      dnode->name(), input_types, output_types, input_names, output_names);
 
-  common.network_map.add(input_dsockets, dummy_node.inputs());
-  common.network_map.add(output_dsockets, dummy_node.outputs());
+  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 DInputSocket *socket : dnode.inputs()) {
+  for (const InputSocketRef *socket : dnode->inputs()) {
     if (socket_is_mf_data_socket(*socket->bsocket()->typeinfo)) {
       return true;
     }
   }
-  for (const DOutputSocket *socket : dnode.outputs()) {
+  for (const OutputSocketRef *socket : dnode->outputs()) {
     if (socket_is_mf_data_socket(*socket->bsocket()->typeinfo)) {
       return true;
     }
@@ -107,70 +107,39 @@ static bool has_data_sockets(const DNode &dnode)
   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)
 {
-  for (const DNode *dnode : common.tree.nodes()) {
-    const bNodeType *node_type = dnode->node_ref().bnode()->typeinfo;
+  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};
+      NodeMFNetworkBuilder builder{common, dnode};
       node_type->expand_in_mf_network(builder);
     }
-    else if (has_data_sockets(*dnode)) {
-      insert_dummy_node(common, *dnode);
-    }
-  }
-}
-
-static void insert_group_inputs(CommonMFNetworkBuilderData &common)
-{
-  for (const DGroupInput *group_input : common.tree.group_inputs()) {
-    bNodeSocket *bsocket = group_input->bsocket();
-    if (socket_is_mf_data_socket(*bsocket->typeinfo)) {
-      bNodeSocketType *socktype = bsocket->typeinfo;
-      BLI_assert(socktype->expand_in_mf_network != nullptr);
-
-      SocketMFNetworkBuilder builder{common, *group_input};
-      socktype->expand_in_mf_network(builder);
-
-      fn::MFOutputSocket *from_socket = builder.built_socket();
-      BLI_assert(from_socket != nullptr);
-      common.network_map.add(*group_input, *from_socket);
-    }
-  }
-}
-
-static fn::MFOutputSocket *try_find_origin(CommonMFNetworkBuilderData &common,
-                                           const DInputSocket &to_dsocket)
-{
-  Span<const DOutputSocket *> from_dsockets = to_dsocket.linked_sockets();
-  Span<const DGroupInput *> from_group_inputs = to_dsocket.linked_group_inputs();
-  int total_linked_amount = from_dsockets.size() + from_group_inputs.size();
-  BLI_assert(total_linked_amount <= 1);
-
-  if (total_linked_amount == 0) {
-    return nullptr;
-  }
-
-  if (from_dsockets.size() == 1) {
-    const DOutputSocket &from_dsocket = *from_dsockets[0];
-    if (!from_dsocket.is_available()) {
-      return nullptr;
-    }
-    if (socket_is_mf_data_socket(*from_dsocket.bsocket()->typeinfo)) {
-      return &common.network_map.lookup(from_dsocket);
+    else if (has_data_sockets(dnode)) {
+      insert_dummy_node(common, dnode);
     }
-    return nullptr;
-  }
-
-  const DGroupInput &from_group_input = *from_group_inputs[0];
-  if (socket_is_mf_data_socket(*from_group_input.bsocket()->typeinfo)) {
-    return &common.network_map.lookup(from_group_input);
-  }
-  return nullptr;
+  });
 }
 
 template<typename From, typename To>
@@ -286,78 +255,82 @@ static fn::MFOutputSocket &insert_default_value_for_type(CommonMFNetworkBuilderD
   return node.output(0);
 }
 
-static void insert_links(CommonMFNetworkBuilderData &common)
+static fn::MFOutputSocket *insert_unlinked_input(CommonMFNetworkBuilderData &common,
+                                                 const DInputSocket &dsocket)
 {
-  for (const DInputSocket *to_dsocket : common.tree.input_sockets()) {
-    if (!to_dsocket->is_available()) {
-      continue;
-    }
-    if (!to_dsocket->is_linked()) {
-      continue;
-    }
-    if (!socket_is_mf_data_socket(*to_dsocket->bsocket()->typeinfo)) {
-      continue;
-    }
-
-    Span<fn::MFInputSocket *> to_sockets = common.network_map.lookup(*to_dsocket);
-    BLI_assert(to_sockets.size() >= 1);
-    fn::MFDataType to_type = to_sockets[0]->data_type();
+  BLI_assert(socket_is_mf_data_socket(*dsocket->typeinfo()));
 
-    fn::MFOutputSocket *from_socket = try_find_origin(common, *to_dsocket);
-    if (from_socket == nullptr) {
-      from_socket = &insert_default_value_for_type(common, to_type);
-    }
-
-    fn::MFDataType from_type = from_socket->data_type();
-
-    if (from_type != to_type) {
-      const fn::MultiFunction *conversion_fn = get_implicit_type_conversions().get_conversion(
-          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);
-      }
-    }
+  SocketMFNetworkBuilder builder{common, dsocket};
+  socket_expand_in_mf_network(builder);
 
-    for (fn::MFInputSocket *to_socket : to_sockets) {
-      common.network.add_link(*from_socket, *to_socket);
-    }
-  }
+  fn::MFOutputSocket *built_socket = builder.built_socket();
+  BLI_assert(built_socket != nullptr);
+  return built_socket;
 }
 
-static void insert_unlinked_input(CommonMFNetworkBuilderData &common, const DInputSocket &dsocket)
+static void insert_links_and_unlinked_inputs(CommonMFNetworkBuilderData &common)
 {
-  bNodeSocket *bsocket = dsocket.bsocket();
-  bNodeSocketType *socktype = bsocket->typeinfo;
-  BLI_assert(socktype->expand_in_mf_network != nullptr);
-
-  SocketMFNetworkBuilder builder{common, dsocket};
-  socktype->expand_in_mf_network(builder);
-
-  fn::MFOutputSocket *from_socket = builder.built_socket();
-  BLI_assert(from_socket != nullptr);
+  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;
+      }
 
-  for (fn::MFInputSocket *to_socket : common.network_map.lookup(dsocket)) {
-    common.network.add_link(*from_socket, *to_socket);
-  }
-}
+      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();
 
-static void insert_unlinked_inputs(CommonMFNetworkBuilderData &common)
-{
-  Vector<const DInputSocket *> unlinked_data_inputs;
-  for (const DInputSocket *dsocket : common.tree.input_sockets()) {
-    if (dsocket->is_available()) {
-      if (socket_is_mf_data_socket(*dsocket->bsocket()->typeinfo)) {
-        if (!dsocket->is_linked()) {
-          insert_unlinked_input(common, *dsocket);
+      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(
+            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);
       }
     }
-  }
+  });
 }
 
 /**
@@ -376,9 +349,7 @@ MFNetworkTreeMap insert_node_tree_into_mf_network(fn::MFNetwork &network,
   CommonMFNetworkBuilderData common{resources, network, network_map, tree};
 
   insert_nodes(common);
-  insert_group_inputs(common);
-  insert_links(common);
-  insert_unlinked_inputs(common);
+  insert_links_and_unlinked_inputs(common);
 
   return network_map;
 }
@@ -420,16 +391,17 @@ static NodeExpandType get_node_expand_type(MFNetworkTreeMap &network_map,
     }
   };
 
-  for (const DInputSocket *dsocket : dnode.inputs()) {
+  for (const InputSocketRef *dsocket : dnode->inputs()) {
     if (dsocket->is_available()) {
-      for (fn::MFInputSocket *mf_input : network_map.lookup(*dsocket)) {
+      for (fn::MFInputSocket *mf_input :
+           network_map.lookup(DInputSocket(dnode.context(), dsocket))) {
         check_mf_node(mf_input->node());
       }
     }
   }
-  for (const DOutputSocket *dsocket : dnode.outputs()) {
+  for (const OutputSocketRef *dsocket : dnode->outputs()) {
     if (dsocket->is_available()) {
-      fn::MFOutputSocket &mf_output = network_map.lookup(*dsocket);
+      fn::MFOutputSocket &mf_output = network_map.lookup(DOutputSocket(dnode.context(), dsocket));
       check_mf_node(mf_output.node());
     }
   }
@@ -451,20 +423,21 @@ static const fn::MultiFunction &create_function_for_node_that_expands_into_multi
     ResourceCollector &resources)
 {
   Vector<const fn::MFOutputSocket *> dummy_fn_inputs;
-  for (const DInputSocket *dsocket : dnode.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(*dsocket)) {
+      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 DOutputSocket *dsocket : dnode.outputs()) {
+  for (const OutputSocketRef *dsocket : dnode->outputs()) {
     if (dsocket->is_available()) {
-      fn::MFOutputSocket &mf_output = network_map.lookup(*dsocket);
+      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);
@@ -492,18 +465,18 @@ MultiFunctionByNode get_multi_function_per_node(const DerivedNodeTree &tree,
 
   CommonMFNetworkBuilderData common{resources, network, network_map, tree};
 
-  for (const DNode *dnode : tree.nodes()) {
+  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. */
-      continue;
+      return;
     }
 
-    NodeMFNetworkBuilder builder{common, *dnode};
+    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);
+    const NodeExpandType expand_type = get_node_expand_type(network_map, dnode, &single_function);
 
     switch (expand_type) {
       case NodeExpandType::HasDummyNodes: {
@@ -519,12 +492,12 @@ MultiFunctionByNode get_multi_function_per_node(const DerivedNodeTree &tree,
         /* 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, resources);
+            dnode, network, network_map, resources);
         functions_by_node.add_new(dnode, &fn);
         break;
       }
     }
-  }
+  });
 
   return functions_by_node;
 }
diff --git a/source/blender/nodes/shader/nodes/node_shader_math.cc b/source/blender/nodes/shader/nodes/node_shader_math.cc
index f54914ceba9..7a846031456 100644
--- a/source/blender/nodes/shader/nodes/node_shader_math.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_math.cc
@@ -116,7 +116,7 @@ static void sh_node_math_expand_in_mf_network(blender::nodes::NodeMFNetworkBuild
   blender::fn::MFNetwork &network = builder.network();
   blender::fn::MFFunctionNode &base_node = network.add_function(base_function);
 
-  builder.network_map().add_try_match(dnode.inputs(), base_node.inputs());
+  builder.network_map().add_try_match(*dnode.context(), dnode->inputs(), base_node.inputs());
 
   const bool clamp_output = builder.bnode().custom2 != 0;
   if (clamp_output) {
@@ -126,10 +126,12 @@ static void sh_node_math_expand_in_mf_network(blender::nodes::NodeMFNetworkBuild
                                                               }};
     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(dnode.output(0), clamp_node.output(0));
+    builder.network_map().add(blender::nodes::DOutputSocket(dnode.context(), &dnode->output(0)),
+                              clamp_node.output(0));
   }
   else {
-    builder.network_map().add(dnode.output(0), base_node.output(0));
+    builder.network_map().add(blender::nodes::DOutputSocket(dnode.context(), &dnode->output(0)),
+                              base_node.output(0));
   }
 }
 
diff --git a/source/blender/nodes/shader/nodes/node_shader_value.cc b/source/blender/nodes/shader/nodes/node_shader_value.cc
index 5a8a1b847cc..495c8d12824 100644
--- a/source/blender/nodes/shader/nodes/node_shader_value.cc
+++ b/source/blender/nodes/shader/nodes/node_shader_value.cc
@@ -41,7 +41,7 @@ static int gpu_shader_value(GPUMaterial *mat,
 
 static void sh_node_value_expand_in_mf_network(blender::nodes::NodeMFNetworkBuilder &builder)
 {
-  const bNodeSocket *bsocket = builder.dnode().output(0).bsocket();
+  const bNodeSocket *bsocket = builder.dnode()->output(0).bsocket();
   const bNodeSocketValueFloat *value = (const bNodeSocketValueFloat *)bsocket->default_value;
   builder.construct_and_set_matching_fn<blender::fn::CustomMF_Constant<float>>(value->value);
 }