3 files changed, 63 insertions, 121 deletions

diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index 1118552b643..8affbf0ca67 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -707,7 +707,10 @@ struct bNodeLink *nodeAddLink(struct bNodeTree *ntree,
                               struct bNodeSocket *tosock);
 void nodeRemLink(struct bNodeTree *ntree, struct bNodeLink *link);
 void nodeRemSocketLinks(struct bNodeTree *ntree, struct bNodeSocket *sock);
-void nodeMuteLinkToggle(struct bNodeTree *ntree, struct bNodeLink *link);
+/**
+ * Set the mute status of a single link.
+ */
+void nodeLinkSetMute(struct bNodeTree *ntree, struct bNodeLink *link, const bool muted);
 bool nodeLinkIsHidden(const struct bNodeLink *link);
 bool nodeLinkIsSelected(const struct bNodeLink *link);
 void nodeInternalRelink(struct bNodeTree *ntree, struct bNode *node);
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index c2ad0a93d83..e1eaed71f37 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -118,9 +118,6 @@ static void node_free_node(bNodeTree *ntree, bNode *node);
 static void node_socket_interface_free(bNodeTree *UNUSED(ntree),
                                        bNodeSocket *sock,
                                        const bool do_id_user);
-static void nodeMuteRerouteOutputLinks(struct bNodeTree *ntree,
-                                       struct bNode *node,
-                                       const bool mute);
 
 static void ntree_init_data(ID *id)
 {
@@ -2350,102 +2347,11 @@ void nodeRemLink(bNodeTree *ntree, bNodeLink *link)
   }
 }
 
