Cleanup: Avoid using runtime node flag, use topology cache

It's easier to keep track of state in these algorithms if it's stored in
a central place like a set. Plus, using flags requires clearing them
beforehand. For the selected linked operators, using the topology
cache means we don't have to iterate over all links.

5 files changed, 57 insertions, 51 deletions

diff --git a/source/blender/blenkernel/BKE_node_runtime.hh b/source/blender/blenkernel/BKE_node_runtime.hh
index 4c13f73848c..49a6953d0a3 100644
--- a/source/blender/blenkernel/BKE_node_runtime.hh
+++ b/source/blender/blenkernel/BKE_node_runtime.hh
@@ -481,6 +481,12 @@ inline blender::Span<const bNodeSocket *> bNodeSocket::directly_linked_sockets()
   return this->runtime->directly_linked_sockets;
 }
 
+inline blender::Span<bNodeSocket *> bNodeSocket::directly_linked_sockets()
+{
+  BLI_assert(blender::bke::node_tree_runtime::topology_cache_is_available(*this));
+  return this->runtime->directly_linked_sockets;
+}
+
 inline bool bNodeSocket::is_directly_linked() const
 {
   return !this->directly_linked_links().is_empty();
diff --git a/source/blender/editors/space_node/node_intern.hh b/source/blender/editors/space_node/node_intern.hh
index c4f3227f274..456cbf5064d 100644
--- a/source/blender/editors/space_node/node_intern.hh
+++ b/source/blender/editors/space_node/node_intern.hh
@@ -9,6 +9,7 @@
 
 #include "BLI_math_vector.h"
 #include "BLI_math_vector.hh"
+#include "BLI_set.hh"
 #include "BLI_vector.hh"
 
 #include "BKE_node.h"
@@ -171,6 +172,7 @@ void node_keymap(wmKeyConfig *keyconf);
 rctf node_frame_rect_inside(const bNode &node);
 bool node_or_socket_isect_event(const bContext &C, const wmEvent &event);
 
+Set<bNode *> get_selected_nodes(bNodeTree &node_tree);
 void node_deselect_all(SpaceNode &snode);
 void node_socket_select(bNode *node, bNodeSocket &sock);
 void node_socket_deselect(bNode *node, bNodeSocket &sock, bool deselect_node);
diff --git a/source/blender/editors/space_node/node_relationships.cc b/source/blender/editors/space_node/node_relationships.cc
index c28b345b111..7dbaa8ccd6d 100644
--- a/source/blender/editors/space_node/node_relationships.cc
+++ b/source/blender/editors/space_node/node_relationships.cc
@@ -1619,7 +1619,9 @@ void NODE_OT_parent_set(wmOperatorType *ot)
 #define NODE_JOIN_DONE 1
 #define NODE_JOIN_IS_DESCENDANT 2
 
-static void node_join_attach_recursive(bNode *node, bNode *frame)
+static void node_join_attach_recursive(bNode *node,
+                                       bNode *frame,
+                                       const Set<bNode *> &selected_nodes)
 {
   node->done |= NODE_JOIN_DONE;
 
@@ -1629,21 +1631,21 @@ static void node_join_attach_recursive(bNode *node, bNode *frame)
   else if (node->parent) {
     /* call recursively */
     if (!(node->parent->done & NODE_JOIN_DONE)) {
-      node_join_attach_recursive(node->parent, frame);
+      node_join_attach_recursive(node->parent, frame, selected_nodes);
     }
 
     /* in any case: if the parent is a descendant, so is the child */
     if (node->parent->done & NODE_JOIN_IS_DESCENDANT) {
       node->done |= NODE_JOIN_IS_DESCENDANT;
     }
-    else if (node->flag & NODE_TEST) {
+    else if (selected_nodes.contains(node)) {
       /* if parent is not an descendant of the frame, reattach the node */
       nodeDetachNode(node);
       nodeAttachNode(node, frame);
       node->done |= NODE_JOIN_IS_DESCENDANT;
     }
   }
-  else if (node->flag & NODE_TEST) {
+  else if (selected_nodes.contains(node)) {
     nodeAttachNode(node, frame);
     node->done |= NODE_JOIN_IS_DESCENDANT;
   }
@@ -1651,21 +1653,13 @@ static void node_join_attach_recursive(bNode *node, bNode *frame)
 
 static int node_join_exec(bContext *C, wmOperator *UNUSED(op))
 {
+  Main &bmain = *CTX_data_main(C);
   SpaceNode &snode = *CTX_wm_space_node(C);
   bNodeTree &ntree = *snode.edittree;
 
-  /* XXX save selection: add_static_node call below sets the new frame as single
-   * active+selected node */
-  LISTBASE_FOREACH (bNode *, node, &ntree.nodes) {
-    if (node->flag & NODE_SELECT) {
-      node->flag |= NODE_TEST;
-    }
-    else {
-      node->flag &= ~NODE_TEST;
-    }
-  }
+  const Set<bNode *> selected_nodes = get_selected_nodes(ntree);
 
-  bNode *frame = add_static_node(*C, NODE_FRAME, float2(0));
+  bNode *frame_node = nodeAddStaticNode(C, &ntree, NODE_FRAME);
 
   /* reset tags */
   LISTBASE_FOREACH (bNode *, node, &ntree.nodes) {
@@ -1674,18 +1668,12 @@ static int node_join_exec(bContext *C, wmOperator *UNUSED(op))
 
   LISTBASE_FOREACH (bNode *, node, &ntree.nodes) {
     if (!(node->done & NODE_JOIN_DONE)) {
-      node_join_attach_recursive(node, frame);
-    }
-  }
-
-  /* restore selection */
-  LISTBASE_FOREACH (bNode *, node, &ntree.nodes) {
-    if (node->flag & NODE_TEST) {
-      node->flag |= NODE_SELECT;
+      node_join_attach_recursive(node, frame_node, selected_nodes);
     }
   }
 
   node_sort(ntree);
+  ED_node_tree_propagate_change(C, &bmain, snode.edittree);
   WM_event_add_notifier(C, NC_NODE | ND_DISPLAY, nullptr);
 
   return OPERATOR_FINISHED;
diff --git a/source/blender/editors/space_node/node_select.cc b/source/blender/editors/space_node/node_select.cc
index 563dadc511a..1f1ce9c0c2b 100644
--- a/source/blender/editors/space_node/node_select.cc
+++ b/source/blender/editors/space_node/node_select.cc
@@ -311,6 +311,17 @@ void node_deselect_all_output_sockets(SpaceNode &snode, const bool deselect_node
   }
 }
 
+Set<bNode *> get_selected_nodes(bNodeTree &node_tree)
+{
+  Set<bNode *> selected_nodes;
+  for (bNode *node : node_tree.all_nodes()) {
+    if (node->flag & NODE_SELECT) {
+      selected_nodes.add(node);
+    }
+  }
+  return selected_nodes;
+}
+
 /** \} */
 
 /* -------------------------------------------------------------------- */
@@ -1112,22 +1123,21 @@ static int node_select_linked_to_exec(bContext *C, wmOperator *UNUSED(op))
   SpaceNode &snode = *CTX_wm_space_node(C);
   bNodeTree &node_tree = *snode.edittree;
 
-  LISTBASE_FOREACH (bNode *, node, &node_tree.nodes) {
-    node->flag &= ~NODE_TEST;
-  }
+  node_tree.ensure_topology_cache();
 
-  LISTBASE_FOREACH (bNodeLink *, link, &node_tree.links) {
-    if (nodeLinkIsHidden(link)) {
-      continue;
-    }
-    if (link->fromnode && link->tonode && (link->fromnode->flag & NODE_SELECT)) {
-      link->tonode->flag |= NODE_TEST;
-    }
-  }
+  Set<bNode *> initial_selection = get_selected_nodes(node_tree);
 
-  LISTBASE_FOREACH (bNode *, node, &node_tree.nodes) {
-    if (node->flag & NODE_TEST) {
-      nodeSetSelected(node, true);
+  for (bNode *node : initial_selection) {
+    for (bNodeSocket *output_socket : node->output_sockets()) {
+      if (!output_socket->is_available()) {
+        continue;
+      }
+      for (bNodeSocket *input_socket : output_socket->directly_linked_sockets()) {
+        if (!input_socket->is_available()) {
+          continue;
+        }
+        nodeSetSelected(&input_socket->owner_node(), true);
+      }
     }
   }
 
@@ -1163,22 +1173,21 @@ static int node_select_linked_from_exec(bContext *C, wmOperator *UNUSED(op))
   SpaceNode &snode = *CTX_wm_space_node(C);
   bNodeTree &node_tree = *snode.edittree;
 
-  LISTBASE_FOREACH (bNode *, node, &node_tree.nodes) {
-    node->flag &= ~NODE_TEST;
-  }
+  node_tree.ensure_topology_cache();
 
-  LISTBASE_FOREACH (bNodeLink *, link, &node_tree.links) {
-    if (nodeLinkIsHidden(link)) {
-      continue;
-    }
-    if (link->fromnode && link->tonode && (link->tonode->flag & NODE_SELECT)) {
-      link->fromnode->flag |= NODE_TEST;
-    }
-  }
+  Set<bNode *> initial_selection = get_selected_nodes(node_tree);
 
-  LISTBASE_FOREACH (bNode *, node, &node_tree.nodes) {
-    if (node->flag & NODE_TEST) {
-      nodeSetSelected(node, true);
+  for (bNode *node : initial_selection) {
+    for (bNodeSocket *input_socket : node->input_sockets()) {
+      if (!input_socket->is_available()) {
+        continue;
+      }
+      for (bNodeSocket *output_socket : input_socket->directly_linked_sockets()) {
+        if (!output_socket->is_available()) {
+          continue;
+        }
+        nodeSetSelected(&output_socket->owner_node(), true);
+      }
     }
   }
 
diff --git a/source/blender/makesdna/DNA_node_types.h b/source/blender/makesdna/DNA_node_types.h
index 0114988a0bc..67ff6586de0 100644
--- a/source/blender/makesdna/DNA_node_types.h
+++ b/source/blender/makesdna/DNA_node_types.h
@@ -210,6 +210,7 @@ typedef struct bNodeSocket {
   blender::Span<bNodeLink *> directly_linked_links();
   blender::Span<const bNodeLink *> directly_linked_links() const;
   /** Sockets which are connected to this socket with a link. */
+  blender::Span<bNodeSocket *> directly_linked_sockets();
   blender::Span<const bNodeSocket *> directly_linked_sockets() const;
   bool is_directly_linked() const;
   /**