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

1 files changed, 118 insertions, 145 deletions

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