-/* Check if all output links are muted or not. */
-static bool nodeMuteFromSocketLinks(const bNodeTree *ntree, const bNodeSocket *sock)
+void nodeLinkSetMute(bNodeTree *ntree, bNodeLink *link, const bool muted)
 {
-  int tot = 0;
-  int muted = 0;
-  LISTBASE_FOREACH (const bNodeLink *, link, &ntree->links) {
-    if (link->fromsock == sock) {
-      tot++;
-      if (link->flag & NODE_LINK_MUTED) {
-        muted++;
-      }
-    }
-  }
-  return tot == muted;
-}
-
-static void nodeMuteLink(bNodeLink *link)
-{
-  link->flag |= NODE_LINK_MUTED;
-  link->flag |= NODE_LINK_TEST;
-  if (!(link->tosock->flag & SOCK_MULTI_INPUT)) {
-    link->tosock->flag &= ~SOCK_IN_USE;
-  }
-}
-
-static void nodeUnMuteLink(bNodeLink *link)
-{
-  link->flag &= ~NODE_LINK_MUTED;
-  link->flag |= NODE_LINK_TEST;
-  link->tosock->flag |= SOCK_IN_USE;
-}
-
-/* Upstream muting. Always happens when unmuting but checks when muting. O(n^2) algorithm. */
-static void nodeMuteRerouteInputLinks(bNodeTree *ntree, bNode *node, const bool mute)
-{
-  if (node->type != NODE_REROUTE) {
-    return;
-  }
-  if (!mute || nodeMuteFromSocketLinks(ntree, (bNodeSocket *)node->outputs.first)) {
-    bNodeSocket *sock = (bNodeSocket *)node->inputs.first;
-    LISTBASE_FOREACH (bNodeLink *, link, &ntree->links) {
-      if (!(link->flag & NODE_LINK_VALID) || (link->tosock != sock)) {
-        continue;
-      }
-      if (mute) {
-        nodeMuteLink(link);
-      }
-      else {
-        nodeUnMuteLink(link);
-      }
-      nodeMuteRerouteInputLinks(ntree, link->fromnode, mute);
-    }
-  }
-}
-
-/* Downstream muting propagates when reaching reroute nodes. O(n^2) algorithm. */
-static void nodeMuteRerouteOutputLinks(bNodeTree *ntree, bNode *node, const bool mute)
-{
-  if (node->type != NODE_REROUTE) {
-    return;
-  }
-  bNodeSocket *sock;
-  sock = (bNodeSocket *)node->outputs.first;
-  LISTBASE_FOREACH (bNodeLink *, link, &ntree->links) {
-    if (!(link->flag & NODE_LINK_VALID) || (link->fromsock != sock)) {
-      continue;
-    }
-    if (mute) {
-      nodeMuteLink(link);
-    }
-    else {
-      nodeUnMuteLink(link);
-    }
-    nodeMuteRerouteOutputLinks(ntree, link->tonode, mute);
-  }
-}
-
-void nodeMuteLinkToggle(bNodeTree *ntree, bNodeLink *link)
-{
-  if (link->tosock) {
-    bool mute = !(link->flag & NODE_LINK_MUTED);
-    if (mute) {
-      nodeMuteLink(link);
-    }
-    else {
-      nodeUnMuteLink(link);
-    }
-    if (link->tonode->type == NODE_REROUTE) {
-      nodeMuteRerouteOutputLinks(ntree, link->tonode, mute);
-    }
-    if (link->fromnode->type == NODE_REROUTE) {
-      nodeMuteRerouteInputLinks(ntree, link->fromnode, mute);
-    }
-  }
-
-  if (ntree) {
+  const bool was_muted = link->flag & NODE_LINK_MUTED;
+  SET_FLAG_FROM_TEST(link->flag, muted, NODE_LINK_MUTED);
+  if (muted != was_muted) {
     BKE_ntree_update_tag_link_mute(ntree, link);
   }
 }
diff --git a/source/blender/editors/space_node/node_relationships.cc b/source/blender/editors/space_node/node_relationships.cc
index 7d59bb6cb0c..e95de7b8e85 100644
--- a/source/blender/editors/space_node/node_relationships.cc
+++ b/source/blender/editors/space_node/node_relationships.cc
@@ -11,6 +11,7 @@
 #include "DNA_node_types.h"
 
 #include "BLI_easing.h"
+#include "BLI_stack.hh"
 
 #include "BKE_anim_data.h"
 #include "BKE_context.h"
@@ -1390,11 +1391,22 @@ void NODE_OT_links_cut(wmOperatorType *ot)
 /** \name Mute Links Operator
  * \{ */
 
+static bool all_links_muted(const bNodeSocket &socket)
+{
+  for (const bNodeLink *link : socket.directly_linked_links()) {
+    if (!(link->flag & NODE_LINK_MUTED)) {
+      return false;
+    }
+  }
+  return true;
+}
+
 static int mute_links_exec(bContext *C, wmOperator *op)
 {
   Main &bmain = *CTX_data_main(C);
   SpaceNode &snode = *CTX_wm_space_node(C);
   const ARegion &region = *CTX_wm_region(C);
+  bNodeTree &ntree = *snode.edittree;
 
   Vector<float2> path;
   RNA_BEGIN (op->ptr, itemptr, "path") {
@@ -1415,41 +1427,62 @@ static int mute_links_exec(bContext *C, wmOperator *op)
 
   ED_preview_kill_jobs(CTX_wm_manager(C), &bmain);
 
-  /* Count intersected links and clear test flag. */
-  int tot = 0;
-  LISTBASE_FOREACH (bNodeLink *, link, &snode.edittree->links) {
+  ntree.ensure_topology_cache();
+
+  Set<bNodeLink *> affected_links;
+  LISTBASE_FOREACH (bNodeLink *, link, &ntree.links) {
     if (node_link_is_hidden_or_dimmed(region.v2d, *link)) {
       continue;
     }
-    link->flag &= ~NODE_LINK_TEST;
-    if (link_path_intersection(*link, path)) {
-      tot++;
+    if (!link_path_intersection(*link, path)) {
+      continue;
     }
+    affected_links.add(link);
   }
-  if (tot == 0) {
+
+  if (affected_links.is_empty()) {
     return OPERATOR_CANCELLED;
   }
 
-  /* Mute links. */
-  LISTBASE_FOREACH (bNodeLink *, link, &snode.edittree->links) {
-    if (node_link_is_hidden_or_dimmed(region.v2d, *link) || (link->flag & NODE_LINK_TEST)) {
-      continue;
-    }
+  bke::node_tree_runtime::AllowUsingOutdatedInfo allow_outdated_info{ntree};
 
-    if (link_path_intersection(*link, path)) {
-      nodeMuteLinkToggle(snode.edittree, link);
-    }
-  }
+  for (bNodeLink *link : affected_links) {
+    nodeLinkSetMute(&ntree, link, !(link->flag & NODE_LINK_MUTED));
+    const bool muted = link->flag & NODE_LINK_MUTED;
 
-  /* Clear remaining test flags. */
-  LISTBASE_FOREACH (bNodeLink *, link, &snode.edittree->links) {
-    if (node_link_is_hidden_or_dimmed(region.v2d, *link)) {
-      continue;
+    /* Propagate mute status downsteam past reroute nodes. */
+    if (link->tonode->is_reroute()) {
+      Stack<bNodeLink *> links;
+      links.push_multiple(link->tonode->output_sockets().first()->directly_linked_links());
+      while (!links.is_empty()) {
+        bNodeLink *link = links.pop();
+        nodeLinkSetMute(&ntree, link, muted);
+        if (!link->tonode->is_reroute()) {
+          continue;
+        }
+        links.push_multiple(link->tonode->output_sockets().first()->directly_linked_links());
+      }
+    }
+    /* Propagate mute status upstream past reroutes, but only if all outputs are muted. */
+    if (link->fromnode->is_reroute()) {
+      if (!muted || all_links_muted(*link->fromsock)) {
+        Stack<bNodeLink *> links;
+        links.push_multiple(link->fromnode->input_sockets().first()->directly_linked_links());
+        while (!links.is_empty()) {
+          bNodeLink *link = links.pop();
+          nodeLinkSetMute(&ntree, link, muted);
+          if (!link->fromnode->is_reroute()) {
+            continue;
+          }
+          if (!muted || all_links_muted(*link->fromsock)) {
+            links.push_multiple(link->fromnode->input_sockets().first()->directly_linked_links());
+          }
+        }
+      }
     }
-    link->flag &= ~NODE_LINK_TEST;
   }
 
-  ED_node_tree_propagate_change(C, CTX_data_main(C), snode.edittree);
+  ED_node_tree_propagate_change(C, CTX_data_main(C), &ntree);
   return OPERATOR_FINISHED;
 }