Nodes: improve NodeTreeRef

This adds a couple more utility methods to various node tree ref types.
Also `InternalLinkRef` has been added to get simpler access to internal links.

4 files changed, 107 insertions, 10 deletions

diff --git a/source/blender/nodes/NOD_derived_node_tree.hh b/source/blender/nodes/NOD_derived_node_tree.hh
index 62affe43895..ea67f23eade 100644
--- a/source/blender/nodes/NOD_derived_node_tree.hh
+++ b/source/blender/nodes/NOD_derived_node_tree.hh
@@ -174,8 +174,6 @@ class DParentNode : NonCopyable, NonMovable {
   int id() const;
 };
 
-using NodeTreeRefMap = Map<bNodeTree *, std::unique_ptr<const NodeTreeRef>>;
-
 class DerivedNodeTree : NonCopyable, NonMovable {
  private:
   LinearAllocator<> allocator_;
diff --git a/source/blender/nodes/NOD_node_tree_ref.hh b/source/blender/nodes/NOD_node_tree_ref.hh
index c1ff7bfe5e0..3c898aaf0ff 100644
--- a/source/blender/nodes/NOD_node_tree_ref.hh
+++ b/source/blender/nodes/NOD_node_tree_ref.hh
@@ -67,6 +67,7 @@ class OutputSocketRef;
 class NodeRef;
 class NodeTreeRef;
 class LinkRef;
+class InternalLinkRef;
 
 class SocketRef : NonCopyable, NonMovable {
  protected:
@@ -106,16 +107,21 @@ class SocketRef : NonCopyable, NonMovable {
   StringRefNull idname() const;
   StringRefNull name() const;
   StringRefNull identifier() const;
+  bNodeSocketType *typeinfo() const;
 
   bNodeSocket *bsocket() const;
   bNode *bnode() const;
   bNodeTree *btree() const;
+
+  bool is_available() const;
 };
 
 class InputSocketRef final : public SocketRef {
  public:
   Span<const OutputSocketRef *> linked_sockets() const;
   Span<const OutputSocketRef *> directly_linked_sockets() const;
+
+  bool is_multi_input_socket() const;
 };
 
 class OutputSocketRef final : public SocketRef {
@@ -132,6 +138,7 @@ class NodeRef : NonCopyable, NonMovable {
   int id_;
   Vector<InputSocketRef *> inputs_;
   Vector<OutputSocketRef *> outputs_;
+  Vector<InternalLinkRef *> internal_links_;
 
   friend NodeTreeRef;
 
@@ -140,6 +147,7 @@ class NodeRef : NonCopyable, NonMovable {
 
   Span<const InputSocketRef *> inputs() const;
   Span<const OutputSocketRef *> outputs() const;
+  Span<const InternalLinkRef *> internal_links() const;
 
   const InputSocketRef &input(int index) const;
   const OutputSocketRef &output(int index) const;
@@ -150,6 +158,7 @@ class NodeRef : NonCopyable, NonMovable {
   PointerRNA *rna() const;
   StringRefNull idname() const;
   StringRefNull name() const;
+  bNodeType *typeinfo() const;
 
   int id() const;
 
@@ -175,6 +184,21 @@ class LinkRef : NonCopyable, NonMovable {
   bNodeLink *blink() const;
 };
 
+class InternalLinkRef : NonCopyable, NonMovable {
+ private:
+  InputSocketRef *from_;
+  OutputSocketRef *to_;
+  bNodeLink *blink_;
+
+  friend NodeTreeRef;
+
+ public:
+  const InputSocketRef &from() const;
+  const OutputSocketRef &to() const;
+
+  bNodeLink *blink() const;
+};
+
 class NodeTreeRef : NonCopyable, NonMovable {
  private:
   LinearAllocator<> allocator_;
@@ -217,6 +241,19 @@ class NodeTreeRef : NonCopyable, NonMovable {
   void find_targets_skipping_reroutes(OutputSocketRef &socket_ref, Vector<SocketRef *> &r_targets);
 };
 
+using NodeTreeRefMap = Map<bNodeTree *, std::unique_ptr<const NodeTreeRef>>;
+
+const NodeTreeRef &get_tree_ref_from_map(NodeTreeRefMap &node_tree_refs, bNodeTree &btree);
+
+namespace node_tree_ref_types {
+using nodes::InputSocketRef;
+using nodes::NodeRef;
+using nodes::NodeTreeRef;
+using nodes::NodeTreeRefMap;
+using nodes::OutputSocketRef;
+using nodes::SocketRef;
+}  // namespace node_tree_ref_types
+
 /* --------------------------------------------------------------------
  * SocketRef inline methods.
  */
@@ -308,6 +345,11 @@ inline StringRefNull SocketRef::identifier() const
   return bsocket_->identifier;
 }
 
+inline bNodeSocketType *SocketRef::typeinfo() const
+{
+  return bsocket_->typeinfo;
+}
+
 inline bNodeSocket *SocketRef::bsocket() const
 {
   return bsocket_;
@@ -323,6 +365,11 @@ inline bNodeTree *SocketRef::btree() const
   return node_->btree();
 }
 
+inline bool SocketRef::is_available() const
+{
+  return (bsocket_->flag & SOCK_UNAVAIL) == 0;
+}
+
 /* --------------------------------------------------------------------
  * InputSocketRef inline methods.
  */
@@ -337,6 +384,11 @@ inline Span<const OutputSocketRef *> InputSocketRef::directly_linked_sockets() c
   return directly_linked_sockets_.as_span().cast<const OutputSocketRef *>();
 }
 
+inline bool InputSocketRef::is_multi_input_socket() const
+{
+  return bsocket_->flag & SOCK_MULTI_INPUT;
+}
+
 /* --------------------------------------------------------------------
  * OutputSocketRef inline methods.
  */
@@ -370,6 +422,11 @@ inline Span<const OutputSocketRef *> NodeRef::outputs() const
   return outputs_;
 }
 
+inline Span<const InternalLinkRef *> NodeRef::internal_links() const
+{
+  return internal_links_;
+}
+
 inline const InputSocketRef &NodeRef::input(int index) const
 {
   return *inputs_[index];
@@ -405,6 +462,11 @@ inline StringRefNull NodeRef::name() const
   return bnode_->name;
 }
 
+inline bNodeType *NodeRef::typeinfo() const
+{
+  return bnode_->typeinfo;
+}
+
 inline int NodeRef::id() const
 {
   return id_;
@@ -455,6 +517,25 @@ inline bNodeLink *LinkRef::blink() const
 }
 
 /* --------------------------------------------------------------------
+ * InternalLinkRef inline methods.
+ */
+
+inline const InputSocketRef &InternalLinkRef::from() const
+{
+  return *from_;
+}
+
+inline const OutputSocketRef &InternalLinkRef::to() const
+{
+  return *to_;
+}
+
+inline bNodeLink *InternalLinkRef::blink() const
+{
+  return blink_;
+}
+
+/* --------------------------------------------------------------------
  * NodeTreeRef inline methods.
  */
 
diff --git a/source/blender/nodes/intern/derived_node_tree.cc b/source/blender/nodes/intern/derived_node_tree.cc
index 7934104f84b..a65641bca2a 100644
--- a/source/blender/nodes/intern/derived_node_tree.cc
+++ b/source/blender/nodes/intern/derived_node_tree.cc
@@ -22,17 +22,11 @@
 
 namespace blender::nodes {
 
-static const NodeTreeRef &get_tree_ref(NodeTreeRefMap &node_tree_refs, bNodeTree *btree)
-{
-  return *node_tree_refs.lookup_or_add_cb(btree,
-                                          [&]() { return std::make_unique<NodeTreeRef>(btree); });
-}
-
 DerivedNodeTree::DerivedNodeTree(bNodeTree *btree, NodeTreeRefMap &node_tree_refs) : btree_(btree)
 {
   BLI_assert(btree != nullptr);
 
-  const NodeTreeRef &main_tree_ref = get_tree_ref(node_tree_refs, btree);
+  const NodeTreeRef &main_tree_ref = get_tree_ref_from_map(node_tree_refs, *btree);
   used_node_tree_refs_.add_new(&main_tree_ref);
 
   Vector<DNode *> all_nodes;
@@ -144,7 +138,7 @@ BLI_NOINLINE void DerivedNodeTree::expand_group_node(DNode &group_node,
     return;
   }
 
-  const NodeTreeRef &group_ref = get_tree_ref(node_tree_refs, btree);
+  const NodeTreeRef &group_ref = get_tree_ref_from_map(node_tree_refs, *btree);
   used_node_tree_refs_.add(&group_ref);
 
   DParentNode &parent = *allocator_.construct<DParentNode>();
diff --git a/source/blender/nodes/intern/node_tree_ref.cc b/source/blender/nodes/intern/node_tree_ref.cc
index e174f096ff7..7fe21ec8582 100644
--- a/source/blender/nodes/intern/node_tree_ref.cc
+++ b/source/blender/nodes/intern/node_tree_ref.cc
@@ -52,6 +52,24 @@ NodeTreeRef::NodeTreeRef(bNodeTree *btree) : btree_(btree)
       RNA_pointer_create(&btree->id, &RNA_NodeSocket, bsocket, &socket.rna_);
     }
 
+    LISTBASE_FOREACH (bNodeLink *, blink, &bnode->internal_links) {
+      InternalLinkRef &internal_link = *allocator_.construct<InternalLinkRef>();
+      internal_link.blink_ = blink;
+      for (InputSocketRef *socket_ref : node.inputs_) {
+        if (socket_ref->bsocket_ == blink->fromsock) {
+          internal_link.from_ = socket_ref;
+          break;
+        }
+      }
+      for (OutputSocketRef *socket_ref : node.outputs_) {
+        if (socket_ref->bsocket_ == blink->tosock) {
+          internal_link.to_ = socket_ref;
+          break;
+        }
+      }
+      node.internal_links_.append(&internal_link);
+    }
+
     input_sockets_.extend(node.inputs_.as_span());
     output_sockets_.extend(node.outputs_.as_span());
 
@@ -229,4 +247,10 @@ std::string NodeTreeRef::to_dot() const
   return digraph.to_dot_string();
 }
 
+const NodeTreeRef &get_tree_ref_from_map(NodeTreeRefMap &node_tree_refs, bNodeTree &btree)
+{
+  return *node_tree_refs.lookup_or_add_cb(&btree,
+                                          [&]() { return std::make_unique<NodeTreeRef>(&btree); });
+}
+
 }  // namespace blender::nodes