Geometry Nodes: support tracing bundles and closures through closure calls

This makes the "sync sockets" operation work on more cases than before.

Pull Request: https://projects.blender.org/blender/blender/pulls/141014

source/blender/editors/space_node/space_node.cc

@@ -469,6 +469,12 @@ static std::optional<const ComputeContext *> compute_context_for_tree_path(
   return current;
 }
 
+static Vector<nodes::SocketInContext> find_origin_sockets_through_contexts(
+    nodes::SocketInContext start_socket,
+    bke::ComputeContextCache &compute_context_cache,
+    StringRef query_node_idname,
+    bool find_all);
+
 static Vector<nodes::SocketInContext> find_target_sockets_through_contexts(
     const nodes::SocketInContext start_socket,
     bke::ComputeContextCache &compute_context_cache,
@@ -569,6 +575,70 @@ static Vector<nodes::SocketInContext> find_target_sockets_through_contexts(
         }
         continue;
       }
+      if (node->is_type("GeometryNodeClosureOutput")) {
+        const auto &closure_storage = *static_cast<const NodeGeometryClosureOutput *>(
+            node->storage);
+        const nodes::SocketInterfaceKey key(
+            closure_storage.output_items.items[socket->index()].name);
+        const Vector<nodes::SocketInContext> target_sockets = find_target_sockets_through_contexts(
+            node.output_socket(0), compute_context_cache, "GeometryNodeEvaluateClosure", true);
+        for (const auto &target_socket : target_sockets) {
+          const nodes::NodeInContext evaluate_node = target_socket.owner_node();
+          const auto &evaluate_storage = *static_cast<const NodeGeometryEvaluateClosure *>(
+              evaluate_node->storage);
+          for (const int i : IndexRange(evaluate_storage.output_items.items_num)) {
+            const NodeGeometryEvaluateClosureOutputItem &item =
+                evaluate_storage.output_items.items[i];
+            if (key.matches(nodes::SocketInterfaceKey(item.name))) {
+              add_if_new(evaluate_node.output_socket(i), bundle_path);
+            }
+          }
+        }
+        continue;
+      }
+      if (node->is_type("GeometryNodeEvaluateClosure")) {
+        if (socket->index() == 0) {
+          continue;
+        }
+        const auto &evaluate_storage = *static_cast<const NodeGeometryEvaluateClosure *>(
+            node->storage);
+        const nodes::SocketInterfaceKey key(
+            evaluate_storage.input_items.items[socket->index() - 1].name);
+        const Vector<nodes::SocketInContext> origin_sockets = find_origin_sockets_through_contexts(
+            node.input_socket(0), compute_context_cache, "GeometryNodeClosureOutput", true);
+        for (const nodes::SocketInContext origin_socket : origin_sockets) {
+          const bNodeTree &closure_tree = origin_socket->owner_tree();
+          const bke::bNodeTreeZones *closure_tree_zones = closure_tree.zones();
+          if (!closure_tree_zones) {
+            continue;
+          }
+          const auto &closure_output_node = origin_socket.owner_node();
+          const bke::bNodeTreeZone *closure_zone = closure_tree_zones->get_zone_by_node(
+              closure_output_node->identifier);
+          if (!closure_zone) {
+            continue;
+          }
+          const bNode *closure_input_node = closure_zone->input_node();
+          if (!closure_input_node) {
+            continue;
+          }
+          const ComputeContext &closure_context = compute_context_cache.for_evaluate_closure(
+              node.context,
+              node->identifier,
+              &node->owner_tree(),
+              nodes::ClosureSourceLocation{
+                  &closure_tree, closure_output_node->identifier, origin_socket.context_hash()});
+          const auto &closure_output_storage = *static_cast<const NodeGeometryClosureOutput *>(
+              closure_output_node->storage);
+          for (const int i : IndexRange(closure_output_storage.input_items.items_num)) {
+            const NodeGeometryClosureInputItem &item = closure_output_storage.input_items.items[i];
+            if (key.matches(nodes::SocketInterfaceKey(item.name))) {
+              add_if_new({&closure_context, &closure_input_node->output_socket(i)}, bundle_path);
+            }
+          }
+        }
+        continue;
+      }
     }
     else {
       const bke::bNodeTreeZones *zones = node->owner_tree().zones();
@@ -728,6 +798,64 @@ static Vector<nodes::SocketInContext> find_origin_sockets_through_contexts(
         }
         continue;
       }
+      if (node->is_type("GeometryNodeEvaluateClosure")) {
+        const auto &evaluate_storage = *static_cast<const NodeGeometryEvaluateClosure *>(
+            node->storage);
+        const nodes::SocketInterfaceKey key(
+            evaluate_storage.output_items.items[socket->index()].name);
+        const Vector<nodes::SocketInContext> origin_sockets = find_origin_sockets_through_contexts(
+            node.input_socket(0), compute_context_cache, "GeometryNodeClosureOutput", true);
+        for (const nodes::SocketInContext origin_socket : origin_sockets) {
+          const bNodeTree &closure_tree = origin_socket->owner_tree();
+          const nodes::NodeInContext closure_output_node = origin_socket.owner_node();
+          const auto &closure_storage = *static_cast<const NodeGeometryClosureOutput *>(
+              closure_output_node->storage);
+          const ComputeContext &closure_context = compute_context_cache.for_evaluate_closure(
+              node.context,
+              node->identifier,
+              &node->owner_tree(),
+              nodes::ClosureSourceLocation{
+                  &closure_tree, closure_output_node->identifier, origin_socket.context_hash()});
+          for (const int i : IndexRange(closure_storage.output_items.items_num)) {
+            const NodeGeometryClosureOutputItem &item = closure_storage.output_items.items[i];
+            if (key.matches(nodes::SocketInterfaceKey(item.name))) {
+              add_if_new({&closure_context, &closure_output_node->input_socket(i)}, bundle_path);
+            }
+          }
+        }
+        continue;
+      }
+      if (node->is_type("GeometryNodeClosureInput")) {
+        const auto &input_storage = *static_cast<const NodeGeometryClosureInput *>(node->storage);
+        const bNode *closure_output_node = node->owner_tree().node_by_id(
+            input_storage.output_node_id);
+        if (!closure_output_node) {
+          continue;
+        }
+        const auto &output_storage = *static_cast<const NodeGeometryClosureOutput *>(
+            closure_output_node->storage);
+        const nodes::SocketInterfaceKey key(
+            output_storage.input_items.items[socket->index()].name);
+        const bNodeSocket &closure_output_socket = closure_output_node->output_socket(0);
+        const Vector<nodes::SocketInContext> target_sockets = find_target_sockets_through_contexts(
+            {socket.context, &closure_output_socket},
+            compute_context_cache,
+            "GeometryNodeEvaluateClosure",
+            true);
+        for (const nodes::SocketInContext &target_socket : target_sockets) {
+          const nodes::NodeInContext target_node = target_socket.owner_node();
+          const auto &evaluate_storage = *static_cast<const NodeGeometryEvaluateClosure *>(
+              target_node.node->storage);
+          for (const int i : IndexRange(evaluate_storage.input_items.items_num)) {
+            const NodeGeometryEvaluateClosureInputItem &item =
+                evaluate_storage.input_items.items[i];
+            if (key.matches(nodes::SocketInterfaceKey(item.name))) {
+              add_if_new(target_node.input_socket(i + 1), bundle_path);
+            }
+          }
+        }
+        continue;
+      }
       if (node->is_type("GeometryNodeCombineBundle")) {
         if (bundle_path.is_empty()) {
           continue;