Nodes: move NodeTreeRef functionality into node runtime data

The purpose of `NodeTreeRef` was to speed up various queries on a read-only
`bNodeTree`. Not that we have runtime data in nodes and sockets, we can also
store the result of some queries there. This has some benefits:
* No need for a read-only separate node tree data structure which increased
  complexity.
* Makes it easier to reuse cached queries in more parts of Blender that can
  benefit from it.

A downside is that we loose some type safety that we got by having different
types for input and output sockets, as well as internal and non-internal links.

This patch also refactors `DerivedNodeTree` so that it does not use
`NodeTreeRef` anymore, but uses `bNodeTree` directly instead.

To provide a convenient API (that is also close to what `NodeTreeRef` has), a
new approach is implemented: `bNodeTree`, `bNode`, `bNodeSocket` and `bNodeLink`
now have C++ methods declared in `DNA_node_types.h` which are implemented in
`BKE_node_runtime.hh`. To make this work, `makesdna` now skips c++ sections when
parsing dna header files.

No user visible changes are expected.

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

3 files changed, 641 insertions, 298 deletions

diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 1a40c0a18ed..4eb48e9edc9 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -72,7 +72,6 @@
 #include "NOD_function.h"
 #include "NOD_geometry.h"
 #include "NOD_node_declaration.hh"
-#include "NOD_node_tree_ref.hh"
 #include "NOD_shader.h"
 #include "NOD_socket.h"
 #include "NOD_texture.h"
@@ -104,7 +103,6 @@ using blender::nodes::NodeDeclaration;
 using blender::nodes::OutputFieldDependency;
 using blender::nodes::OutputSocketFieldType;
 using blender::nodes::SocketDeclaration;
-using namespace blender::nodes::node_tree_ref_types;
 
 /* Fallback types for undefined tree, nodes, sockets */
 static bNodeTreeType NodeTreeTypeUndefined;
