Nodes: store references to bNodeLinks in NodeTreeRef

Sometimes it is useful to have this information available to be able
to read information from links.

2 files changed, 65 insertions, 1 deletions

diff --git a/source/blender/nodes/NOD_node_tree_ref.hh b/source/blender/nodes/NOD_node_tree_ref.hh
index 7fcc117ba93..c1ff7bfe5e0 100644
--- a/source/blender/nodes/NOD_node_tree_ref.hh
+++ b/source/blender/nodes/NOD_node_tree_ref.hh
@@ -25,6 +25,7 @@
  * The following queries are supported efficiently:
  *  - socket -> index of socket
  *  - socket -> directly linked sockets
+ *  - socket -> directly linked links
  *  - socket -> linked sockets when skipping reroutes
  *  - socket -> node
  *  - socket/node -> rna pointer
@@ -65,6 +66,7 @@ class InputSocketRef;
 class OutputSocketRef;
 class NodeRef;
 class NodeTreeRef;
+class LinkRef;
 
 class SocketRef : NonCopyable, NonMovable {
  protected:
@@ -76,12 +78,14 @@ class SocketRef : NonCopyable, NonMovable {
   PointerRNA rna_;
   Vector<SocketRef *> linked_sockets_;
   Vector<SocketRef *> directly_linked_sockets_;
+  Vector<LinkRef *> directly_linked_links_;
 
   friend NodeTreeRef;
 
  public:
   Span<const SocketRef *> linked_sockets() const;
   Span<const SocketRef *> directly_linked_sockets() const;
+  Span<const LinkRef *> directly_linked_links() const;
   bool is_linked() const;
 
   const NodeRef &node() const;
@@ -156,6 +160,21 @@ class NodeRef : NonCopyable, NonMovable {
   bool is_muted() const;
 };
 
+class LinkRef : NonCopyable, NonMovable {
+ private:
+  OutputSocketRef *from_;
+  InputSocketRef *to_;
+  bNodeLink *blink_;
+
+  friend NodeTreeRef;
+
+ public:
+  const OutputSocketRef &from() const;
+  const InputSocketRef &to() const;
+
+  bNodeLink *blink() const;
+};
+
 class NodeTreeRef : NonCopyable, NonMovable {
  private:
   LinearAllocator<> allocator_;
@@ -164,6 +183,7 @@ class NodeTreeRef : NonCopyable, NonMovable {
   Vector<SocketRef *> sockets_by_id_;
   Vector<InputSocketRef *> input_sockets_;
   Vector<OutputSocketRef *> output_sockets_;
+  Vector<LinkRef *> links_;
   MultiValueMap<const bNodeType *, NodeRef *> nodes_by_type_;
 
  public:
@@ -178,6 +198,8 @@ class NodeTreeRef : NonCopyable, NonMovable {
   Span<const InputSocketRef *> input_sockets() const;
   Span<const OutputSocketRef *> output_sockets() const;
 
+  Span<const LinkRef *> links() const;
+
   bool has_link_cycles() const;
 
   bNodeTree *btree() const;
@@ -209,6 +231,11 @@ inline Span<const SocketRef *> SocketRef::directly_linked_sockets() const
   return directly_linked_sockets_;
 }
 
+inline Span<const LinkRef *> SocketRef::directly_linked_links() const
+{
+  return directly_linked_links_;
+}
+
 inline bool SocketRef::is_linked() const
 {
   return linked_sockets_.size() > 0;
@@ -409,7 +436,26 @@ inline bool NodeRef::is_muted() const
 }
 
 /* --------------------------------------------------------------------
- * NodeRef inline methods.
+ * LinkRef inline methods.
+ */
+
+inline const OutputSocketRef &LinkRef::from() const
+{
+  return *from_;
+}
+
+inline const InputSocketRef &LinkRef::to() const
+{
+  return *to_;
+}
+
+inline bNodeLink *LinkRef::blink() const
+{
+  return blink_;
+}
+
+/* --------------------------------------------------------------------
+ * NodeTreeRef inline methods.
  */
 
 inline Span<const NodeRef *> NodeTreeRef::nodes() const
@@ -443,6 +489,11 @@ inline Span<const OutputSocketRef *> NodeTreeRef::output_sockets() const
   return output_sockets_;
 }
 
+inline Span<const LinkRef *> NodeTreeRef::links() const
+{
+  return links_;
+}
+
 inline bNodeTree *NodeTreeRef::btree() const
 {
   return btree_;
diff --git a/source/blender/nodes/intern/node_tree_ref.cc b/source/blender/nodes/intern/node_tree_ref.cc
index 9dcd90f9f50..e174f096ff7 100644
--- a/source/blender/nodes/intern/node_tree_ref.cc
+++ b/source/blender/nodes/intern/node_tree_ref.cc
@@ -64,8 +64,16 @@ NodeTreeRef::NodeTreeRef(bNodeTree *btree) : btree_(btree)
     InputSocketRef &to_socket = this->find_input_socket(
         node_mapping, blink->tonode, blink->tosock);
 
+    LinkRef &link = *allocator_.construct<LinkRef>();
+    link.from_ = &from_socket;
+    link.to_ = &to_socket;
+    link.blink_ = blink;
+
+    links_.append(&link);
     from_socket.directly_linked_sockets_.append(&to_socket);
     to_socket.directly_linked_sockets_.append(&from_socket);
+    from_socket.directly_linked_links_.append(&link);
+    to_socket.directly_linked_links_.append(&link);
   }
 
   for (OutputSocketRef *socket : output_sockets_) {
@@ -85,6 +93,8 @@ NodeTreeRef::NodeTreeRef(bNodeTree *btree) : btree_(btree)
 
 NodeTreeRef::~NodeTreeRef()
 {
+  /* The destructor has to be called manually, because these types are allocated in a linear
+   * allocator. */
   for (NodeRef *node : nodes_by_id_) {
     node->~NodeRef();
   }
@@ -94,6 +104,9 @@ NodeTreeRef::~NodeTreeRef()
   for (OutputSocketRef *socket : output_sockets_) {
     socket->~OutputSocketRef();
   }
+  for (LinkRef *link : links_) {
+    link->~LinkRef();
+  }
 }
 
 InputSocketRef &NodeTreeRef::find_input_socket(Map<bNode *, NodeRef *> &node_mapping,