Nodes: make distinction between directly and logically linked sockets more clear

This also moves the handling of muted nodes from derived node tree to
node tree ref.

4 files changed, 87 insertions, 79 deletions

diff --git a/source/blender/nodes/NOD_derived_node_tree.hh b/source/blender/nodes/NOD_derived_node_tree.hh
index 3529336baf6..f1dbb2caf29 100644
--- a/source/blender/nodes/NOD_derived_node_tree.hh
+++ b/source/blender/nodes/NOD_derived_node_tree.hh
@@ -132,8 +132,7 @@ class DInputSocket : public DSocket {
   DOutputSocket get_corresponding_group_node_output() const;
   Vector<DOutputSocket, 4> get_corresponding_group_input_sockets() const;
 
-  void foreach_origin_socket(FunctionRef<void(DSocket)> callback,
-                             const bool follow_only_first_incoming_link = false) const;
+  void foreach_origin_socket(FunctionRef<void(DSocket)> callback) const;
 };
 
 /* A (nullable) reference to an output socket and the context it is in. */
diff --git a/source/blender/nodes/NOD_node_tree_ref.hh b/source/blender/nodes/NOD_node_tree_ref.hh
index b9c5d41658a..9302f14c3a4 100644
--- a/source/blender/nodes/NOD_node_tree_ref.hh
+++ b/source/blender/nodes/NOD_node_tree_ref.hh
@@ -79,17 +79,22 @@ class SocketRef : NonCopyable, NonMovable {
   int index_;
   PointerRNA rna_;
   Vector<LinkRef *> directly_linked_links_;
-  /* This is derived data that is cached for easy and fast access. */
+
+  /* These sockets are linked directly, i.e. with a single link inbetween. */
   MutableSpan<SocketRef *> directly_linked_sockets_;
-  MutableSpan<SocketRef *> linked_sockets_without_reroutes_and_muted_links_;
+  /* These sockets are linked when reroutes, muted links and muted nodes have been taken into
+   * account. */
+  MutableSpan<SocketRef *> logically_linked_sockets_;
 
   friend NodeTreeRef;
 
  public:
-  Span<const SocketRef *> linked_sockets() const;
+  Span<const SocketRef *> logically_linked_sockets() const;
   Span<const SocketRef *> directly_linked_sockets() const;
   Span<const LinkRef *> directly_linked_links() const;
-  bool is_linked() const;
+
+  bool is_directly_linked() const;
+  bool is_logically_linked() const;
 
   const NodeRef &node() const;
   const NodeTreeRef &tree() const;
@@ -123,7 +128,7 @@ class SocketRef : NonCopyable, NonMovable {
 
 class InputSocketRef final : public SocketRef {
  public:
-  Span<const OutputSocketRef *> linked_sockets() const;
+  Span<const OutputSocketRef *> logically_linked_sockets() const;
   Span<const OutputSocketRef *> directly_linked_sockets() const;
 
   bool is_multi_input_socket() const;
@@ -131,7 +136,7 @@ class InputSocketRef final : public SocketRef {
 
 class OutputSocketRef final : public SocketRef {
  public:
-  Span<const InputSocketRef *> linked_sockets() const;
+  Span<const InputSocketRef *> logically_linked_sockets() const;
   Span<const InputSocketRef *> directly_linked_sockets() const;
 };
 
@@ -250,10 +255,12 @@ class NodeTreeRef : NonCopyable, NonMovable {
                                       bNodeSocket *bsocket);
 
   void create_linked_socket_caches();
-  void foreach_origin_skipping_reroutes_and_muted_links(
-      InputSocketRef &socket, FunctionRef<void(OutputSocketRef &)> callback);
-  void foreach_target_skipping_reroutes_and_muted_links(
-      OutputSocketRef &socket, FunctionRef<void(InputSocketRef &)> callback);
+
+  void foreach_logical_origin(InputSocketRef &socket,
+                              FunctionRef<void(OutputSocketRef &)> callback,
+                              bool only_follow_first_input_link = false);
+  void foreach_logical_target(OutputSocketRef &socket,
+                              FunctionRef<void(InputSocketRef &)> callback);
 };
 
 using NodeTreeRefMap = Map<bNodeTree *, std::unique_ptr<const NodeTreeRef>>;
@@ -273,9 +280,9 @@ using nodes::SocketRef;
  * SocketRef inline methods.
  */
 
-inline Span<const SocketRef *> SocketRef::linked_sockets() const
+inline Span<const SocketRef *> SocketRef::logically_linked_sockets() const
 {
-  return linked_sockets_without_reroutes_and_muted_links_;
+  return logically_linked_sockets_;
 }
 
 inline Span<const SocketRef *> SocketRef::directly_linked_sockets() const
@@ -288,9 +295,14 @@ inline Span<const LinkRef *> SocketRef::directly_linked_links() const
   return directly_linked_links_;
 }
 
-inline bool SocketRef::is_linked() const
+inline bool SocketRef::is_directly_linked() const
+{
+  return directly_linked_sockets_.size() > 0;
+}
+
+inline bool SocketRef::is_logically_linked() const
 {
-  return linked_sockets_without_reroutes_and_muted_links_.size() > 0;
+  return logically_linked_sockets_.size() > 0;
 }
 
 inline const NodeRef &SocketRef::node() const
@@ -399,10 +411,9 @@ template<typename T> inline T *SocketRef::default_value() const
  * InputSocketRef inline methods.
  */
 
-inline Span<const OutputSocketRef *> InputSocketRef::linked_sockets() const
+inline Span<const OutputSocketRef *> InputSocketRef::logically_linked_sockets() const
 {
-  return linked_sockets_without_reroutes_and_muted_links_.as_span()
-      .cast<const OutputSocketRef *>();
+  return logically_linked_sockets_.as_span().cast<const OutputSocketRef *>();
 }
 
 inline Span<const OutputSocketRef *> InputSocketRef::directly_linked_sockets() const
@@ -419,9 +430,9 @@ inline bool InputSocketRef::is_multi_input_socket() const
  * OutputSocketRef inline methods.
  */
 
-inline Span<const InputSocketRef *> OutputSocketRef::linked_sockets() const
+inline Span<const InputSocketRef *> OutputSocketRef::logically_linked_sockets() const
 {
-  return linked_sockets_without_reroutes_and_muted_links_.as_span().cast<const InputSocketRef *>();
+  return logically_linked_sockets_.as_span().cast<const InputSocketRef *>();
 }
 
 inline Span<const InputSocketRef *> OutputSocketRef::directly_linked_sockets() const
diff --git a/source/blender/nodes/intern/derived_node_tree.cc b/source/blender/nodes/intern/derived_node_tree.cc
index 36c64b00f47..18605e93d94 100644
--- a/source/blender/nodes/intern/derived_node_tree.cc
+++ b/source/blender/nodes/intern/derived_node_tree.cc
@@ -168,35 +168,21 @@ DInputSocket DOutputSocket::get_active_corresponding_group_output_socket() const
 /* Call the given callback for every "real" origin socket. "Real" means that reroutes, muted nodes
  * and node groups are handled by this function. Origin sockets are ones where a node gets its
  * inputs from. */
-void DInputSocket::foreach_origin_socket(FunctionRef<void(DSocket)> callback,
-                                         const bool follow_only_first_incoming_link) const
+void DInputSocket::foreach_origin_socket(FunctionRef<void(DSocket)> callback) const
 {
   BLI_assert(*this);
-  Span<const OutputSocketRef *> linked_sockets_to_check = socket_ref_->as_input().linked_sockets();
-  if (follow_only_first_incoming_link) {
-    linked_sockets_to_check = linked_sockets_to_check.take_front(1);
-  }
-  for (const OutputSocketRef *linked_socket : linked_sockets_to_check) {
+  for (const OutputSocketRef *linked_socket : socket_ref_->as_input().logically_linked_sockets()) {
     const NodeRef &linked_node = linked_socket->node();
     DOutputSocket linked_dsocket{context_, linked_socket};
 
-    if (linked_node.is_muted()) {
-      /* If the node is muted, follow the internal links of the node. */
-      for (const InternalLinkRef *internal_link : linked_node.internal_links()) {
-        if (&internal_link->to() == linked_socket) {
-          DInputSocket input_of_muted_node{context_, &internal_link->from()};
-          input_of_muted_node.foreach_origin_socket(callback, true);
-        }
-      }
-    }
-    else if (linked_node.is_group_input_node()) {
+    if (linked_node.is_group_input_node()) {
       if (context_->is_root()) {
         /* This is a group input in the root node group. */
         callback(linked_dsocket);
       }
       else {
         DInputSocket socket_in_parent_group = linked_dsocket.get_corresponding_group_node_input();
-        if (socket_in_parent_group->is_linked()) {
+        if (socket_in_parent_group->is_logically_linked()) {
           /* Follow the links coming into the corresponding socket on the parent group node. */
           socket_in_parent_group.foreach_origin_socket(callback);
         }
@@ -210,7 +196,7 @@ void DInputSocket::foreach_origin_socket(FunctionRef<void(DSocket)> callback,
     else if (linked_node.is_group_node()) {
       DInputSocket socket_in_group = linked_dsocket.get_active_corresponding_group_output_socket();
       if (socket_in_group) {
-        if (socket_in_group->is_linked()) {
+        if (socket_in_group->is_logically_linked()) {
           /* Follow the links coming into the group output node of the child node group. */
           socket_in_group.foreach_origin_socket(callback);
         }
@@ -233,20 +219,11 @@ void DInputSocket::foreach_origin_socket(FunctionRef<void(DSocket)> callback,
  * from this socket.   */
 void DOutputSocket::foreach_target_socket(FunctionRef<void(DInputSocket)> callback) const
 {
-  for (const InputSocketRef *linked_socket : socket_ref_->as_output().linked_sockets()) {
+  for (const InputSocketRef *linked_socket : socket_ref_->as_output().logically_linked_sockets()) {
     const NodeRef &linked_node = linked_socket->node();
     DInputSocket linked_dsocket{context_, linked_socket};
 
-    if (linked_node.is_muted()) {
-      /* If the target node is muted, follow its internal links. */
-      for (const InternalLinkRef *internal_link : linked_node.internal_links()) {
-        if (&internal_link->from() == linked_socket) {
-          DOutputSocket output_of_muted_node{context_, &internal_link->to()};
-          output_of_muted_node.foreach_target_socket(callback);
-        }
-      }
-    }
-    else if (linked_node.is_group_output_node()) {
+    if (linked_node.is_group_output_node()) {
       if (context_->is_root()) {
         /* This is a group output in the root node group. */
         callback(linked_dsocket);
diff --git a/source/blender/nodes/intern/node_tree_ref.cc b/source/blender/nodes/intern/node_tree_ref.cc
index 4b326929dbb..3735cfe3247 100644
--- a/source/blender/nodes/intern/node_tree_ref.cc
+++ b/source/blender/nodes/intern/node_tree_ref.cc
@@ -170,16 +170,16 @@ void NodeTreeRef::create_linked_socket_caches()
     socket->directly_linked_sockets_ = allocator_.construct_array_copy(
         directly_linked_sockets.as_span());
 
-    /* Find linked sockets when skipping reroutes. */
-    Vector<SocketRef *> linked_sockets;
-    this->foreach_origin_skipping_reroutes_and_muted_links(
-        *socket, [&](OutputSocketRef &origin) { linked_sockets.append(&origin); });
-    if (linked_sockets == directly_linked_sockets) {
-      socket->linked_sockets_without_reroutes_and_muted_links_ = socket->directly_linked_sockets_;
+    /* Find logically linked sockets. */
+    Vector<SocketRef *> logically_linked_sockets;
+    this->foreach_logical_origin(
+        *socket, [&](OutputSocketRef &origin) { logically_linked_sockets.append(&origin); });
+    if (logically_linked_sockets == directly_linked_sockets) {
+      socket->logically_linked_sockets_ = socket->directly_linked_sockets_;
     }
     else {
-      socket->linked_sockets_without_reroutes_and_muted_links_ = allocator_.construct_array_copy(
-          linked_sockets.as_span());
+      socket->logically_linked_sockets_ = allocator_.construct_array_copy(
+          logically_linked_sockets.as_span());
     }
   }
 
@@ -192,51 +192,72 @@ void NodeTreeRef::create_linked_socket_caches()
     socket->directly_linked_sockets_ = allocator_.construct_array_copy(
         directly_linked_sockets.as_span());
 
-    /* Find linked sockets when skipping reroutes. */
-    Vector<SocketRef *> linked_sockets;
-    this->foreach_target_skipping_reroutes_and_muted_links(
-        *socket, [&](InputSocketRef &target) { linked_sockets.append(&target); });
-    if (linked_sockets == directly_linked_sockets) {
-      socket->linked_sockets_without_reroutes_and_muted_links_ = socket->directly_linked_sockets_;
+    /* Find logically linked sockets. */
+    Vector<SocketRef *> logically_linked_sockets;
+    this->foreach_logical_target(
+        *socket, [&](InputSocketRef &target) { logically_linked_sockets.append(&target); });
+    if (logically_linked_sockets == directly_linked_sockets) {
+      socket->logically_linked_sockets_ = socket->directly_linked_sockets_;
     }
     else {
-      socket->linked_sockets_without_reroutes_and_muted_links_ = allocator_.construct_array_copy(
-          linked_sockets.as_span());
+      socket->logically_linked_sockets_ = allocator_.construct_array_copy(
+          logically_linked_sockets.as_span());
     }
   }
 }
 
-void NodeTreeRef::foreach_origin_skipping_reroutes_and_muted_links(
-    InputSocketRef &socket, FunctionRef<void(OutputSocketRef &)> callback)
+void NodeTreeRef::foreach_logical_origin(InputSocketRef &socket,
+                                         FunctionRef<void(OutputSocketRef &)> callback,
+                                         bool only_follow_first_input_link)
 {
-  for (LinkRef *link : socket.directly_linked_links_) {
+  Span<LinkRef *> links_to_check = socket.directly_linked_links_;
+  if (only_follow_first_input_link) {
+    links_to_check = links_to_check.take_front(1);
+  }
+  for (LinkRef *link : links_to_check) {
     if (link->is_muted()) {
       continue;
     }
     OutputSocketRef *origin = link->from_;
-    if (origin->node_->is_reroute_node()) {
-      this->foreach_origin_skipping_reroutes_and_muted_links(*origin->node_->inputs_[0], callback);
+    NodeRef *origin_node = origin->node_;
+    if (origin_node->is_reroute_node()) {
+      this->foreach_logical_origin(*origin_node->inputs_[0], callback, false);
+    }
+    else if (origin_node->is_muted()) {
+      for (InternalLinkRef *internal_link : origin_node->internal_links_) {
+        if (internal_link->to_ == origin) {
+          this->foreach_logical_origin(*internal_link->from_, callback, true);
+          break;
+        }
+      }
     }
     else {
-      callback(*(OutputSocketRef *)origin);
+      callback(*origin);
     }
   }
 }
 
-void NodeTreeRef::foreach_target_skipping_reroutes_and_muted_links(
-    OutputSocketRef &socket, FunctionRef<void(InputSocketRef &)> callback)
+void NodeTreeRef::foreach_logical_target(OutputSocketRef &socket,
+                                         FunctionRef<void(InputSocketRef &)> callback)
 {
   for (LinkRef *link : socket.directly_linked_links_) {
     if (link->is_muted()) {
       continue;
     }
     InputSocketRef *target = link->to_;
-    if (target->node_->is_reroute_node()) {
-      this->foreach_target_skipping_reroutes_and_muted_links(*target->node_->outputs_[0],
-                                                             callback);
+    NodeRef *target_node = target->node_;
+    if (target_node->is_reroute_node()) {
+      this->foreach_logical_target(*target_node->outputs_[0], callback);
+    }
+    else if (target_node->is_muted()) {
+      for (InternalLinkRef *internal_link : target_node->internal_links_) {
+        if (internal_link->from_ == target) {
+          this->foreach_logical_target(*internal_link->to_, callback);
+        }
+      }
     }
     else {
-      callback(*(InputSocketRef *)target);
+      callback(*target);
     }
   }
 }