diff --git a/source/blender/blenkernel/intern/node_runtime.cc b/source/blender/blenkernel/intern/node_runtime.cc
new file mode 100644
index 00000000000..20ee3c41534
--- /dev/null
+++ b/source/blender/blenkernel/intern/node_runtime.cc
@@ -0,0 +1,403 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+
+#include "BKE_node.h"
+#include "BKE_node_runtime.hh"
+
+#include "DNA_node_types.h"
+
+#include "BLI_function_ref.hh"
+#include "BLI_stack.hh"
+#include "BLI_task.hh"
+#include "BLI_timeit.hh"
+
+namespace blender::bke::node_tree_runtime {
+
+static void double_checked_lock(std::mutex &mutex, bool &data_is_dirty, FunctionRef<void()> fn)
+{
+  if (!data_is_dirty) {
+    return;
+  }
+  std::lock_guard lock{mutex};
+  if (!data_is_dirty) {
+    return;
+  }
+  fn();
+  data_is_dirty = false;
+}
+
+static void double_checked_lock_with_task_isolation(std::mutex &mutex,
+                                                    bool &data_is_dirty,
+                                                    FunctionRef<void()> fn)
+{
+  double_checked_lock(mutex, data_is_dirty, [&]() { threading::isolate_task(fn); });
+}
+
+static void update_node_vector(const bNodeTree &ntree)
+{
+  bNodeTreeRuntime &tree_runtime = *ntree.runtime;
+  tree_runtime.nodes.clear();
+  tree_runtime.has_undefined_nodes_or_sockets = false;
+  LISTBASE_FOREACH (bNode *, node, &ntree.nodes) {
+    node->runtime->index_in_tree = tree_runtime.nodes.append_and_get_index(node);
+    node->runtime->owner_tree = const_cast<bNodeTree *>(&ntree);
+    tree_runtime.has_undefined_nodes_or_sockets |= node->typeinfo == &NodeTypeUndefined;
+  }
+}
+
+static void update_link_vector(const bNodeTree &ntree)
+{
+  bNodeTreeRuntime &tree_runtime = *ntree.runtime;
+  tree_runtime.links.clear();
+  LISTBASE_FOREACH (bNodeLink *, link, &ntree.links) {
+    tree_runtime.links.append(link);
+  }
+}
+
+static void update_internal_links(const bNodeTree &ntree)
+{
+  bNodeTreeRuntime &tree_runtime = *ntree.runtime;
+  for (bNode *node : tree_runtime.nodes) {
+    node->runtime->internal_links.clear();
+    for (bNodeSocket *socket : node->runtime->outputs) {
+      socket->runtime->internal_link_input = nullptr;
+    }
+    LISTBASE_FOREACH (bNodeLink *, link, &node->internal_links) {
+      node->runtime->internal_links.append(link);
+      link->tosock->runtime->internal_link_input = link->fromsock;
+    }
+  }
+}
+
+static void update_socket_vectors_and_owner_node(const bNodeTree &ntree)
+{
+  bNodeTreeRuntime &tree_runtime = *ntree.runtime;
+  tree_runtime.sockets.clear();
+  tree_runtime.input_sockets.clear();
+  tree_runtime.output_sockets.clear();
+  for (bNode *node : tree_runtime.nodes) {
+    bNodeRuntime &node_runtime = *node->runtime;
+    node_runtime.inputs.clear();
+    node_runtime.outputs.clear();
+    LISTBASE_FOREACH (bNodeSocket *, socket, &node->inputs) {
+      socket->runtime->index_in_node = node_runtime.inputs.append_and_get_index(socket);
+      socket->runtime->index_in_all_sockets = tree_runtime.sockets.append_and_get_index(socket);
+      socket->runtime->index_in_inout_sockets = tree_runtime.input_sockets.append_and_get_index(
+          socket);
+      socket->runtime->owner_node = node;
+      tree_runtime.has_undefined_nodes_or_sockets |= socket->typeinfo == &NodeSocketTypeUndefined;
+    }
+    LISTBASE_FOREACH (bNodeSocket *, socket, &node->outputs) {
+      socket->runtime->index_in_node = node_runtime.outputs.append_and_get_index(socket);
+      socket->runtime->index_in_all_sockets = tree_runtime.sockets.append_and_get_index(socket);
+      socket->runtime->index_in_inout_sockets = tree_runtime.output_sockets.append_and_get_index(
+          socket);
+      socket->runtime->owner_node = node;
+      tree_runtime.has_undefined_nodes_or_sockets |= socket->typeinfo == &NodeSocketTypeUndefined;
+    }
+  }
+}
+
+static void update_directly_linked_links_and_sockets(const bNodeTree &ntree)
+{
+  bNodeTreeRuntime &tree_runtime = *ntree.runtime;
+  for (bNode *node : tree_runtime.nodes) {
+    for (bNodeSocket *socket : node->runtime->inputs) {
+      socket->runtime->directly_linked_links.clear();
+      socket->runtime->directly_linked_sockets.clear();
+    }
+    for (bNodeSocket *socket : node->runtime->outputs) {
+      socket->runtime->directly_linked_links.clear();
+      socket->runtime->directly_linked_sockets.clear();
+    }
+    node->runtime->has_linked_inputs = false;
+    node->runtime->has_linked_outputs = false;
+  }
+  for (bNodeLink *link : tree_runtime.links) {
+    link->fromsock->runtime->directly_linked_links.append(link);
+    link->fromsock->runtime->directly_linked_sockets.append(link->tosock);
+    link->tosock->runtime->directly_linked_links.append(link);
+    link->fromnode->runtime->has_linked_outputs = true;
+    link->tonode->runtime->has_linked_inputs = true;
+  }
+  for (bNodeSocket *socket : tree_runtime.input_sockets) {
+    if (socket->flag & SOCK_MULTI_INPUT) {
+      std::sort(socket->runtime->directly_linked_links.begin(),
+                socket->runtime->directly_linked_links.end(),
+                [&](const bNodeLink *a, const bNodeLink *b) {
+                  return a->multi_input_socket_index > b->multi_input_socket_index;
+                });
+    }
+  }
+  for (bNodeSocket *socket : tree_runtime.input_sockets) {
+    for (bNodeLink *link : socket->runtime->directly_linked_links) {
+      /* Do this after sorting the input links. */
+      socket->runtime->directly_linked_sockets.append(link->fromsock);
+    }
+  }
+}
+
+static void find_logical_origins_for_socket_recursive(
+    bNodeSocket &input_socket,
+    bool only_follow_first_input_link,
+    Vector<bNodeSocket *, 16> &sockets_in_current_chain,
+    Vector<bNodeSocket *> &r_logical_origins,
+    Vector<bNodeSocket *> &r_skipped_origins)
+{
+  if (sockets_in_current_chain.contains(&input_socket)) {
+    /* Protect against reroute recursions. */
+    return;
+  }
+  sockets_in_current_chain.append(&input_socket);
+
+  Span<bNodeLink *> links_to_check = input_socket.runtime->directly_linked_links;
+  if (only_follow_first_input_link) {
+    links_to_check = links_to_check.take_front(1);
+  }
+  for (bNodeLink *link : links_to_check) {
+    if (link->flag & NODE_LINK_MUTED) {
+      continue;
+    }
+    bNodeSocket &origin_socket = *link->fromsock;
+    bNode &origin_node = *link->fromnode;
+    if (!origin_socket.is_available()) {
+      /* Non available sockets are ignored. */
+      continue;
+    }
+    if (origin_node.type == NODE_REROUTE) {
+      bNodeSocket &reroute_input = *origin_node.runtime->inputs[0];
+      bNodeSocket &reroute_output = *origin_node.runtime->outputs[0];
+      r_skipped_origins.append(&reroute_input);
+      r_skipped_origins.append(&reroute_output);
+      find_logical_origins_for_socket_recursive(
+          reroute_input, false, sockets_in_current_chain, r_logical_origins, r_skipped_origins);
+      continue;
+    }
+    if (origin_node.is_muted()) {
+      if (bNodeSocket *mute_input = origin_socket.runtime->internal_link_input) {
+        r_skipped_origins.append(&origin_socket);
+        r_skipped_origins.append(mute_input);
+        find_logical_origins_for_socket_recursive(
+            *mute_input, true, sockets_in_current_chain, r_logical_origins, r_skipped_origins);
+      }
+      continue;
+    }
+    r_logical_origins.append(&origin_socket);
+  }
+
+  sockets_in_current_chain.pop_last();
+}
+
+static void update_logical_origins(const bNodeTree &ntree)
+{
+  bNodeTreeRuntime &tree_runtime = *ntree.runtime;
+  threading::parallel_for(tree_runtime.nodes.index_range(), 128, [&](const IndexRange range) {
+    for (const int i : range) {
+      bNode &node = *tree_runtime.nodes[i];
+      for (bNodeSocket *socket : node.runtime->inputs) {
+        Vector<bNodeSocket *, 16> sockets_in_current_chain;
+        find_logical_origins_for_socket_recursive(
+            *socket,
+            false,
+            sockets_in_current_chain,
+            socket->runtime->logically_linked_sockets,
+            socket->runtime->logically_linked_skipped_sockets);
+      }
+    }
+  });
+}
+
+static void update_nodes_by_type(const bNodeTree &ntree)
+{
+  bNodeTreeRuntime &tree_runtime = *ntree.runtime;
+  tree_runtime.nodes_by_type.clear();
+  for (bNode *node : tree_runtime.nodes) {
+    tree_runtime.nodes_by_type.add(node->typeinfo, node);
+  }
+}
+
+static void update_sockets_by_identifier(const bNodeTree &ntree)
+{
+  bNodeTreeRuntime &tree_runtime = *ntree.runtime;
+  threading::parallel_for(tree_runtime.nodes.index_range(), 128, [&](const IndexRange range) {
+    for (bNode *node : tree_runtime.nodes.as_span().slice(range)) {
+      node->runtime->inputs_by_identifier.clear();
+      node->runtime->outputs_by_identifier.clear();
+      for (bNodeSocket *socket : node->runtime->inputs) {
+        node->runtime->inputs_by_identifier.add_new(socket->identifier, socket);
+      }
+      for (bNodeSocket *socket : node->runtime->outputs) {
+        node->runtime->outputs_by_identifier.add_new(socket->identifier, socket);
+      }
+    }
+  });
+}
+
+enum class ToposortDirection {
+  LeftToRight,
+  RightToLeft,
+};
+
+struct ToposortNodeState {
+  bool is_done = false;
+  bool is_in_stack = false;
+};
+
+static void toposort_from_start_node(const ToposortDirection direction,
+                                     bNode &start_node,
+                                     MutableSpan<ToposortNodeState> node_states,
+                                     Vector<bNode *> &r_sorted_nodes,
+                                     bool &r_cycle_detected)
+{
+  struct Item {
+    bNode *node;
+    int socket_index = 0;
+    int link_index = 0;
+  };
+
+  Stack<Item, 64> nodes_to_check;
+  nodes_to_check.push({&start_node});
+  while (!nodes_to_check.is_empty()) {
+    Item &item = nodes_to_check.peek();
+    bNode &node = *item.node;
+    const Span<bNodeSocket *> sockets = (direction == ToposortDirection::LeftToRight) ?
+                                            node.runtime->inputs :
+                                            node.runtime->outputs;
+    while (true) {
+      if (item.socket_index == sockets.size()) {
+        /* All sockets have already been visited. */
+        break;
+      }
+      bNodeSocket &socket = *sockets[item.socket_index];
+      const Span<bNodeSocket *> linked_sockets = socket.runtime->directly_linked_sockets;
+      if (item.link_index == linked_sockets.size()) {
+        /* All links connected to this socket have already been visited. */
+        item.socket_index++;
+        item.link_index = 0;
+        continue;
+      }
+      bNodeSocket &linked_socket = *linked_sockets[item.link_index];
+      bNode &linked_node = *linked_socket.runtime->owner_node;
+      ToposortNodeState &linked_node_state = node_states[linked_node.runtime->index_in_tree];
+      if (linked_node_state.is_done) {
+        /* The linked node has already been visited. */
+        item.link_index++;
+        continue;
+      }
+      if (linked_node_state.is_in_stack) {
+        r_cycle_detected = true;
+      }
+      else {
+        nodes_to_check.push({&linked_node});
+        linked_node_state.is_in_stack = true;
+      }
+      break;
+    }
+
+    /* If no other element has been pushed, the current node can be pushed to the sorted list. */
+    if (&item == &nodes_to_check.peek()) {
+      ToposortNodeState &node_state = node_states[node.runtime->index_in_tree];
+      node_state.is_done = true;
+      node_state.is_in_stack = false;
+      r_sorted_nodes.append(&node);
+      nodes_to_check.pop();
+    }
+  }
+}
+
+static void update_toposort(const bNodeTree &ntree,
+                            const ToposortDirection direction,
+                            Vector<bNode *> &r_sorted_nodes,
+                            bool &r_cycle_detected)
+{
+  bNodeTreeRuntime &tree_runtime = *ntree.runtime;
+  r_sorted_nodes.clear();
+  r_sorted_nodes.reserve(tree_runtime.nodes.size());
+  r_cycle_detected = false;
+
+  Array<ToposortNodeState> node_states(tree_runtime.nodes.size());
+  for (bNode *node : tree_runtime.nodes) {
+    if (node_states[node->runtime->index_in_tree].is_done) {
+      /* Ignore nodes that are done already. */
+      continue;
+    }
+    if ((direction == ToposortDirection::LeftToRight) ? node->runtime->has_linked_outputs :
+                                                        node->runtime->has_linked_inputs) {
+      /* Ignore non-start nodes. */
+      continue;
+    }
+    toposort_from_start_node(direction, *node, node_states, r_sorted_nodes, r_cycle_detected);
+  }
+
+  if (r_sorted_nodes.size() < tree_runtime.nodes.size()) {
+    r_cycle_detected = true;
+    for (bNode *node : tree_runtime.nodes) {
+      if (node_states[node->runtime->index_in_tree].is_done) {
+        /* Ignore nodes that are done already. */
+        continue;
+      }
+      /* Start toposort at this node which is somewhere in the middle of a loop. */
+      toposort_from_start_node(direction, *node, node_states, r_sorted_nodes, r_cycle_detected);
+    }
+  }
+
+  BLI_assert(tree_runtime.nodes.size() == r_sorted_nodes.size());
+}
+
+static void update_group_output_node(const bNodeTree &ntree)
+{
+  bNodeTreeRuntime &tree_runtime = *ntree.runtime;
+  const bNodeType *node_type = nodeTypeFind("NodeGroupOutput");
+  const Span<bNode *> group_output_nodes = tree_runtime.nodes_by_type.lookup(node_type);
+  if (group_output_nodes.is_empty()) {
+    tree_runtime.group_output_node = nullptr;
+  }
+  else if (group_output_nodes.size() == 1) {
+    tree_runtime.group_output_node = group_output_nodes[0];
+  }
+  else {
+    for (bNode *group_output : group_output_nodes) {
+      if (group_output->flag & NODE_DO_OUTPUT) {
+        tree_runtime.group_output_node = group_output;
+        break;
+      }
+    }
+  }
+}
+
+static void ensure_topology_cache(const bNodeTree &ntree)
+{
+  bNodeTreeRuntime &tree_runtime = *ntree.runtime;
+  double_checked_lock_with_task_isolation(
+      tree_runtime.topology_cache_mutex, tree_runtime.topology_cache_is_dirty, [&]() {
+        update_node_vector(ntree);
+        update_link_vector(ntree);
+        update_socket_vectors_and_owner_node(ntree);
+        update_internal_links(ntree);
+        update_directly_linked_links_and_sockets(ntree);
+        threading::parallel_invoke([&]() { update_logical_origins(ntree); },
+                                   [&]() { update_nodes_by_type(ntree); },
+                                   [&]() { update_sockets_by_identifier(ntree); },
+                                   [&]() {
+                                     update_toposort(ntree,
+                                                     ToposortDirection::LeftToRight,
+                                                     tree_runtime.toposort_left_to_right,
+                                                     tree_runtime.has_link_cycle);
+                                   },
+                                   [&]() {
+                                     bool dummy;
+                                     update_toposort(ntree,
+                                                     ToposortDirection::RightToLeft,
+                                                     tree_runtime.toposort_right_to_left,
+                                                     dummy);
+                                   });
+        update_group_output_node(ntree);
+        tree_runtime.topology_cache_exists = true;
+      });
+}
+
+}  // namespace blender::bke::node_tree_runtime
+
+void bNodeTree::ensure_topology_cache() const
+{
+  blender::bke::node_tree_runtime::ensure_topology_cache(*this);
+}
diff --git a/source/blender/blenkernel/intern/node_tree_update.cc b/source/blender/blenkernel/intern/node_tree_update.cc
index 58084226b00..716b8ce31d3 100644
--- a/source/blender/blenkernel/intern/node_tree_update.cc
+++ b/source/blender/blenkernel/intern/node_tree_update.cc
@@ -5,6 +5,7 @@
 #include "BLI_noise.hh"
 #include "BLI_set.hh"
 #include "BLI_stack.hh"
