1 files changed, 58 insertions, 11 deletions

diff --git a/source/blender/blenkernel/BKE_derived_node_tree.hh b/source/blender/blenkernel/BKE_derived_node_tree.hh
index 009d9b747f9..84b36f675e0 100644
--- a/source/blender/blenkernel/BKE_derived_node_tree.hh
+++ b/source/blender/blenkernel/BKE_derived_node_tree.hh
@@ -32,7 +32,8 @@
 
 #include "BKE_node_tree_ref.hh"
 
-namespace BKE {
+namespace blender {
+namespace bke {
 
 class DSocket;
 class DInputSocket;
@@ -42,7 +43,7 @@ class DParentNode;
 class DGroupInput;
 class DerivedNodeTree;
 
-class DSocket : blender::NonCopyable, blender::NonMovable {
+class DSocket : NonCopyable, NonMovable {
  protected:
   DNode *m_node;
   const SocketRef *m_socket_ref;
@@ -66,6 +67,11 @@ class DSocket : blender::NonCopyable, blender::NonMovable {
   PointerRNA *rna() const;
   StringRefNull idname() const;
   StringRefNull name() const;
+
+  const SocketRef &socket_ref() const;
+  bNodeSocket *bsocket() const;
+
+  bool is_available() const;
 };
 
 class DInputSocket : public DSocket {
@@ -95,7 +101,7 @@ class DOutputSocket : public DSocket {
   Span<const DInputSocket *> linked_sockets() const;
 };
 
-class DGroupInput : blender::NonCopyable, blender::NonMovable {
+class DGroupInput : NonCopyable, NonMovable {
  private:
   const InputSocketRef *m_socket_ref;
   DParentNode *m_parent;
@@ -106,13 +112,14 @@ class DGroupInput : blender::NonCopyable, blender::NonMovable {
 
  public:
   const InputSocketRef &socket_ref() const;
+  bNodeSocket *bsocket() const;
   const DParentNode *parent() const;
   Span<const DInputSocket *> linked_sockets() const;
   uint id() const;
   StringRefNull name() const;
 };
 
-class DNode : blender::NonCopyable, blender::NonMovable {
+class DNode : NonCopyable, NonMovable {
  private:
   const NodeRef *m_node_ref;
   DParentNode *m_parent;
@@ -144,7 +151,7 @@ class DNode : blender::NonCopyable, blender::NonMovable {
   void destruct_with_sockets();
 };
 
-class DParentNode : blender::NonCopyable, blender::NonMovable {
+class DParentNode : NonCopyable, NonMovable {
  private:
   const NodeRef *m_node_ref;
   DParentNode *m_parent;
@@ -160,7 +167,7 @@ class DParentNode : blender::NonCopyable, blender::NonMovable {
 
 using NodeTreeRefMap = Map<bNodeTree *, std::unique_ptr<const NodeTreeRef>>;
 
-class DerivedNodeTree : blender::NonCopyable, blender::NonMovable {
+class DerivedNodeTree : NonCopyable, NonMovable {
  private:
   LinearAllocator<> m_allocator;
   bNodeTree *m_btree;
@@ -172,19 +179,22 @@ class DerivedNodeTree : blender::NonCopyable, blender::NonMovable {
   Vector<DInputSocket *> m_input_sockets;
   Vector<DOutputSocket *> m_output_sockets;
 
-  Map<std::string, Vector<DNode *>> m_nodes_by_idname;
+  Map<const bNodeType *, Vector<DNode *>> m_nodes_by_type;
 
  public:
   DerivedNodeTree(bNodeTree *btree, NodeTreeRefMap &node_tree_refs);
   ~DerivedNodeTree();
 
   Span<const DNode *> nodes() const;
-  Span<const DNode *> nodes_with_idname(StringRef idname) const;
+  Span<const DNode *> nodes_by_type(StringRefNull idname) const;
+  Span<const DNode *> nodes_by_type(const bNodeType *nodetype) const;
 
   Span<const DSocket *> sockets() const;
   Span<const DInputSocket *> input_sockets() const;
   Span<const DOutputSocket *> output_sockets() const;
 
+  Span<const DGroupInput *> group_inputs() const;
+
   std::string to_dot() const;
 
  private:
@@ -278,6 +288,21 @@ inline StringRefNull DSocket::name() const
   return m_socket_ref->name();
 }
 
+inline const SocketRef &DSocket::socket_ref() const
+{
+  return *m_socket_ref;
+}
+
+inline bNodeSocket *DSocket::bsocket() const
+{
+  return m_socket_ref->bsocket();
+}
+
+inline bool DSocket::is_available() const
+{
+  return (m_socket_ref->bsocket()->flag & SOCK_UNAVAIL) == 0;
+}
+
 /* --------------------------------------------------------------------
  * DInputSocket inline methods.
  */
@@ -325,6 +350,11 @@ inline const InputSocketRef &DGroupInput::socket_ref() const
   return *m_socket_ref;
 }
 
+inline bNodeSocket *DGroupInput::bsocket() const
+{
+  return m_socket_ref->bsocket();
+}
+
 inline const DParentNode *DGroupInput::parent() const
 {
   return m_parent;
@@ -427,9 +457,15 @@ inline Span<const DNode *> DerivedNodeTree::nodes() const
   return m_nodes_by_id.as_span();
 }
 
-inline Span<const DNode *> DerivedNodeTree::nodes_with_idname(StringRef idname) const
+inline Span<const DNode *> DerivedNodeTree::nodes_by_type(StringRefNull idname) const
+{
+  const bNodeType *nodetype = nodeTypeFind(idname.data());
+  return this->nodes_by_type(nodetype);
+}
+
+inline Span<const DNode *> DerivedNodeTree::nodes_by_type(const bNodeType *nodetype) const
 {
-  const Vector<DNode *> *nodes = m_nodes_by_idname.lookup_ptr(idname);
+  const Vector<DNode *> *nodes = m_nodes_by_type.lookup_ptr(nodetype);
   if (nodes == nullptr) {
     return {};
   }
@@ -438,6 +474,11 @@ inline Span<const DNode *> DerivedNodeTree::nodes_with_idname(StringRef idname)
   }
 }
 
+inline Span<const DSocket *> DerivedNodeTree::sockets() const
+{
+  return m_sockets_by_id.as_span();
+}
+
 inline Span<const DInputSocket *> DerivedNodeTree::input_sockets() const
 {
   return m_input_sockets.as_span();
@@ -448,6 +489,12 @@ inline Span<const DOutputSocket *> DerivedNodeTree::output_sockets() const
   return m_output_sockets.as_span();
 }
 
-}  // namespace BKE
+inline Span<const DGroupInput *> DerivedNodeTree::group_inputs() const
+{
+  return m_group_inputs.as_span();
+}
+
+}  // namespace bke
+}  // namespace blender
 
 #endif /* __BKE_DERIVED_NODE_TREE_HH__ */