Fix: GPU: Fix closure evaluation order

EEVEE-Next can only store data for a single (global) closure at a time,
which, when combined with ShaderToRGB nodes, requires extra care in the
order that closures are evaluated.
For example:
![image](/attachments/0c56613f-3515-40a2-bf0e-282a8a99d64e)
Here, after `ntree_shader_shader_to_rgba_branch` there will be 2 Diffuse
nodes (the original and a copy for the ShaderToRGB branch).
However, the generated code order will be something like this:
```
Diffuse (original)
Diffuse (copy)
ShaderToRGB // This resets closures
Mix
```
So while the original node is technically "evaluated", the closure data
is reset after ShaderToRGB.
This patch updates the code generation to ensure closure evaluation is
ordered taking ShaderToRGB branches into account, so the generated code
looks like this:
```
Diffuse (copy)
ShaderToRGB // This resets closures
Diffuse (original)
Mix
```
This also fixes ShaderToRGB support for AOVs, removes unused code, and
fixes several bugs that I've found along the way that were harmless for
EEVEE but broke EEVEE Next.

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

source/blender/blenkernel/BKE_node.hh

@@ -163,6 +163,9 @@ void nodeChainIter(const bNodeTree *ntree,
  * Can be called recursively (using another nodeChainIterBackwards) by
  * setting the recursion_lvl accordingly.
  *
+ * WARN: No node is guaranteed to be iterated as a to_node,
+ * since it could have been iterated earlier as a from_node.
+ *
  * \note Needs updated socket links (ntreeUpdateTree).
  * \note Recursive
  */

source/blender/nodes/shader/node_shader_tree.cc

@@ -21,6 +21,7 @@
 #include "BLI_linklist.h"
 #include "BLI_listbase.h"
 #include "BLI_math_vector.h"
+#include "BLI_set.hh"
 #include "BLI_string.h"
 #include "BLI_threads.h"
 #include "BLI_utildefines.h"
@@ -580,69 +581,87 @@ static bool ntree_branch_count_and_tag_nodes(bNode *fromnode, bNode *tonode, voi
   return true;
 }
 
-/* Create a copy of a branch starting from a given node.
- * callback is executed once for every copied node.
- * Returns input node copy. */
-static bNode *ntree_shader_copy_branch(bNodeTree *ntree,
-                                       bNode *start_node,
-                                       bool (*node_filter)(const bNode *node),
-                                       void (*callback)(bNode *node, int user_data),
-                                       int user_data)
+/* Create a copy of a branch starting from a given node. */
+static void ntree_shader_copy_branch(bNodeTree *ntree,
+                                     bNode *start_node,
+                                     bool (*node_filter)(const bNode *node))
 {
+  auto gather_branch_nodes = [](bNode *fromnode, bNode * /*tonode*/, void *userdata) {
+    blender::Set<bNode *> *set = static_cast<blender::Set<bNode *> *>(userdata);
+    set->add(fromnode);
+    return true;
+  };
+  blender::Set<bNode *> branch_nodes = {start_node};
+  blender::bke::nodeChainIterBackwards(ntree, start_node, gather_branch_nodes, &branch_nodes, 0);
+
   /* Initialize `runtime->tmp_flag`. */
   LISTBASE_FOREACH (bNode *, node, &ntree->nodes) {
     node->runtime->tmp_flag = -1;
   }
   /* Count and tag all nodes inside the displacement branch of the tree. */
-  start_node->runtime->tmp_flag = 0;
   branchIterData iter_data;
   iter_data.node_filter = node_filter;
-  iter_data.node_count = 1;
+  iter_data.node_count = 0;
   blender::bke::nodeChainIterBackwards(
       ntree, start_node, ntree_branch_count_and_tag_nodes, &iter_data, 1);
-  /* Make a full copy of the branch */
+  /* Copies of the non-filtered nodes on the branch. */
   Array<bNode *> nodes_copy(iter_data.node_count);
+
   LISTBASE_FOREACH (bNode *, node, &ntree->nodes) {
     if (node->runtime->tmp_flag >= 0) {
       int id = node->runtime->tmp_flag;
-      /* Avoid creating unique names in the new tree, since it is very slow. The names on the new
-       * nodes will be invalid. But identifiers must be created for the `bNodeTree::all_nodes()`
-       * vector, though they won't match the original. */
+      /* Avoid creating unique names in the new tree, since it is very slow.
+       * The names on the new nodes will be invalid. */
       nodes_copy[id] = blender::bke::node_copy(
           ntree, *node, LIB_ID_CREATE_NO_USER_REFCOUNT | LIB_ID_CREATE_NO_MAIN, false);
+      /* But identifiers must be created for the `bNodeTree::all_nodes()` vector,
+       * so they won't match the original. */
       nodeUniqueID(ntree, nodes_copy[id]);
 
-      nodes_copy[id]->runtime->tmp_flag = -2; /* Copy */
-      nodes_copy[id]->runtime->original = node->runtime->original;
+      bNode *copy = nodes_copy[id];
+      copy->runtime->tmp_flag = -2; /* Copy */
+      copy->runtime->original = node->runtime->original;
       /* Make sure to clear all sockets links as they are invalid. */
-      LISTBASE_FOREACH (bNodeSocket *, sock, &nodes_copy[id]->inputs) {
+      LISTBASE_FOREACH (bNodeSocket *, sock, &copy->inputs) {
         sock->link = nullptr;
       }
-      LISTBASE_FOREACH (bNodeSocket *, sock, &nodes_copy[id]->outputs) {
+      LISTBASE_FOREACH (bNodeSocket *, sock, &copy->outputs) {
         sock->link = nullptr;
       }
     }
   }
-  /* Recreate links between copied nodes AND incoming links to the copied nodes. */
-  LISTBASE_FOREACH (bNodeLink *, link, &ntree->links) {
-    if (link->tonode->runtime->tmp_flag >= 0) {
-      bool from_node_copied = link->fromnode->runtime->tmp_flag >= 0;
-      bNode *fromnode = from_node_copied ? nodes_copy[link->fromnode->runtime->tmp_flag] :
-                                           link->fromnode;
-      bNode *tonode = nodes_copy[link->tonode->runtime->tmp_flag];
-      bNodeSocket *fromsock = ntree_shader_node_find_output(fromnode, link->fromsock->identifier);
-      bNodeSocket *tosock = ntree_shader_node_find_input(tonode, link->tosock->identifier);
-      nodeAddLink(ntree, fromnode, fromsock, tonode, tosock);
+
+  /* Unlink the original nodes from this branch and link the copies. */
+  LISTBASE_FOREACH_MUTABLE (bNodeLink *, link, &ntree->links) {
+    bool from_copy = link->fromnode->runtime->tmp_flag >= 0;
+    bool to_copy = link->tonode->runtime->tmp_flag >= 0;
+    if (from_copy && to_copy) {
+      bNode *from_node = nodes_copy[link->fromnode->runtime->tmp_flag];
+      bNode *to_node = nodes_copy[link->tonode->runtime->tmp_flag];
+      nodeAddLink(ntree,
+                  from_node,
+                  ntree_shader_node_find_output(from_node, link->fromsock->identifier),
+                  to_node,
+                  ntree_shader_node_find_input(to_node, link->tosock->identifier));
+    }
+    else if (to_copy) {
+      bNode *to_node = nodes_copy[link->tonode->runtime->tmp_flag];
+      nodeAddLink(ntree,
+                  link->fromnode,
+                  link->fromsock,
+                  to_node,
+                  ntree_shader_node_find_input(to_node, link->tosock->identifier));
+    }
+    else if (from_copy && branch_nodes.contains(link->tonode)) {
+      bNode *from_node = nodes_copy[link->fromnode->runtime->tmp_flag];
+      nodeAddLink(ntree,
+                  from_node,
+                  ntree_shader_node_find_output(from_node, link->fromsock->identifier),
+                  link->tonode,
+                  link->tosock);
+      nodeRemLink(ntree, link);
     }
   }
-  /* Per node callback. */
-  if (callback) {
-    for (int i = 0; i < iter_data.node_count; i++) {
-      callback(nodes_copy[i], user_data);
-    }
-  }
-  bNode *start_node_copy = nodes_copy[start_node->runtime->tmp_flag];
-  return start_node_copy;
 }
 
 /* Generate emission node to convert regular data to closure sockets.
@@ -985,47 +1004,51 @@ static bool closure_node_filter(const bNode *node)
   }
 }
 
-static bool shader_to_rgba_node_gather(bNode * /*fromnode*/, bNode *tonode, void *userdata)
-{
-  Vector<bNode *> &shader_to_rgba_nodes = *(Vector<bNode *> *)userdata;
-  if (tonode->runtime->tmp_flag == -1 && tonode->type == SH_NODE_SHADERTORGB) {
-    tonode->runtime->tmp_flag = 0;
-    shader_to_rgba_nodes.append(tonode);
-  }
-  return true;
-}
-
 /* Shader to rgba needs their associated closure duplicated and the weight tree generated for. */
 static void ntree_shader_shader_to_rgba_branch(bNodeTree *ntree, bNode *output_node)
 {
-  LISTBASE_FOREACH (bNode *, node, &ntree->nodes) {
-    node->runtime->tmp_flag = -1;
-  }
-  /* First gather the shader_to_rgba nodes linked to the output. This is separate to avoid
-   * conflicting usage of the `node->runtime->tmp_flag`. */
   Vector<bNode *> shader_to_rgba_nodes;
-  blender::bke::nodeChainIterBackwards(
-      ntree, output_node, shader_to_rgba_node_gather, &shader_to_rgba_nodes, 0);
+  LISTBASE_FOREACH (bNode *, node, &ntree->nodes) {
+    if (node->type == SH_NODE_SHADERTORGB) {
+      shader_to_rgba_nodes.append(node);
+    }
+  }
 
   for (bNode *shader_to_rgba : shader_to_rgba_nodes) {
     bNodeSocket *closure_input = ntree_shader_node_input_get(shader_to_rgba, 0);
     if (closure_input->link == nullptr) {
       continue;
     }
-    bNode *start_node = closure_input->link->fromnode;
-    bNode *start_node_copy = ntree_shader_copy_branch(
-        ntree, start_node, closure_node_filter, nullptr, 0);
-    /* Replace node copy link. This assumes that every node possibly connected to the closure input
-     * has only one output. */
-    bNodeSocket *closure_output = ntree_shader_node_output_get(start_node_copy, 0);
-    nodeRemLink(ntree, closure_input->link);
-    nodeAddLink(ntree, start_node_copy, closure_output, shader_to_rgba, closure_input);
+    ntree_shader_copy_branch(ntree, shader_to_rgba, closure_node_filter);
     BKE_ntree_update_main_tree(G.main, ntree, nullptr);
 
     ntree_shader_weight_tree_invert(ntree, shader_to_rgba);
   }
 }
 
+static void iter_shader_to_rgba_depth_count(bNode *node,
+                                            int16_t &max_depth,
+                                            int16_t depth_level = 0)
+{
+  if (node->type == SH_NODE_SHADERTORGB) {
+    depth_level++;
+    max_depth = std::max(max_depth, depth_level);
+  }
+  node->runtime->tmp_flag = std::max(node->runtime->tmp_flag, depth_level);
+
+  LISTBASE_FOREACH (bNodeSocket *, sock, &node->inputs) {
+    bNodeLink *link = sock->link;
+    if (link == nullptr) {
+      continue;
+    }
+    if ((link->flag & NODE_LINK_VALID) == 0) {
+      /* Skip links marked as cyclic. */
+      continue;
+    }
+    iter_shader_to_rgba_depth_count(link->fromnode, max_depth, depth_level);
+  }
+}
+
 static void shader_node_disconnect_input(bNodeTree *ntree, bNode *node, int index)
 {
   bNodeLink *link = ntree_shader_node_input_get(node, index)->link;
@@ -1180,10 +1203,24 @@ void ntreeGPUMaterialNodes(bNodeTree *localtree, GPUMaterial *mat)
   }
 
   exec = ntreeShaderBeginExecTree(localtree);
-  ntreeExecGPUNodes(exec, mat, output);
+  /* Execute nodes ordered by the number of ShaderToRGB nodes found in their path,
+   * so all closures can be properly evaluated. */
+  int16_t max_depth = 0;
+  LISTBASE_FOREACH (bNode *, node, &localtree->nodes) {
+    node->runtime->tmp_flag = -1;
+  }
+  iter_shader_to_rgba_depth_count(output, max_depth);
   LISTBASE_FOREACH (bNode *, node, &localtree->nodes) {
     if (node->type == SH_NODE_OUTPUT_AOV) {
-      ntreeExecGPUNodes(exec, mat, node);
+      iter_shader_to_rgba_depth_count(node, max_depth);
+    }
+  }
+  for (int depth = max_depth; depth >= 0; depth--) {
+    ntreeExecGPUNodes(exec, mat, output, &depth);
+    LISTBASE_FOREACH (bNode *, node, &localtree->nodes) {
+      if (node->type == SH_NODE_OUTPUT_AOV) {
+        ntreeExecGPUNodes(exec, mat, node, &depth);
+      }
     }
   }
   ntreeShaderEndExecTree(exec);

source/blender/nodes/shader/node_shader_util.cc

@@ -305,7 +305,7 @@ bNode *nodeGetActivePaintCanvas(bNodeTree *ntree)
 }
 }  // namespace blender::bke
 
-void ntreeExecGPUNodes(bNodeTreeExec *exec, GPUMaterial *mat, bNode *output_node)
+void ntreeExecGPUNodes(bNodeTreeExec *exec, GPUMaterial *mat, bNode *output_node, int *depth_level)
 {
   bNodeExec *nodeexec;
   bNode *node;
@@ -321,6 +321,10 @@ void ntreeExecGPUNodes(bNodeTreeExec *exec, GPUMaterial *mat, bNode *output_node
   for (n = 0, nodeexec = exec->nodeexec; n < exec->totnodes; n++, nodeexec++) {
     node = nodeexec->node;
 
+    if (depth_level && node->runtime->tmp_flag != *depth_level) {
+      continue;
+    }
+
     do_it = false;
     /* for groups, only execute outputs for edited group */
     if (node->typeinfo->nclass == NODE_CLASS_OUTPUT) {
@@ -334,6 +338,7 @@ void ntreeExecGPUNodes(bNodeTreeExec *exec, GPUMaterial *mat, bNode *output_node
     }
 
     if (do_it) {
+      BLI_assert(!depth_level || node->runtime->tmp_flag >= 0);
       if (node->typeinfo->gpu_fn) {
         node_get_stack(node, stack, nsin, nsout);
         gpu_stack_from_data_list(gpuin, &node->inputs, nsin);

source/blender/nodes/shader/node_shader_util.hh

@@ -68,7 +68,14 @@ bNodeTreeExec *ntreeShaderBeginExecTree_internal(bNodeExecContext *context,
                                                  bNodeInstanceKey parent_key);
 void ntreeShaderEndExecTree_internal(bNodeTreeExec *exec);
 
-void ntreeExecGPUNodes(bNodeTreeExec *exec, GPUMaterial *mat, bNode *output_node);
+/* If depth_level is not null, only nodes where `node->runtime->tmp_flag == depth_level` will be
+ * executed. This allows finer control over node execution order without modifying the tree
+ * topology. */
+void ntreeExecGPUNodes(bNodeTreeExec *exec,
+                       GPUMaterial *mat,
+                       bNode *output_node,
+                       int *depth_level = nullptr);
+
 void get_XYZ_to_RGB_for_gpu(XYZ_to_RGB *data);
 
 bool node_socket_not_zero(const GPUNodeStack &socket);