+#include "BLI_timeit.hh"
 #include "BLI_vector_set.hh"
 
 #include "DNA_anim_types.h"
@@ -21,7 +22,6 @@
 #include "MOD_nodes.h"
 
 #include "NOD_node_declaration.hh"
-#include "NOD_node_tree_ref.hh"
 #include "NOD_texture.h"
 
 #include "DEG_depsgraph_query.h"
@@ -50,6 +50,7 @@ enum eNodeTreeChangedFlag {
 static void add_tree_tag(bNodeTree *ntree, const eNodeTreeChangedFlag flag)
 {
   ntree->runtime->changed_flag |= flag;
+  ntree->runtime->topology_cache_is_dirty = true;
 }
 
 static void add_node_tag(bNodeTree *ntree, bNode *node, const eNodeTreeChangedFlag flag)
@@ -73,31 +74,32 @@ static bool is_field_socket_type(eNodeSocketDatatype type)
   return ELEM(type, SOCK_FLOAT, SOCK_INT, SOCK_BOOLEAN, SOCK_VECTOR, SOCK_RGBA);
 }
 
-static bool is_field_socket_type(const SocketRef &socket)
+static bool is_field_socket_type(const bNodeSocket &socket)
 {
-  return is_field_socket_type((eNodeSocketDatatype)socket.typeinfo()->type);
+  return is_field_socket_type((eNodeSocketDatatype)socket.typeinfo->type);
 }
 
