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

3 files changed, 311 insertions, 601 deletions

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;
 }