griefith/test
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
628a10a9fb587cd323fe399718ae5f00413c1921
test/source/blender/nodes/intern/trace_values.cc
Jacques Lucke 77563587cd Geometry Nodes: trace bundles/closures through repeat zone, simulation zone and bake node 
This adds support for tracing bundle and closure structures through repeat zones,
simulation zones and bake nodes. Previously, syncing through these nodes just
didn't work.

Pull Request: https://projects.blender.org/blender/blender/pulls/143860
2025-08-02 22:28:17 +02:00

642 lines
26 KiB
C++
Raw Blame History

/* SPDX-FileCopyrightText: 2025 Blender Authors
*
* SPDX-License-Identifier: GPL-2.0-or-later */
#include "BLI_stack.hh"
#include "NOD_geometry_nodes_bundle_signature.hh"
#include "NOD_geometry_nodes_closure_location.hh"
#include "NOD_geometry_nodes_closure_signature.hh"
#include "NOD_node_in_compute_context.hh"
#include "NOD_socket_declarations.hh"
#include "NOD_trace_values.hh"
#include "BKE_compute_context_cache.hh"
#include "BKE_node_tree_zones.hh"
#include "ED_node.hh"
namespace blender::nodes {
static bool is_evaluate_closure_node_input(const SocketInContext &socket)
{
return socket->is_input() && socket->index() == 0 &&
socket.owner_node()->is_type("NodeEvaluateClosure");
}
static bool is_closure_zone_output_socket(const SocketInContext &socket)
{
return socket->owner_node().is_type("NodeClosureOutput") && socket->is_output();
}
static bool use_link_for_tracing(const bNodeLink &link)
{
if (!link.is_used()) {
return false;
}
const bNodeTree &tree = link.fromnode->owner_tree();
if (tree.typeinfo->validate_link &&
!tree.typeinfo->validate_link(eNodeSocketDatatype(link.fromsock->type),
eNodeSocketDatatype(link.tosock->type)))
{
return false;
}
return true;
}
static Vector<SocketInContext> find_origin_sockets_through_contexts(
SocketInContext start_socket,
bke::ComputeContextCache &compute_context_cache,
FunctionRef<bool(const SocketInContext &)> handle_possible_origin_socket_fn,
bool find_all);
static Vector<SocketInContext> find_target_sockets_through_contexts(
const SocketInContext start_socket,
bke::ComputeContextCache &compute_context_cache,
const FunctionRef<bool(const SocketInContext &)> handle_possible_target_socket_fn,
const bool find_all)
{
using BundlePath = Vector<std::string, 0>;
struct SocketToCheck {
SocketInContext socket;
BundlePath bundle_path;
};
Stack<SocketToCheck> sockets_to_check;
Set<SocketInContext> added_sockets;
auto add_if_new = [&](const SocketInContext &socket, BundlePath bundle_path) {
if (added_sockets.add(socket)) {
sockets_to_check.push({socket, std::move(bundle_path)});
}
};
add_if_new(start_socket, {});
VectorSet<SocketInContext> found_targets;
while (!sockets_to_check.is_empty()) {
const SocketToCheck socket_to_check = sockets_to_check.pop();
const SocketInContext socket = socket_to_check.socket;
const BundlePath &bundle_path = socket_to_check.bundle_path;
const NodeInContext &node = socket.owner_node();
if (socket->is_input()) {
if (node->is_muted()) {
for (const bNodeLink &link : node->internal_links()) {
if (link.fromsock == socket.socket) {
add_if_new({socket.context, link.tosock}, bundle_path);
}
}
continue;
}
if (bundle_path.is_empty() && handle_possible_target_socket_fn(socket)) {
found_targets.add(socket);
if (!find_all) {
break;
}
continue;
}
if (node->is_reroute()) {
add_if_new(node.output_socket(0), bundle_path);
continue;
}
if (node->is_group()) {
if (const bNodeTree *group = reinterpret_cast<const bNodeTree *>(node->id)) {
group->ensure_topology_cache();
const ComputeContext &group_compute_context = compute_context_cache.for_group_node(
socket.context, node->identifier, &node->owner_tree());
for (const bNode *input_node : group->group_input_nodes()) {
if (const bNodeSocket *group_input_socket = input_node->output_by_identifier(
socket->identifier))
{
if (group_input_socket->is_directly_linked()) {
add_if_new({&group_compute_context, group_input_socket}, bundle_path);
}
}
}
}
continue;
}
if (node->is_group_output()) {
if (const auto *group_context = dynamic_cast<const bke::GroupNodeComputeContext *>(
socket.context))
{
const bNodeTree *caller_group = group_context->tree();
const bNode *caller_group_node = group_context->node();
if (caller_group && caller_group_node) {
caller_group->ensure_topology_cache();
if (const bNodeSocket *output_socket = caller_group_node->output_by_identifier(
socket->identifier))
{
add_if_new({group_context->parent(), output_socket}, bundle_path);
}
}
}
continue;
}
if (node->is_type("NodeCombineBundle")) {
const auto &storage = *static_cast<const NodeCombineBundle *>(node->storage);
BundlePath new_bundle_path = bundle_path;
new_bundle_path.append(storage.items[socket->index()].name);
add_if_new(node.output_socket(0), std::move(new_bundle_path));
continue;
}
if (node->is_type("NodeSeparateBundle")) {
if (bundle_path.is_empty()) {
continue;
}
const StringRef last_key = bundle_path.last();
const auto &storage = *static_cast<const NodeSeparateBundle *>(node->storage);
for (const int output_i : IndexRange(storage.items_num)) {
if (last_key == storage.items[output_i].name) {
add_if_new(node.output_socket(output_i), bundle_path.as_span().drop_back(1));
}
}
continue;
}
if (node->is_type("NodeClosureOutput")) {
const auto &closure_storage = *static_cast<const NodeClosureOutput *>(node->storage);
const StringRef key = closure_storage.output_items.items[socket->index()].name;
const Vector<SocketInContext> target_sockets = find_target_sockets_through_contexts(
node.output_socket(0), compute_context_cache, is_evaluate_closure_node_input, true);
for (const auto &target_socket : target_sockets) {
const NodeInContext evaluate_node = target_socket.owner_node();
const auto &evaluate_storage = *static_cast<const NodeEvaluateClosure *>(
evaluate_node->storage);
for (const int i : IndexRange(evaluate_storage.output_items.items_num)) {
const NodeEvaluateClosureOutputItem &item = evaluate_storage.output_items.items[i];
if (key == item.name) {
add_if_new(evaluate_node.output_socket(i), bundle_path);
}
}
}
continue;
}
if (node->is_type("NodeEvaluateClosure")) {
if (socket->index() == 0) {
continue;
}
const auto &evaluate_storage = *static_cast<const NodeEvaluateClosure *>(node->storage);
const StringRef key = evaluate_storage.input_items.items[socket->index() - 1].name;
const Vector<SocketInContext> origin_sockets = find_origin_sockets_through_contexts(
node.input_socket(0), compute_context_cache, is_closure_zone_output_socket, true);
for (const 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(),
ClosureSourceLocation{
&closure_tree, closure_output_node->identifier, origin_socket.context_hash()});
const auto &closure_output_storage = *static_cast<const NodeClosureOutput *>(
closure_output_node->storage);
for (const int i : IndexRange(closure_output_storage.input_items.items_num)) {
const NodeClosureInputItem &item = closure_output_storage.input_items.items[i];
if (key == item.name) {
add_if_new({&closure_context, &closure_input_node->output_socket(i)}, bundle_path);
}
}
}
continue;
}
if (node->is_type("GeometryNodeSimulationInput")) {
const ComputeContext &simulation_compute_context =
compute_context_cache.for_simulation_zone(socket.context, *node);
add_if_new({&simulation_compute_context, &node->output_socket(socket->index() + 1)},
bundle_path);
continue;
}
if (node->is_type("GeometryNodeSimulationOutput")) {
const int output_index = socket->index();
if (output_index >= 1) {
BLI_assert(dynamic_cast<const bke::SimulationZoneComputeContext *>(socket.context));
add_if_new({socket.context->parent(), &node->output_socket(output_index - 1)},
bundle_path);
}
continue;
}
if (node->is_type("GeometryNodeRepeatInput")) {
const int index = socket->index();
if (index >= 1) {
const ComputeContext &repeat_compute_context = compute_context_cache.for_repeat_zone(
socket.context, *node, 0);
add_if_new({&repeat_compute_context, &node->output_socket(index)}, bundle_path);
const auto &storage = *static_cast<NodeGeometryRepeatInput *>(node->storage);
if (const bNode *repeat_output_node = node->owner_tree().node_by_id(
storage.output_node_id))
{
add_if_new({socket.context, &repeat_output_node->output_socket(index - 1)},
bundle_path);
}
}
continue;
}
if (node->is_type("GeometryNodeRepeatOutput")) {
BLI_assert(dynamic_cast<const bke::RepeatZoneComputeContext *>(socket.context));
add_if_new({socket.context->parent(), &node->output_socket(socket->index())}, bundle_path);
continue;
}
for (const bNodeSocket *output_socket : node->output_sockets()) {
const SocketDeclaration *output_decl = output_socket->runtime->declaration;
if (!output_decl) {
continue;
}
if (const decl::Bundle *bundle_decl = dynamic_cast<const decl::Bundle *>(output_decl)) {
if (bundle_decl->pass_through_input_index == socket->index()) {
add_if_new({socket.context, output_socket}, bundle_path);
}
}
}
}
else {
const bke::bNodeTreeZones *zones = node->owner_tree().zones();
if (!zones) {
continue;
}
const bke::bNodeTreeZone *from_zone = zones->get_zone_by_socket(*socket.socket);
for (const bNodeLink *link : socket->directly_linked_links()) {
if (!use_link_for_tracing(*link)) {
continue;
}
bNodeSocket *to_socket = link->tosock;
const bke::bNodeTreeZone *to_zone = zones->get_zone_by_socket(*to_socket);
if (!zones->link_between_zones_is_allowed(from_zone, to_zone)) {
continue;
}
const Vector<const bke::bNodeTreeZone *> zones_to_enter = zones->get_zones_to_enter(
from_zone, to_zone);
const ComputeContext *compute_context = ed::space_node::compute_context_for_zones(
zones_to_enter, compute_context_cache, socket.context);
if (!compute_context) {
continue;
}
add_if_new({compute_context, to_socket}, bundle_path);
}
}
}
return found_targets.extract_vector();
}
[[nodiscard]] const ComputeContext *compute_context_for_closure_evaluation(
const ComputeContext *closure_socket_context,
const bNodeSocket &closure_socket,
bke::ComputeContextCache &compute_context_cache,
const std::optional<ClosureSourceLocation> &source_location)
{
const Vector<SocketInContext> target_sockets = find_target_sockets_through_contexts(
{closure_socket_context, &closure_socket},
compute_context_cache,
is_evaluate_closure_node_input,
false);
if (target_sockets.is_empty()) {
return nullptr;
}
const SocketInContext target_socket = target_sockets[0];
const NodeInContext target_node = target_socket.owner_node();
return &compute_context_cache.for_evaluate_closure(target_socket.context,
target_node->identifier,
&target_socket->owner_tree(),
source_location);
}
static Vector<SocketInContext> find_origin_sockets_through_contexts(
const SocketInContext start_socket,
bke::ComputeContextCache &compute_context_cache,
const FunctionRef<bool(const SocketInContext &)> handle_possible_origin_socket_fn,
const bool find_all)
{
using BundlePath = Vector<std::string, 0>;
struct SocketToCheck {
SocketInContext socket;
BundlePath bundle_path;
};
Stack<SocketToCheck> sockets_to_check;
Set<SocketInContext> added_sockets;
auto add_if_new = [&](const SocketInContext &socket, BundlePath bundle_path) {
if (added_sockets.add(socket)) {
sockets_to_check.push({socket, std::move(bundle_path)});
}
};
add_if_new(start_socket, {});
VectorSet<SocketInContext> found_origins;
while (!sockets_to_check.is_empty()) {
const SocketToCheck socket_to_check = sockets_to_check.pop();
const SocketInContext socket = socket_to_check.socket;
const BundlePath &bundle_path = socket_to_check.bundle_path;
const NodeInContext &node = socket.owner_node();
const SocketDeclaration *socket_decl = socket->runtime->declaration;
if (socket->is_input()) {
if (bundle_path.is_empty() && handle_possible_origin_socket_fn(socket)) {
found_origins.add(socket);
if (!find_all) {
break;
}
continue;
}
const bke::bNodeTreeZones *zones = node->owner_tree().zones();
if (!zones) {
continue;
}
const bke::bNodeTreeZone *to_zone = zones->get_zone_by_socket(*socket.socket);
for (const bNodeLink *link : socket->directly_linked_links()) {
if (!use_link_for_tracing(*link)) {
continue;
}
const bNodeSocket *from_socket = link->fromsock;
const bke::bNodeTreeZone *from_zone = zones->get_zone_by_socket(*from_socket);
if (!zones->link_between_zones_is_allowed(from_zone, to_zone)) {
continue;
}
const Vector<const bke::bNodeTreeZone *> zones_to_enter = zones->get_zones_to_enter(
from_zone, to_zone);
const ComputeContext *compute_context = socket.context;
for (int i = zones_to_enter.size() - 1; i >= 0; i--) {
if (!compute_context) {
/* There must be a compute context when we are in a zone. */
BLI_assert_unreachable();
return found_origins.extract_vector();
}
/* Each zone corresponds to one compute context level. */
compute_context = compute_context->parent();
}
add_if_new({compute_context, from_socket}, bundle_path);
}
}
else {
if (node->is_muted()) {
for (const bNodeLink &link : node->internal_links()) {
if (link.tosock == socket.socket) {
add_if_new({socket.context, link.fromsock}, bundle_path);
}
}
continue;
}
if (bundle_path.is_empty() && handle_possible_origin_socket_fn(socket)) {
found_origins.add(socket);
if (!find_all) {
break;
}
continue;
}
if (node->is_reroute()) {
add_if_new(node.input_socket(0), bundle_path);
continue;
}
if (node->is_group()) {
if (const bNodeTree *group = reinterpret_cast<const bNodeTree *>(node->id)) {
group->ensure_topology_cache();
if (const bNode *group_output_node = group->group_output_node()) {
const ComputeContext &group_compute_context = compute_context_cache.for_group_node(
socket.context, node->identifier, &node->owner_tree());
if (const bNodeSocket *group_output_socket = group_output_node->input_by_identifier(
socket->identifier))
{
add_if_new({&group_compute_context, group_output_socket}, bundle_path);
}
}
}
continue;
}
if (node->is_group_input()) {
if (const auto *group_context = dynamic_cast<const bke::GroupNodeComputeContext *>(
socket.context))
{
const bNodeTree *caller_group = group_context->tree();
const bNode *caller_group_node = group_context->node();
if (caller_group && caller_group_node) {
caller_group->ensure_topology_cache();
if (const bNodeSocket *input_socket = caller_group_node->input_by_identifier(
socket->identifier))
{
add_if_new({group_context->parent(), input_socket}, bundle_path);
}
}
}
continue;
}
if (node->is_type("NodeEvaluateClosure")) {
const auto &evaluate_storage = *static_cast<const NodeEvaluateClosure *>(node->storage);
const StringRef key = evaluate_storage.output_items.items[socket->index()].name;
const Vector<SocketInContext> origin_sockets = find_origin_sockets_through_contexts(
node.input_socket(0), compute_context_cache, is_closure_zone_output_socket, true);
for (const SocketInContext origin_socket : origin_sockets) {
const bNodeTree &closure_tree = origin_socket->owner_tree();
const NodeInContext closure_output_node = origin_socket.owner_node();
const auto &closure_storage = *static_cast<const NodeClosureOutput *>(
closure_output_node->storage);
const ComputeContext &closure_context = compute_context_cache.for_evaluate_closure(
node.context,
node->identifier,
&node->owner_tree(),
ClosureSourceLocation{
&closure_tree, closure_output_node->identifier, origin_socket.context_hash()});
for (const int i : IndexRange(closure_storage.output_items.items_num)) {
const NodeClosureOutputItem &item = closure_storage.output_items.items[i];
if (key == item.name) {
add_if_new({&closure_context, &closure_output_node->input_socket(i)}, bundle_path);
}
}
}
continue;
}
if (node->is_type("NodeClosureInput")) {
const auto &input_storage = *static_cast<const NodeClosureInput *>(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 NodeClosureOutput *>(
closure_output_node->storage);
const StringRef key = output_storage.input_items.items[socket->index()].name;
const bNodeSocket &closure_output_socket = closure_output_node->output_socket(0);
const Vector<SocketInContext> target_sockets = find_target_sockets_through_contexts(
{socket.context, &closure_output_socket},
compute_context_cache,
is_evaluate_closure_node_input,
true);
for (const SocketInContext &target_socket : target_sockets) {
const NodeInContext target_node = target_socket.owner_node();
const auto &evaluate_storage = *static_cast<const NodeEvaluateClosure *>(
target_node.node->storage);
for (const int i : IndexRange(evaluate_storage.input_items.items_num)) {
const NodeEvaluateClosureInputItem &item = evaluate_storage.input_items.items[i];
if (key == item.name) {
add_if_new(target_node.input_socket(i + 1), bundle_path);
}
}
}
continue;
}
if (node->is_type("NodeCombineBundle")) {
if (bundle_path.is_empty()) {
continue;
}
const StringRef last_key = bundle_path.last();
const auto &storage = *static_cast<const NodeCombineBundle *>(node->storage);
for (const int input_i : IndexRange(storage.items_num)) {
if (last_key == storage.items[input_i].name) {
add_if_new(node.input_socket(input_i), bundle_path.as_span().drop_back(1));
}
}
continue;
}
if (node->is_type("NodeSeparateBundle")) {
const auto &storage = *static_cast<const NodeSeparateBundle *>(node->storage);
BundlePath new_bundle_path = bundle_path;
new_bundle_path.append(storage.items[socket->index()].name);
add_if_new(node.input_socket(0), std::move(new_bundle_path));
continue;
}
if (node->is_type("GeometryNodeSimulationInput")) {
const int output_index = socket->index();
if (output_index >= 1) {
BLI_assert(dynamic_cast<const bke::SimulationZoneComputeContext *>(socket.context));
add_if_new({socket.context->parent(), &node->input_socket(output_index - 1)},
bundle_path);
}
continue;
}
if (node->is_type("GeometryNodeSimulationOutput")) {
const ComputeContext &simulation_compute_context =
compute_context_cache.for_simulation_zone(socket.context, *node);
add_if_new({&simulation_compute_context, &node->input_socket(socket->index() + 1)},
bundle_path);
continue;
}
if (node->is_type("GeometryNodeRepeatInput")) {
const int index = socket->index();
if (index >= 1) {
BLI_assert(dynamic_cast<const bke::RepeatZoneComputeContext *>(socket.context));
add_if_new({socket.context->parent(), &node->input_socket(index)}, bundle_path);
}
continue;
}
if (node->is_type("GeometryNodeRepeatOutput")) {
const int index = socket->index();
const ComputeContext &repeat_compute_context = compute_context_cache.for_repeat_zone(
socket.context, *node, 0);
add_if_new({&repeat_compute_context, &node->input_socket(index)}, bundle_path);
const bke::bNodeZoneType &zone_type = *bke::zone_type_by_node_type(node->type_legacy);
if (const bNode *repeat_input_node = zone_type.get_corresponding_input(node->owner_tree(),
*node))
{
add_if_new({socket.context, &repeat_input_node->input_socket(index + 1)}, bundle_path);
}
continue;
}
if (socket_decl) {
if (const decl::Bundle *bundle_decl = dynamic_cast<const decl::Bundle *>(socket_decl)) {
if (bundle_decl->pass_through_input_index) {
const int input_index = *bundle_decl->pass_through_input_index;
add_if_new(node.input_socket(input_index), bundle_path);
}
}
}
}
}
return found_origins.extract_vector();
}
Vector<BundleSignature> gather_linked_target_bundle_signatures(
const ComputeContext *bundle_socket_context,
const bNodeSocket &bundle_socket,
bke::ComputeContextCache &compute_context_cache)
{
const Vector<SocketInContext> target_sockets = find_target_sockets_through_contexts(
{bundle_socket_context, &bundle_socket},
compute_context_cache,
[](const SocketInContext &socket) {
return socket->is_input() && socket->owner_node().is_type("NodeSeparateBundle");
},
true);
Vector<BundleSignature> signatures;
for (const SocketInContext &target_socket : target_sockets) {
const NodeInContext &target_node = target_socket.owner_node();
signatures.append(BundleSignature::from_separate_bundle_node(*target_node.node));
}
return signatures;
}
Vector<BundleSignature> gather_linked_origin_bundle_signatures(
const ComputeContext *bundle_socket_context,
const bNodeSocket &bundle_socket,
bke::ComputeContextCache &compute_context_cache)
{
const Vector<SocketInContext> origin_sockets = find_origin_sockets_through_contexts(
{bundle_socket_context, &bundle_socket},
compute_context_cache,
[](const SocketInContext &socket) {
return socket->is_output() && socket->owner_node().is_type("NodeCombineBundle");
},
true);
Vector<BundleSignature> signatures;
for (const SocketInContext &origin_socket : origin_sockets) {
const NodeInContext &origin_node = origin_socket.owner_node();
signatures.append(BundleSignature::from_combine_bundle_node(*origin_node.node));
}
return signatures;
}
Vector<ClosureSignature> gather_linked_target_closure_signatures(
const ComputeContext *closure_socket_context,
const bNodeSocket &closure_socket,
bke::ComputeContextCache &compute_context_cache)
{
const Vector<SocketInContext> target_sockets = find_target_sockets_through_contexts(
{closure_socket_context, &closure_socket},
compute_context_cache,
is_evaluate_closure_node_input,
true);
Vector<ClosureSignature> signatures;
for (const SocketInContext &target_socket : target_sockets) {
const NodeInContext &target_node = target_socket.owner_node();
signatures.append(ClosureSignature::from_evaluate_closure_node(*target_node.node));
}
return signatures;
}
Vector<ClosureSignature> gather_linked_origin_closure_signatures(
const ComputeContext *closure_socket_context,
const bNodeSocket &closure_socket,
bke::ComputeContextCache &compute_context_cache)
{
Vector<ClosureSignature> signatures;
find_origin_sockets_through_contexts(
{closure_socket_context, &closure_socket},
compute_context_cache,
[&](const SocketInContext &socket) {
if (is_closure_zone_output_socket(socket)) {
signatures.append(ClosureSignature::from_closure_output_node(socket->owner_node()));
return true;
}
return false;
},
true);
return signatures;
}
} // namespace blender::nodes
Reference in New Issue View Git Blame Copy Permalink