-static InputSocketFieldType get_interface_input_field_type(const NodeRef &node,
-                                                           const InputSocketRef &socket)
+static InputSocketFieldType get_interface_input_field_type(const bNode &node,
+                                                           const bNodeSocket &socket)
 {
   if (!is_field_socket_type(socket)) {
     return InputSocketFieldType::None;
   }
-  if (node.is_reroute_node()) {
+  if (node.type == NODE_REROUTE) {
     return InputSocketFieldType::IsSupported;
   }
-  if (node.is_group_output_node()) {
+  if (node.type == NODE_GROUP_OUTPUT) {
     /* Outputs always support fields when the data type is correct. */
     return InputSocketFieldType::IsSupported;
   }
-  if (node.is_undefined()) {
+  if (node.typeinfo == &NodeTypeUndefined) {
     return InputSocketFieldType::None;
   }
-  if (node.bnode()->type == NODE_CUSTOM) {
+  if (node.type == NODE_CUSTOM) {
     return InputSocketFieldType::None;
   }
 
+  /* TODO: Ensure declaration exists. */
   const NodeDeclaration *node_decl = node.declaration();
 
   /* Node declarations should be implemented for nodes involved here. */
@@ -116,25 +118,25 @@ static InputSocketFieldType get_interface_input_field_type(const NodeRef &node,
   return field_type;
 }
 
-static OutputFieldDependency get_interface_output_field_dependency(const NodeRef &node,
-                                                                   const OutputSocketRef &socket)
+static OutputFieldDependency get_interface_output_field_dependency(const bNode &node,
+                                                                   const bNodeSocket &socket)
 {
   if (!is_field_socket_type(socket)) {
     /* Non-field sockets always output data. */
     return OutputFieldDependency::ForDataSource();
   }
-  if (node.is_reroute_node()) {
+  if (node.type == NODE_REROUTE) {
     /* The reroute just forwards what is passed in. */
     return OutputFieldDependency::ForDependentField();
   }
-  if (node.is_group_input_node()) {
+  if (node.type == NODE_GROUP_INPUT) {
     /* Input nodes get special treatment in #determine_group_input_states. */
     return OutputFieldDependency::ForDependentField();
   }
-  if (node.is_undefined()) {
+  if (node.typeinfo == &NodeTypeUndefined) {
     return OutputFieldDependency::ForDataSource();
   }
-  if (node.bnode()->type == NODE_CUSTOM) {
+  if (node.type == NODE_CUSTOM) {
     return OutputFieldDependency::ForDataSource();
   }
 
@@ -153,12 +155,13 @@ static OutputFieldDependency get_interface_output_field_dependency(const NodeRef
   return socket_decl.output_field_dependency();
 }
 
-static FieldInferencingInterface get_dummy_field_inferencing_interface(const NodeRef &node)
+static FieldInferencingInterface get_dummy_field_inferencing_interface(const bNode &node)
 {
   FieldInferencingInterface inferencing_interface;
-  inferencing_interface.inputs.append_n_times(InputSocketFieldType::None, node.inputs().size());
+  inferencing_interface.inputs.append_n_times(InputSocketFieldType::None,
+                                              node.input_sockets().size());
   inferencing_interface.outputs.append_n_times(OutputFieldDependency::ForDataSource(),
-                                               node.outputs().size());
+                                               node.output_sockets().size());
   return inferencing_interface;
 }
 
@@ -167,11 +170,11 @@ static FieldInferencingInterface get_dummy_field_inferencing_interface(const Nod
  * In the future, this information can be stored in the node declaration. This would allow this
  * function to return a reference, making it more efficient.
  */
-static FieldInferencingInterface get_node_field_inferencing_interface(const NodeRef &node)
+static FieldInferencingInterface get_node_field_inferencing_interface(const bNode &node)
 {
   /* Node groups already reference all required information, so just return that. */
-  if (node.is_group_node()) {
-    bNodeTree *group = (bNodeTree *)node.bnode()->id;
+  if (node.is_group()) {
+    bNodeTree *group = (bNodeTree *)node.id;
     if (group == nullptr) {
       return FieldInferencingInterface();
     }
@@ -187,11 +190,11 @@ static FieldInferencingInterface get_node_field_inferencing_interface(const Node
   }
 
   FieldInferencingInterface inferencing_interface;
-  for (const InputSocketRef *input_socket : node.inputs()) {
+  for (const bNodeSocket *input_socket : node.input_sockets()) {
     inferencing_interface.inputs.append(get_interface_input_field_type(node, *input_socket));
   }
 
-  for (const OutputSocketRef *output_socket : node.outputs()) {
+  for (const bNodeSocket *output_socket : node.output_sockets()) {
     inferencing_interface.outputs.append(
         get_interface_output_field_dependency(node, *output_socket));
   }
@@ -215,11 +218,11 @@ struct SocketFieldState {
   bool requires_single = false;
 };
 
-static Vector<const InputSocketRef *> gather_input_socket_dependencies(
-    const OutputFieldDependency &field_dependency, const NodeRef &node)
+static Vector<const bNodeSocket *> gather_input_socket_dependencies(
+    const OutputFieldDependency &field_dependency, const bNode &node)
 {
   const OutputSocketFieldType type = field_dependency.field_type();
-  Vector<const InputSocketRef *> input_sockets;
+  Vector<const bNodeSocket *> input_sockets;
   switch (type) {
     case OutputSocketFieldType::FieldSource:
     case OutputSocketFieldType::None: {
@@ -227,13 +230,13 @@ static Vector<const InputSocketRef *> gather_input_socket_dependencies(
     }
     case OutputSocketFieldType::DependentField: {
       /* This output depends on all inputs. */
-      input_sockets.extend(node.inputs());
+      input_sockets.extend(node.input_sockets());
       break;
     }
     case OutputSocketFieldType::PartiallyDependent: {
       /* This output depends only on a few inputs. */
       for (const int i : field_dependency.linked_input_indices()) {
-        input_sockets.append(&node.input(i));
+        input_sockets.append(&node.input_socket(i));
       }
       break;
     }
@@ -246,8 +249,7 @@ static Vector<const InputSocketRef *> gather_input_socket_dependencies(
  * to figure out if it is always a field or if it depends on any group inputs.
  */
 static OutputFieldDependency find_group_output_dependencies(
-    const InputSocketRef &group_output_socket,
-    const Span<SocketFieldState> field_state_by_socket_id)
+    const bNodeSocket &group_output_socket, const Span<SocketFieldState> field_state_by_socket_id)
 {
   if (!is_field_socket_type(group_output_socket)) {
     return OutputFieldDependency::ForDataSource();
@@ -255,8 +257,8 @@ static OutputFieldDependency find_group_output_dependencies(
 
   /* Use a Set here instead of an array indexed by socket id, because we my only need to look at
    * very few sockets. */
-  Set<const InputSocketRef *> handled_sockets;
-  Stack<const InputSocketRef *> sockets_to_check;
+  Set<const bNodeSocket *> handled_sockets;
+  Stack<const bNodeSocket *> sockets_to_check;
 
   handled_sockets.add(&group_output_socket);
   sockets_to_check.push(&group_output_socket);
@@ -265,20 +267,21 @@ static OutputFieldDependency find_group_output_dependencies(
   Vector<int> linked_input_indices;
 
   while (!sockets_to_check.is_empty()) {
-    const InputSocketRef *input_socket = sockets_to_check.pop();
+    const bNodeSocket *input_socket = sockets_to_check.pop();
 
     if (!input_socket->is_directly_linked() &&
-        !field_state_by_socket_id[input_socket->id()].is_single) {
+        !field_state_by_socket_id[input_socket->index_in_tree()].is_single) {
       /* This socket uses a field as input by default. */
       return OutputFieldDependency::ForFieldSource();
     }
 
-    for (const OutputSocketRef *origin_socket : input_socket->directly_linked_sockets()) {
-      const NodeRef &origin_node = origin_socket->node();
-      const SocketFieldState &origin_state = field_state_by_socket_id[origin_socket->id()];
+    for (const bNodeSocket *origin_socket : input_socket->directly_linked_sockets()) {
+      const bNode &origin_node = origin_socket->owner_node();
+      const SocketFieldState &origin_state =
+          field_state_by_socket_id[origin_socket->index_in_tree()];
 
       if (origin_state.is_field_source) {
-        if (origin_node.is_group_input_node()) {
+        if (origin_node.type == NODE_GROUP_INPUT) {
           /* Found a group input that the group output depends on. */
           linked_input_indices.append_non_duplicates(origin_socket->index());
         }
@@ -294,12 +297,12 @@ static OutputFieldDependency find_group_output_dependencies(
             inferencing_interface.outputs[origin_socket->index()];
 
         /* Propagate search further to the left. */
-        for (const InputSocketRef *origin_input_socket :
+        for (const bNodeSocket *origin_input_socket :
              gather_input_socket_dependencies(field_dependency, origin_node)) {
           if (!origin_input_socket->is_available()) {
             continue;
           }
-          if (!field_state_by_socket_id[origin_input_socket->id()].is_single) {
+          if (!field_state_by_socket_id[origin_input_socket->index_in_tree()].is_single) {
             if (handled_sockets.add(origin_input_socket)) {
               sockets_to_check.push(origin_input_socket);
             }
@@ -312,17 +315,16 @@ static OutputFieldDependency find_group_output_dependencies(
 }
 
 static void propagate_data_requirements_from_right_to_left(
-    const NodeTreeRef &tree, const MutableSpan<SocketFieldState> field_state_by_socket_id)
+    const bNodeTree &tree, const MutableSpan<SocketFieldState> field_state_by_socket_id)
 {
-  const NodeTreeRef::ToposortResult toposort_result = tree.toposort(
-      NodeTreeRef::ToposortDirection::RightToLeft);
+  const Span<const bNode *> toposort_result = tree.toposort_right_to_left();
 
-  for (const NodeRef *node : toposort_result.sorted_nodes) {
+  for (const bNode *node : toposort_result) {
     const FieldInferencingInterface inferencing_interface = get_node_field_inferencing_interface(
         *node);
 
-    for (const OutputSocketRef *output_socket : node->outputs()) {
-      SocketFieldState &state = field_state_by_socket_id[output_socket->id()];
+    for (const bNodeSocket *output_socket : node->output_sockets()) {
+      SocketFieldState &state = field_state_by_socket_id[output_socket->index_in_tree()];
 
       const OutputFieldDependency &field_dependency =
           inferencing_interface.outputs[output_socket->index()];
@@ -338,17 +340,18 @@ static void propagate_data_requirements_from_right_to_left(
 
       /* The output is required to be a single value when it is connected to any input that does
        * not support fields. */
-      for (const InputSocketRef *target_socket : output_socket->directly_linked_sockets()) {
+      for (const bNodeSocket *target_socket : output_socket->directly_linked_sockets()) {
         if (target_socket->is_available()) {
-          state.requires_single |= field_state_by_socket_id[target_socket->id()].requires_single;
+          state.requires_single |=
+              field_state_by_socket_id[target_socket->index_in_tree()].requires_single;
         }
       }
 
       if (state.requires_single) {
         bool any_input_is_field_implicitly = false;
-        const Vector<const InputSocketRef *> connected_inputs = gather_input_socket_dependencies(
+        const Vector<const bNodeSocket *> connected_inputs = gather_input_socket_dependencies(
             field_dependency, *node);
-        for (const InputSocketRef *input_socket : connected_inputs) {
+        for (const bNodeSocket *input_socket : connected_inputs) {
           if (!input_socket->is_available()) {
             continue;
           }
@@ -367,16 +370,16 @@ static void propagate_data_requirements_from_right_to_left(
         else {
           /* If the output is required to be a single value, the connected inputs in the same node
            * must not be fields as well. */
-          for (const InputSocketRef *input_socket : connected_inputs) {
-            field_state_by_socket_id[input_socket->id()].requires_single = true;
+          for (const bNodeSocket *input_socket : connected_inputs) {
+            field_state_by_socket_id[input_socket->index_in_tree()].requires_single = true;
           }
         }
       }
     }
 
     /* Some inputs do not require fields independent of what the outputs are connected to. */
-    for (const InputSocketRef *input_socket : node->inputs()) {
-      SocketFieldState &state = field_state_by_socket_id[input_socket->id()];
+    for (const bNodeSocket *input_socket : node->input_sockets()) {
+      SocketFieldState &state = field_state_by_socket_id[input_socket->index_in_tree()];
       if (inferencing_interface.inputs[input_socket->index()] == InputSocketFieldType::None) {
         state.requires_single = true;
         state.is_always_single = true;
@@ -386,14 +389,14 @@ static void propagate_data_requirements_from_right_to_left(
 }
 
 static void determine_group_input_states(
-    const NodeTreeRef &tree,
+    const bNodeTree &tree,
     FieldInferencingInterface &new_inferencing_interface,
     const MutableSpan<SocketFieldState> field_state_by_socket_id)
 {
   {
     /* Non-field inputs never support fields. */
     int index;
-    LISTBASE_FOREACH_INDEX (bNodeSocket *, group_input, &tree.btree()->inputs, index) {
+    LISTBASE_FOREACH_INDEX (bNodeSocket *, group_input, &tree.inputs, index) {
       if (!is_field_socket_type((eNodeSocketDatatype)group_input->type)) {
         new_inferencing_interface.inputs[index] = InputSocketFieldType::None;
       }
@@ -401,18 +404,18 @@ static void determine_group_input_states(
   }
   /* Check if group inputs are required to be single values, because they are (indirectly)
    * connected to some socket that does not support fields. */
-  for (const NodeRef *node : tree.nodes_by_type("NodeGroupInput")) {
-    for (const OutputSocketRef *output_socket : node->outputs().drop_back(1)) {
-      SocketFieldState &state = field_state_by_socket_id[output_socket->id()];
+  for (const bNode *node : tree.nodes_by_type("NodeGroupInput")) {
+    for (const bNodeSocket *output_socket : node->output_sockets().drop_back(1)) {
+      SocketFieldState &state = field_state_by_socket_id[output_socket->index_in_tree()];
       if (state.requires_single) {
         new_inferencing_interface.inputs[output_socket->index()] = InputSocketFieldType::None;
       }
     }
   }
   /* If an input does not support fields, this should be reflected in all Group Input nodes. */
-  for (const NodeRef *node : tree.nodes_by_type("NodeGroupInput")) {
-    for (const OutputSocketRef *output_socket : node->outputs().drop_back(1)) {
-      SocketFieldState &state = field_state_by_socket_id[output_socket->id()];
+  for (const bNode *node : tree.nodes_by_type("NodeGroupInput")) {
+    for (const bNodeSocket *output_socket : node->output_sockets().drop_back(1)) {
+      SocketFieldState &state = field_state_by_socket_id[output_socket->index_in_tree()];
       const bool supports_field = new_inferencing_interface.inputs[output_socket->index()] !=
                                   InputSocketFieldType::None;
       if (supports_field) {
@@ -423,19 +426,19 @@ static void determine_group_input_states(
         state.requires_single = true;
       }
     }
-    SocketFieldState &dummy_socket_state = field_state_by_socket_id[node->outputs().last()->id()];
+    SocketFieldState &dummy_socket_state =
+        field_state_by_socket_id[node->output_sockets().last()->index_in_tree()];
     dummy_socket_state.requires_single = true;
   }
 }
 
 static void propagate_field_status_from_left_to_right(
-    const NodeTreeRef &tree, const MutableSpan<SocketFieldState> field_state_by_socket_id)
+    const bNodeTree &tree, const MutableSpan<SocketFieldState> field_state_by_socket_id)
 {
-  const NodeTreeRef::ToposortResult toposort_result = tree.toposort(
-      NodeTreeRef::ToposortDirection::LeftToRight);
+  const Span<const bNode *> toposort_result = tree.toposort_left_to_right();
 
-  for (const NodeRef *node : toposort_result.sorted_nodes) {
-    if (node->is_group_input_node()) {
+  for (const bNode *node : toposort_result) {
+    if (node->type == NODE_GROUP_INPUT) {
       continue;
     }
 
@@ -443,22 +446,22 @@ static void propagate_field_status_from_left_to_right(
         *node);
 
     /* Update field state of input sockets, also taking into account linked origin sockets. */
-    for (const InputSocketRef *input_socket : node->inputs()) {
-      SocketFieldState &state = field_state_by_socket_id[input_socket->id()];
+    for (const bNodeSocket *input_socket : node->input_sockets()) {
+      SocketFieldState &state = field_state_by_socket_id[input_socket->index_in_tree()];
       if (state.is_always_single) {
         state.is_single = true;
         continue;
       }
       state.is_single = true;
-      if (input_socket->directly_linked_sockets().is_empty()) {
+      if (!input_socket->is_directly_linked()) {
         if (inferencing_interface.inputs[input_socket->index()] ==
             InputSocketFieldType::Implicit) {
           state.is_single = false;
         }
       }
       else {
-        for (const OutputSocketRef *origin_socket : input_socket->directly_linked_sockets()) {
-          if (!field_state_by_socket_id[origin_socket->id()].is_single) {
+        for (const bNodeSocket *origin_socket : input_socket->directly_linked_sockets()) {
+          if (!field_state_by_socket_id[origin_socket->index_in_tree()].is_single) {
             state.is_single = false;
             break;
           }
@@ -467,8 +470,8 @@ static void propagate_field_status_from_left_to_right(
     }
 
     /* Update field state of output sockets, also taking into account input sockets. */
-    for (const OutputSocketRef *output_socket : node->outputs()) {
-      SocketFieldState &state = field_state_by_socket_id[output_socket->id()];
+    for (const bNodeSocket *output_socket : node->output_sockets()) {
+      SocketFieldState &state = field_state_by_socket_id[output_socket->index_in_tree()];
       const OutputFieldDependency &field_dependency =
           inferencing_interface.outputs[output_socket->index()];
 
@@ -484,12 +487,12 @@ static void propagate_field_status_from_left_to_right(
         }
         case OutputSocketFieldType::PartiallyDependent:
         case OutputSocketFieldType::DependentField: {
-          for (const InputSocketRef *input_socket :
+          for (const bNodeSocket *input_socket :
                gather_input_socket_dependencies(field_dependency, *node)) {
             if (!input_socket->is_available()) {
               continue;
             }
-            if (!field_state_by_socket_id[input_socket->id()].is_single) {
+            if (!field_state_by_socket_id[input_socket->index_in_tree()].is_single) {
               state.is_single = false;
               break;
             }
@@ -501,17 +504,18 @@ static void propagate_field_status_from_left_to_right(
   }
 }
 
-static void determine_group_output_states(const NodeTreeRef &tree,
+static void determine_group_output_states(const bNodeTree &tree,
                                           FieldInferencingInterface &new_inferencing_interface,
                                           const Span<SocketFieldState> field_state_by_socket_id)
 {
-  for (const NodeRef *group_output_node : tree.nodes_by_type("NodeGroupOutput")) {
+  for (const bNode *group_output_node : tree.nodes_by_type("NodeGroupOutput")) {
     /* Ignore inactive group output nodes. */
-    if (!(group_output_node->bnode()->flag & NODE_DO_OUTPUT)) {
+    if (!(group_output_node->flag & NODE_DO_OUTPUT)) {
       continue;
     }
     /* Determine dependencies of all group outputs. */
-    for (const InputSocketRef *group_output_socket : group_output_node->inputs().drop_back(1)) {
+    for (const bNodeSocket *group_output_socket :
+         group_output_node->input_sockets().drop_back(1)) {
       OutputFieldDependency field_dependency = find_group_output_dependencies(
           *group_output_socket, field_state_by_socket_id);
       new_inferencing_interface.outputs[group_output_socket->index()] = std::move(
@@ -521,7 +525,7 @@ static void determine_group_output_states(const NodeTreeRef &tree,
   }
 }
 
-static void update_socket_shapes(const NodeTreeRef &tree,
+static void update_socket_shapes(const bNodeTree &tree,
                                  const Span<SocketFieldState> field_state_by_socket_id)
 {
   const eNodeSocketDisplayShape requires_data_shape = SOCK_DISPLAY_SHAPE_CIRCLE;
@@ -541,32 +545,30 @@ static void update_socket_shapes(const NodeTreeRef &tree,
     return data_but_can_be_field_shape;
   };
 
-  for (const InputSocketRef *socket : tree.input_sockets()) {
-    bNodeSocket *bsocket = socket->bsocket();
-    const SocketFieldState &state = field_state_by_socket_id[socket->id()];
-    bsocket->display_shape = get_shape_for_state(state);
+  for (const bNodeSocket *socket : tree.all_input_sockets()) {
+    const SocketFieldState &state = field_state_by_socket_id[socket->index_in_tree()];
+    const_cast<bNodeSocket *>(socket)->display_shape = get_shape_for_state(state);
   }
-  for (const OutputSocketRef *socket : tree.output_sockets()) {
-    bNodeSocket *bsocket = socket->bsocket();
-    const SocketFieldState &state = field_state_by_socket_id[socket->id()];
-    bsocket->display_shape = get_shape_for_state(state);
+  for (const bNodeSocket *socket : tree.all_sockets()) {
+    const SocketFieldState &state = field_state_by_socket_id[socket->index_in_tree()];
+    const_cast<bNodeSocket *>(socket)->display_shape = get_shape_for_state(state);
   }
 }
 
-static bool update_field_inferencing(const NodeTreeRef &tree)
+static bool update_field_inferencing(const bNodeTree &tree)
 {
-  bNodeTree &btree = *tree.btree();
+  tree.ensure_topology_cache();
 
   /* Create new inferencing interface for this node group. */
   std::unique_ptr<FieldInferencingInterface> new_inferencing_interface =
       std::make_unique<FieldInferencingInterface>();
-  new_inferencing_interface->inputs.resize(BLI_listbase_count(&btree.inputs),
+  new_inferencing_interface->inputs.resize(BLI_listbase_count(&tree.inputs),
                                            InputSocketFieldType::IsSupported);
-  new_inferencing_interface->outputs.resize(BLI_listbase_count(&btree.outputs),
+  new_inferencing_interface->outputs.resize(BLI_listbase_count(&tree.outputs),
                                             OutputFieldDependency::ForDataSource());
 
   /* Keep track of the state of all sockets. The index into this array is #SocketRef::id(). */
-  Array<SocketFieldState> field_state_by_socket_id(tree.sockets().size());
+  Array<SocketFieldState> field_state_by_socket_id(tree.all_sockets().size());
 
   propagate_data_requirements_from_right_to_left(tree, field_state_by_socket_id);
   determine_group_input_states(tree, *new_inferencing_interface, field_state_by_socket_id);
@@ -575,10 +577,10 @@ static bool update_field_inferencing(const NodeTreeRef &tree)
   update_socket_shapes(tree, field_state_by_socket_id);
 
   /* Update the previous group interface. */
-  const bool group_interface_changed = !btree.runtime->field_inferencing_interface ||
-                                       *btree.runtime->field_inferencing_interface !=
+  const bool group_interface_changed = !tree.runtime->field_inferencing_interface ||
+                                       *tree.runtime->field_inferencing_interface !=
                                            *new_inferencing_interface;
-  btree.runtime->field_inferencing_interface = std::move(new_inferencing_interface);
+  tree.runtime->field_inferencing_interface = std::move(new_inferencing_interface);
 
   return group_interface_changed;
 }
@@ -979,29 +981,22 @@ class NodeTreeMainUpdater {
   {
     TreeUpdateResult result;
 
-    /* Use a #NodeTreeRef to speedup certain queries. It is rebuilt whenever the node tree topology
-     * changes, which typically happens zero or one times during the entire update of the node
-     * tree. */
-    std::unique_ptr<NodeTreeRef> tree_ref;
-    this->ensure_tree_ref(ntree, tree_ref);
-
-    this->update_socket_link_and_use(*tree_ref);
-    this->update_individual_nodes(ntree, tree_ref);
-    this->update_internal_links(ntree, tree_ref);
-    this->update_generic_callback(ntree, tree_ref);
+    this->update_socket_link_and_use(ntree);
+    this->update_individual_nodes(ntree);
+    this->update_internal_links(ntree);
+    this->update_generic_callback(ntree);
     this->remove_unused_previews_when_necessary(ntree);
 
-    this->ensure_tree_ref(ntree, tree_ref);
-    this->propagate_runtime_flags(*tree_ref);
+    this->propagate_runtime_flags(ntree);
     if (ntree.type == NTREE_GEOMETRY) {
-      if (node_field_inferencing::update_field_inferencing(*tree_ref)) {
+      if (node_field_inferencing::update_field_inferencing(ntree)) {
         result.interface_changed = true;
       }
     }
 
-    result.output_changed = this->check_if_output_changed(*tree_ref);
+    result.output_changed = this->check_if_output_changed(ntree);
 
-    this->update_socket_link_and_use(*tree_ref);
+    this->update_socket_link_and_use(ntree);
     this->update_node_levels(ntree);
     this->update_link_validation(ntree);
 
@@ -1021,86 +1016,69 @@ class NodeTreeMainUpdater {
     return result;
   }
 
-  void ensure_tree_ref(bNodeTree &ntree, std::unique_ptr<NodeTreeRef> &tree_ref)
-  {
-    if (!tree_ref) {
-      tree_ref = std::make_unique<NodeTreeRef>(&ntree);
-    }
-  }
-
-  void update_socket_link_and_use(const NodeTreeRef &tree)
+  void update_socket_link_and_use(bNodeTree &tree)
   {
-    for (const InputSocketRef *socket : tree.input_sockets()) {
-      bNodeSocket *bsocket = socket->bsocket();
+    tree.ensure_topology_cache();
+    for (bNodeSocket *socket : tree.all_input_sockets()) {
       if (socket->directly_linked_links().is_empty()) {
-        bsocket->link = nullptr;
+        socket->link = nullptr;
       }
       else {
-        bsocket->link = socket->directly_linked_links()[0]->blink();
+        socket->link = socket->directly_linked_links()[0];
       }
     }
 
     this->update_socket_used_tags(tree);
   }
 
-  void update_socket_used_tags(const NodeTreeRef &tree)
+  void update_socket_used_tags(bNodeTree &tree)
   {
-    for (const SocketRef *socket : tree.sockets()) {
-      bNodeSocket *bsocket = socket->bsocket();
-      bsocket->flag &= ~SOCK_IN_USE;
-      for (const LinkRef *link : socket->directly_linked_links()) {
-        if (!link->is_muted()) {
-          bsocket->flag |= SOCK_IN_USE;
+    tree.ensure_topology_cache();
+    for (bNodeSocket *socket : tree.all_sockets()) {
+      socket->flag &= ~SOCK_IN_USE;
+      for (const bNodeLink *link : socket->directly_linked_links()) {
+        if ((link->flag & NODE_LINK_MUTED) != 0) {
+          socket->flag |= SOCK_IN_USE;
           break;
         }
       }
     }
   }
 
-  void update_individual_nodes(bNodeTree &ntree, std::unique_ptr<NodeTreeRef> &tree_ref)
+  void update_individual_nodes(bNodeTree &ntree)
   {
-    /* Iterate over nodes instead of #NodeTreeRef, because the #tree_ref might be outdated after
-     * some update functions. */
-    LISTBASE_FOREACH (bNode *, bnode, &ntree.nodes) {
-      this->ensure_tree_ref(ntree, tree_ref);
-      const NodeRef &node = *tree_ref->find_node(*bnode);
-      if (this->should_update_individual_node(node)) {
-        const uint32_t old_changed_flag = ntree.runtime->changed_flag;
-        ntree.runtime->changed_flag = NTREE_CHANGED_NOTHING;
-
-        /* This may set #ntree.runtime->changed_flag which is detected below. */
-        this->update_individual_node(node);
-
-        if (ntree.runtime->changed_flag != NTREE_CHANGED_NOTHING) {
-          /* The tree ref is outdated and needs to be rebuilt. Generally, only very few update
-           * functions change the node. Typically zero or one nodes change after an update. */
-          tree_ref.reset();
+    LISTBASE_FOREACH (bNode *, node, &ntree.nodes) {
+      if (this->should_update_individual_node(ntree, *node)) {
+        bNodeType &ntype = *node->typeinfo;
+        if (ntype.group_update_func) {
+          ntype.group_update_func(&ntree, node);
+        }
+        if (ntype.updatefunc) {
+          ntype.updatefunc(&ntree, node);
         }
-        ntree.runtime->changed_flag |= old_changed_flag;
       }
     }
   }
 
-  bool should_update_individual_node(const NodeRef &node)
+  bool should_update_individual_node(const bNodeTree &ntree, const bNode &node)
   {
-    bNodeTree &ntree = *node.btree();
-    bNode &bnode = *node.bnode();
     if (ntree.runtime->changed_flag & NTREE_CHANGED_ANY) {
       return true;
     }
-    if (bnode.runtime->changed_flag & NTREE_CHANGED_NODE_PROPERTY) {
+    if (node.runtime->changed_flag & NTREE_CHANGED_NODE_PROPERTY) {
       return true;
     }
     if (ntree.runtime->changed_flag & NTREE_CHANGED_LINK) {
+      ntree.ensure_topology_cache();
       /* Node groups currently always rebuilt their sockets when they are updated.
        * So avoid calling the update method when no new link was added to it. */
-      if (node.is_group_input_node()) {
-        if (node.outputs().last()->is_directly_linked()) {
+      if (node.type == NODE_GROUP_INPUT) {
+        if (node.output_sockets().last()->is_directly_linked()) {
           return true;
         }
       }
-      else if (node.is_group_output_node()) {
-        if (node.inputs().last()->is_directly_linked()) {
+      else if (node.type == NODE_GROUP_OUTPUT) {
+        if (node.input_sockets().last()->is_directly_linked()) {
           return true;
         }
       }
@@ -1110,95 +1088,76 @@ class NodeTreeMainUpdater {
       }
     }
     if (ntree.runtime->changed_flag & NTREE_CHANGED_INTERFACE) {
-      if (node.is_group_input_node() || node.is_group_output_node()) {
+      if (ELEM(node.type, NODE_GROUP_INPUT, NODE_GROUP_OUTPUT)) {
         return true;
       }
     }
     return false;
   }
 
-  void update_individual_node(const NodeRef &node)
-  {
-    bNodeTree &ntree = *node.btree();
-    bNode &bnode = *node.bnode();
-    bNodeType &ntype = *bnode.typeinfo;
-    if (ntype.group_update_func) {
-      ntype.group_update_func(&ntree, &bnode);
-    }
-    if (ntype.updatefunc) {
-      ntype.updatefunc(&ntree, &bnode);
-    }
-  }
-
-  void update_internal_links(bNodeTree &ntree, std::unique_ptr<NodeTreeRef> &tree_ref)
+  void update_internal_links(bNodeTree &ntree)
   {
-    bool any_internal_links_updated = false;
-    this->ensure_tree_ref(ntree, tree_ref);
-    for (const NodeRef *node : tree_ref->nodes()) {
-      if (!this->should_update_individual_node(*node)) {
+    bke::node_tree_runtime::AllowUsingOutdatedInfo allow_outdated_info{ntree};
+    ntree.ensure_topology_cache();
+    for (bNode *node : ntree.all_nodes()) {
+      if (!this->should_update_individual_node(ntree, *node)) {
         continue;
       }
       /* Find all expected internal links. */
       Vector<std::pair<bNodeSocket *, bNodeSocket *>> expected_internal_links;
-      for (const OutputSocketRef *output_socket : node->outputs()) {
+      for (const bNodeSocket *output_socket : node->output_sockets()) {
         if (!output_socket->is_available()) {
           continue;
         }
         if (!output_socket->is_directly_linked()) {
           continue;
         }
-        if (output_socket->bsocket()->flag & SOCK_NO_INTERNAL_LINK) {
+        if (output_socket->flag & SOCK_NO_INTERNAL_LINK) {
           continue;
         }
-        const InputSocketRef *input_socket = this->find_internally_linked_input(output_socket);
+        const bNodeSocket *input_socket = this->find_internally_linked_input(output_socket);
         if (input_socket != nullptr) {
-          expected_internal_links.append({input_socket->bsocket(), output_socket->bsocket()});
+          expected_internal_links.append(
+              {const_cast<bNodeSocket *>(input_socket), const_cast<bNodeSocket *>(output_socket)});
         }
       }
-      /* rebuilt internal links if they have changed. */
-      if (node->internal_links().size() != expected_internal_links.size()) {
-        this->update_internal_links_in_node(ntree, *node->bnode(), expected_internal_links);
-        any_internal_links_updated = true;
+      /* Rebuilt internal links if they have changed. */
+      if (node->internal_links_span().size() != expected_internal_links.size()) {
+        this->update_internal_links_in_node(ntree, *node, expected_internal_links);
       }
       else {
         for (auto &item : expected_internal_links) {
           const bNodeSocket *from_socket = item.first;
           const bNodeSocket *to_socket = item.second;
           bool found = false;
-          for (const InternalLinkRef *internal_link : node->internal_links()) {
-            if (from_socket == internal_link->from().bsocket() &&
-                to_socket == internal_link->to().bsocket()) {
+          for (const bNodeLink *internal_link : node->internal_links_span()) {
+            if (from_socket == internal_link->fromsock && to_socket == internal_link->tosock) {
               found = true;
             }
           }
           if (!found) {
-            this->update_internal_links_in_node(ntree, *node->bnode(), expected_internal_links);
-            any_internal_links_updated = true;
+            this->update_internal_links_in_node(ntree, *node, expected_internal_links);
             break;
           }
         }
       }
     }
-
-    if (any_internal_links_updated) {
-      tree_ref.reset();
-    }
   }
 
-  const InputSocketRef *find_internally_linked_input(const OutputSocketRef *output_socket)
+  const bNodeSocket *find_internally_linked_input(const bNodeSocket *output_socket)
   {
-    const InputSocketRef *selected_socket = nullptr;
+    const bNodeSocket *selected_socket = nullptr;
     int selected_priority = -1;
     bool selected_is_linked = false;
-    for (const InputSocketRef *input_socket : output_socket->node().inputs()) {
+    for (const bNodeSocket *input_socket : output_socket->owner_node().input_sockets()) {
       if (!input_socket->is_available()) {
         continue;
       }
-      if (input_socket->bsocket()->flag & SOCK_NO_INTERNAL_LINK) {
+      if (input_socket->flag & SOCK_NO_INTERNAL_LINK) {
         continue;
       }
-      const int priority = get_internal_link_type_priority(input_socket->bsocket()->typeinfo,
-                                                           output_socket->bsocket()->typeinfo);
+      const int priority = get_internal_link_type_priority(input_socket->typeinfo,
+                                                           output_socket->typeinfo);
       if (priority < 0) {
         continue;
       }
@@ -1233,23 +1192,12 @@ class NodeTreeMainUpdater {
     BKE_ntree_update_tag_node_internal_link(&ntree, &node);
   }
 
-  void update_generic_callback(bNodeTree &ntree, std::unique_ptr<NodeTreeRef> &tree_ref)
+  void update_generic_callback(bNodeTree &ntree)
   {
     if (ntree.typeinfo->update == nullptr) {
       return;
     }
-
-    /* Reset the changed_flag to allow detecting when the update callback changed the node tree. */
-    const uint32_t old_changed_flag = ntree.runtime->changed_flag;
-    ntree.runtime->changed_flag = NTREE_CHANGED_NOTHING;
-
     ntree.typeinfo->update(&ntree);
-
-    if (ntree.runtime->changed_flag != NTREE_CHANGED_NOTHING) {
-      /* The tree ref is outdated and needs to be rebuilt. */
-      tree_ref.reset();
-    }
-    ntree.runtime->changed_flag |= old_changed_flag;
   }
 
   void remove_unused_previews_when_necessary(bNodeTree &ntree)
@@ -1264,25 +1212,26 @@ class NodeTreeMainUpdater {
     BKE_node_preview_remove_unused(&ntree);
   }
 
-  void propagate_runtime_flags(const NodeTreeRef &tree_ref)
+  void propagate_runtime_flags(const bNodeTree &ntree)
   {
-    bNodeTree &ntree = *tree_ref.btree();
+    ntree.ensure_topology_cache();
+
     ntree.runtime->runtime_flag = 0;
     if (ntree.type != NTREE_SHADER) {
       return;
     }
 
     /* Check if a used node group has an animated image. */
-    for (const NodeRef *group_node : tree_ref.nodes_by_type("ShaderNodeGroup")) {
-      const bNodeTree *group = reinterpret_cast<bNodeTree *>(group_node->bnode()->id);
+    for (const bNode *group_node : ntree.nodes_by_type("ShaderNodeGroup")) {
+      const bNodeTree *group = reinterpret_cast<bNodeTree *>(group_node->id);
       if (group != nullptr) {
         ntree.runtime->runtime_flag |= group->runtime->runtime_flag;
       }
     }
     /* Check if the tree itself has an animated image. */
     for (const StringRefNull idname : {"ShaderNodeTexImage", "ShaderNodeTexEnvironment"}) {
-      for (const NodeRef *node : tree_ref.nodes_by_type(idname)) {
-        Image *image = reinterpret_cast<Image *>(node->bnode()->id);
+      for (const bNode *node : ntree.nodes_by_type(idname)) {
+        Image *image = reinterpret_cast<Image *>(node->id);
         if (image != nullptr && BKE_image_is_animated(image)) {
           ntree.runtime->runtime_flag |= NTREE_RUNTIME_FLAG_HAS_IMAGE_ANIMATION;
           break;
@@ -1294,7 +1243,7 @@ class NodeTreeMainUpdater {
                                        "ShaderNodeOutputLight",
                                        "ShaderNodeOutputWorld",
                                        "ShaderNodeOutputAOV"}) {
-      const Span<const NodeRef *> nodes = tree_ref.nodes_by_type(idname);
+      const Span<const bNode *> nodes = ntree.nodes_by_type(idname);
       if (!nodes.is_empty()) {
         ntree.runtime->runtime_flag |= NTREE_RUNTIME_FLAG_HAS_MATERIAL_OUTPUT;
         break;
@@ -1325,20 +1274,20 @@ class NodeTreeMainUpdater {
     }
   }
 
-  bool check_if_output_changed(const NodeTreeRef &tree)
+  bool check_if_output_changed(const bNodeTree &tree)
   {
-    bNodeTree &btree = *tree.btree();
+    tree.ensure_topology_cache();
 
     /* Compute a hash that represents the node topology connected to the output. This always has to
      * be updated even if it is not used to detect changes right now. Otherwise
      * #btree.runtime.output_topology_hash will go out of date. */
-    const Vector<const SocketRef *> tree_output_sockets = this->find_output_sockets(tree);
-    const uint32_t old_topology_hash = btree.runtime->output_topology_hash;
+    const Vector<const bNodeSocket *> tree_output_sockets = this->find_output_sockets(tree);
+    const uint32_t old_topology_hash = tree.runtime->output_topology_hash;
     const uint32_t new_topology_hash = this->get_combined_socket_topology_hash(
         tree, tree_output_sockets);
-    btree.runtime->output_topology_hash = new_topology_hash;
+    tree.runtime->output_topology_hash = new_topology_hash;
 
-    if (const AnimData *adt = BKE_animdata_from_id(&btree.id)) {
+    if (const AnimData *adt = BKE_animdata_from_id(&tree.id)) {
       /* Drivers may copy values in the node tree around arbitrarily and may cause the output to
        * change even if it wouldn't without drivers. Only some special drivers like `frame/5` can
        * be used without causing updates all the time currently. In the future we could try to
@@ -1360,7 +1309,7 @@ class NodeTreeMainUpdater {
       }
     }
 
-    if (btree.runtime->changed_flag & NTREE_CHANGED_ANY) {
+    if (tree.runtime->changed_flag & NTREE_CHANGED_ANY) {
       return true;
     }
 
@@ -1369,8 +1318,8 @@ class NodeTreeMainUpdater {
     }
 
     /* The topology hash can only be used when only topology-changing operations have been done. */
-    if (btree.runtime->changed_flag ==
-        (btree.runtime->changed_flag & (NTREE_CHANGED_LINK | NTREE_CHANGED_REMOVED_NODE))) {
+    if (tree.runtime->changed_flag ==
+        (tree.runtime->changed_flag & (NTREE_CHANGED_LINK | NTREE_CHANGED_REMOVED_NODE))) {
       if (old_topology_hash == new_topology_hash) {
         return false;
       }
@@ -1383,15 +1332,15 @@ class NodeTreeMainUpdater {
     return true;
   }
 
-  Vector<const SocketRef *> find_output_sockets(const NodeTreeRef &tree)
+  Vector<const bNodeSocket *> find_output_sockets(const bNodeTree &tree)
   {
-    Vector<const SocketRef *> sockets;
-    for (const NodeRef *node : tree.nodes()) {
+    Vector<const bNodeSocket *> sockets;
+    for (const bNode *node : tree.all_nodes()) {
       if (!this->is_output_node(*node)) {
         continue;
       }
-      for (const InputSocketRef *socket : node->inputs()) {
-        if (socket->idname() != "NodeSocketVirtual") {
+      for (const bNodeSocket *socket : node->input_sockets()) {
+        if (!STREQ(socket->idname, "NodeSocketVirtual")) {
           sockets.append(socket);
         }
       }
@@ -1399,18 +1348,17 @@ class NodeTreeMainUpdater {
     return sockets;
   }
 
-  bool is_output_node(const NodeRef &node) const
+  bool is_output_node(const bNode &node) const
   {
-    const bNode &bnode = *node.bnode();
-    if (bnode.typeinfo->nclass == NODE_CLASS_OUTPUT) {
+    if (node.typeinfo->nclass == NODE_CLASS_OUTPUT) {
       return true;
     }
-    if (bnode.type == NODE_GROUP_OUTPUT) {
+    if (node.type == NODE_GROUP_OUTPUT) {
       return true;
     }
     /* Assume node groups without output sockets are outputs. */
-    if (bnode.type == NODE_GROUP) {
-      const bNodeTree *node_group = reinterpret_cast<const bNodeTree *>(bnode.id);
+    if (node.type == NODE_GROUP) {
+      const bNodeTree *node_group = reinterpret_cast<const bNodeTree *>(node.id);
       if (node_group != nullptr &&
           node_group->runtime->runtime_flag & NTREE_RUNTIME_FLAG_HAS_MATERIAL_OUTPUT) {
         return true;
@@ -1423,10 +1371,10 @@ class NodeTreeMainUpdater {
    * Computes a hash that changes when the node tree topology connected to an output node changes.
    * Adding reroutes does not have an effect on the hash.
    */
-  uint32_t get_combined_socket_topology_hash(const NodeTreeRef &tree,
-                                             Span<const SocketRef *> sockets)
+  uint32_t get_combined_socket_topology_hash(const bNodeTree &tree,
+                                             Span<const bNodeSocket *> sockets)
   {
-    if (tree.has_link_cycles()) {
+    if (tree.has_link_cycle()) {
       /* Return dummy value when the link has any cycles. The algorithm below could be improved to
        * handle cycles more gracefully. */
       return 0;
@@ -1439,29 +1387,28 @@ class NodeTreeMainUpdater {
     return combined_hash;
   }
 
-  Array<uint32_t> get_socket_topology_hashes(const NodeTreeRef &tree,
-                                             Span<const SocketRef *> sockets)
+  Array<uint32_t> get_socket_topology_hashes(const bNodeTree &tree,
+                                             Span<const bNodeSocket *> sockets)
   {
-    BLI_assert(!tree.has_link_cycles());
-    Array<std::optional<uint32_t>> hash_by_socket_id(tree.sockets().size());
-    Stack<const SocketRef *> sockets_to_check = sockets;
+    BLI_assert(!tree.has_link_cycle());
+    Array<std::optional<uint32_t>> hash_by_socket_id(tree.all_sockets().size());
+    Stack<const bNodeSocket *> sockets_to_check = sockets;
 
     while (!sockets_to_check.is_empty()) {
-      const SocketRef &in_out_socket = *sockets_to_check.peek();
-      const NodeRef &node = in_out_socket.node();
+      const bNodeSocket &socket = *sockets_to_check.peek();
+      const bNode &node = socket.owner_node();
 
-      if (hash_by_socket_id[in_out_socket.id()].has_value()) {
+      if (hash_by_socket_id[socket.index_in_tree()].has_value()) {
         sockets_to_check.pop();
         /* Socket is handled already. */
         continue;
       }
 
-      if (in_out_socket.is_input()) {
+      if (socket.is_input()) {
         /* For input sockets, first compute the hashes of all linked sockets. */
-        const InputSocketRef &socket = in_out_socket.as_input();
         bool all_origins_computed = true;
-        for (const OutputSocketRef *origin_socket : socket.logically_linked_sockets()) {
-          if (!hash_by_socket_id[origin_socket->id()].has_value()) {
+        for (const bNodeSocket *origin_socket : socket.logically_linked_sockets()) {
+          if (!hash_by_socket_id[origin_socket->index_in_tree()].has_value()) {
             sockets_to_check.push(origin_socket);
             all_origins_computed = false;
           }
@@ -1471,22 +1418,21 @@ class NodeTreeMainUpdater {
         }
         /* When the hashes for the linked sockets are ready, combine them into a hash for the input
          * socket. */
-        const uint64_t socket_ptr = (uintptr_t)socket.bsocket();
+        const uint64_t socket_ptr = (uintptr_t)&socket;
         uint32_t socket_hash = noise::hash(socket_ptr, socket_ptr >> 32);
-        for (const OutputSocketRef *origin_socket : socket.logically_linked_sockets()) {
-          const uint32_t origin_socket_hash = *hash_by_socket_id[origin_socket->id()];
+        for (const bNodeSocket *origin_socket : socket.logically_linked_sockets()) {
+          const uint32_t origin_socket_hash = *hash_by_socket_id[origin_socket->index_in_tree()];
           socket_hash = noise::hash(socket_hash, origin_socket_hash);
         }
-        hash_by_socket_id[socket.id()] = socket_hash;
+        hash_by_socket_id[socket.index_in_tree()] = socket_hash;
         sockets_to_check.pop();
       }
       else {
         /* For output sockets, first compute the hashes of all available input sockets. */
-        const OutputSocketRef &socket = in_out_socket.as_output();
         bool all_available_inputs_computed = true;
-        for (const InputSocketRef *input_socket : node.inputs()) {
+        for (const bNodeSocket *input_socket : node.input_sockets()) {
           if (input_socket->is_available()) {
-            if (!hash_by_socket_id[input_socket->id()].has_value()) {
+            if (!hash_by_socket_id[input_socket->index_in_tree()].has_value()) {
               sockets_to_check.push(input_socket);
               all_available_inputs_computed = false;
             }
@@ -1497,25 +1443,25 @@ class NodeTreeMainUpdater {
         }
         /* When all input socket hashes have been computed, combine them into a hash for the output
          * socket. */
-        const uint64_t socket_ptr = (uintptr_t)socket.bsocket();
+        const uint64_t socket_ptr = (uintptr_t)&socket;
         uint32_t socket_hash = noise::hash(socket_ptr, socket_ptr >> 32);
-        for (const InputSocketRef *input_socket : node.inputs()) {
+        for (const bNodeSocket *input_socket : node.input_sockets()) {
           if (input_socket->is_available()) {
-            const uint32_t input_socket_hash = *hash_by_socket_id[input_socket->id()];
+            const uint32_t input_socket_hash = *hash_by_socket_id[input_socket->index_in_tree()];
             socket_hash = noise::hash(socket_hash, input_socket_hash);
           }
         }
         /* The Image Texture node has a special case. The behavior of the color output changes
          * depending on whether the Alpha output is linked. */
-        if (node.bnode()->type == SH_NODE_TEX_IMAGE && socket.index() == 0) {
-          BLI_assert(socket.name() == "Color");
-          const OutputSocketRef &alpha_socket = node.output(1);
-          BLI_assert(alpha_socket.name() == "Alpha");
+        if (node.type == SH_NODE_TEX_IMAGE && socket.index() == 0) {
+          BLI_assert(STREQ(socket.name, "Color"));
+          const bNodeSocket &alpha_socket = node.output_socket(1);
+          BLI_assert(STREQ(alpha_socket.name, "Alpha"));
           if (alpha_socket.is_directly_linked()) {
             socket_hash = noise::hash(socket_hash);
           }
         }
-        hash_by_socket_id[socket.id()] = socket_hash;
+        hash_by_socket_id[socket.index_in_tree()] = socket_hash;
         sockets_to_check.pop();
       }
     }
@@ -1523,7 +1469,7 @@ class NodeTreeMainUpdater {
     /* Create output array. */
     Array<uint32_t> hashes(sockets.size());
     for (const int i : sockets.index_range()) {
-      hashes[i] = *hash_by_socket_id[sockets[i]->id()];
+      hashes[i] = *hash_by_socket_id[sockets[i]->index_in_tree()];
     }
     return hashes;
   }
@@ -1532,37 +1478,34 @@ class NodeTreeMainUpdater {
    * Returns true when any of the provided sockets changed its values. A change is detected by
    * checking the #changed_flag on connected sockets and nodes.
    */
-  bool check_if_socket_outputs_changed_based_on_flags(const NodeTreeRef &tree,
-                                                      Span<const SocketRef *> sockets)
+  bool check_if_socket_outputs_changed_based_on_flags(const bNodeTree &tree,
+                                                      Span<const bNodeSocket *> sockets)
   {
     /* Avoid visiting the same socket twice when multiple links point to the same socket. */
-    Array<bool> pushed_by_socket_id(tree.sockets().size(), false);
-    Stack<const SocketRef *> sockets_to_check = sockets;
+    Array<bool> pushed_by_socket_id(tree.all_sockets().size(), false);
+    Stack<const bNodeSocket *> sockets_to_check = sockets;
 
-    for (const SocketRef *socket : sockets) {
-      pushed_by_socket_id[socket->id()] = true;
+    for (const bNodeSocket *socket : sockets) {
+      pushed_by_socket_id[socket->index_in_tree()] = true;
     }
 
     while (!sockets_to_check.is_empty()) {
-      const SocketRef &in_out_socket = *sockets_to_check.pop();
-      const NodeRef &node = in_out_socket.node();
-      const bNode &bnode = *node.bnode();
-      const bNodeSocket &bsocket = *in_out_socket.bsocket();
-      if (bsocket.runtime->changed_flag != NTREE_CHANGED_NOTHING) {
+      const bNodeSocket &socket = *sockets_to_check.pop();
+      const bNode &node = socket.owner_node();
+      if (socket.runtime->changed_flag != NTREE_CHANGED_NOTHING) {
         return true;
       }
-      if (bnode.runtime->changed_flag != NTREE_CHANGED_NOTHING) {
-        const bool only_unused_internal_link_changed = (bnode.flag & NODE_MUTED) == 0 &&
-                                                       bnode.runtime->changed_flag ==
+      if (node.runtime->changed_flag != NTREE_CHANGED_NOTHING) {
+        const bool only_unused_internal_link_changed = !node.is_muted() &&
+                                                       node.runtime->changed_flag ==
                                                            NTREE_CHANGED_INTERNAL_LINK;
         if (!only_unused_internal_link_changed) {
           return true;
         }
       }
-      if (in_out_socket.is_input()) {
-        const InputSocketRef &socket = in_out_socket.as_input();
-        for (const OutputSocketRef *origin_socket : socket.logically_linked_sockets()) {
-          bool &pushed = pushed_by_socket_id[origin_socket->id()];
+      if (socket.is_input()) {
+        for (const bNodeSocket *origin_socket : socket.logically_linked_sockets()) {
+          bool &pushed = pushed_by_socket_id[origin_socket->index_in_tree()];
           if (!pushed) {
             sockets_to_check.push(origin_socket);
             pushed = true;
@@ -1570,10 +1513,9 @@ class NodeTreeMainUpdater {
         }
       }
       else {
-        const OutputSocketRef &socket = in_out_socket.as_output();
-        for (const InputSocketRef *input_socket : node.inputs()) {
+        for (const bNodeSocket *input_socket : node.input_sockets()) {
           if (input_socket->is_available()) {
-            bool &pushed = pushed_by_socket_id[input_socket->id()];
+            bool &pushed = pushed_by_socket_id[input_socket->index_in_tree()];
             if (!pushed) {
               sockets_to_check.push(input_socket);
               pushed = true;
@@ -1582,11 +1524,11 @@ class NodeTreeMainUpdater {
         }
         /* The Normal node has a special case, because the value stored in the first output socket
          * is used as input in the node. */
-        if (bnode.type == SH_NODE_NORMAL && socket.index() == 1) {
-          BLI_assert(socket.name() == "Dot");
-          const OutputSocketRef &normal_output = node.output(0);
-          BLI_assert(normal_output.name() == "Normal");
-          bool &pushed = pushed_by_socket_id[normal_output.id()];
+        if (node.type == SH_NODE_NORMAL && socket.index() == 1) {
+          BLI_assert(STREQ(socket.name, "Dot"));
+          const bNodeSocket &normal_output = node.output_socket(0);
+          BLI_assert(STREQ(normal_output.name, "Normal"));
+          bool &pushed = pushed_by_socket_id[normal_output.index_in_tree()];
           if (!pushed) {
             sockets_to_check.push(&normal_output);
             pushed = true;