Nodes: store socket declaration reference in socket

Previously, to get the declaration of a socket, one had to go
through `node->declaration`. Now this indirection is not necessary
anymore. This makes it easier to add more per-socket information
into the declaration and accessing it in various places.

Currently, this system is used by socket descriptions and node warnings
for unsupported geometry component types.

8 files changed, 70 insertions, 47 deletions

diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index 58fea6d462c..645b4410623 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -735,7 +735,9 @@ void nodeSetSocketAvailability(struct bNodeSocket *sock, bool is_available);
 
 int nodeSocketLinkLimit(const struct bNodeSocket *sock);
 
-void nodeDeclarationEnsure(struct bNodeTree *ntree, struct bNode *node);
+bool nodeDeclarationEnsure(struct bNodeTree *ntree, struct bNode *node);
+bool nodeDeclarationEnsureOnOutdatedNode(struct bNodeTree *ntree, struct bNode *node);
+void nodeSocketDeclarationsUpdate(struct bNode *node);
 
 /* Node Clipboard */
 void BKE_node_clipboard_init(const struct bNodeTree *ntree);
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index bfb2491803d..60a9ce48a0a 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -681,6 +681,7 @@ static void direct_link_node_socket(BlendDataReader *reader, bNodeSocket *sock)
   BLO_read_data_address(reader, &sock->default_value);
   sock->total_inputs = 0; /* Clear runtime data set before drawing. */
   sock->cache = nullptr;
+  sock->declaration = nullptr;
 }
 
 /* ntree itself has been read! */
@@ -1070,8 +1071,7 @@ IDTypeInfo IDType_ID_NT = {
 static void node_add_sockets_from_type(bNodeTree *ntree, bNode *node, bNodeType *ntype)
 {
   if (ntype->declare != nullptr) {
-    nodeDeclarationEnsure(ntree, node);
-    node->declaration->build(*ntree, *node);
+    node_verify_sockets(ntree, node, true);
     return;
   }
   bNodeSocketTemplate *sockdef;
@@ -4018,17 +4018,38 @@ int nodeSocketLinkLimit(const bNodeSocket *sock)
   return sock->limit;
 }
 
+static void update_socket_declarations(ListBase *sockets,
+                                       Span<blender::nodes::SocketDeclarationPtr> declarations)
+{
+  int index;
+  LISTBASE_FOREACH_INDEX (bNodeSocket *, socket, sockets, index) {
+    const SocketDeclaration &socket_decl = *declarations[index];
+    socket->declaration = &socket_decl;
+  }
+}
+
 /**
- * If the node implements a `declare` function, this function makes sure that `node->declaration`
- * is up to date.
+ * Update `socket->declaration` for all sockets in the node. This assumes that the node declaration
+ * and sockets are up to date already.
+ */
+void nodeSocketDeclarationsUpdate(bNode *node)
+{
+  BLI_assert(node->declaration != nullptr);
+  update_socket_declarations(&node->inputs, node->declaration->inputs());
+  update_socket_declarations(&node->outputs, node->declaration->outputs());
+}
+
+/**
+ * Just update `node->declaration` if necessary. This can also be called on nodes that may not be
+ * up to date (e.g. because the need versioning or are dynamic).
  */
-void nodeDeclarationEnsure(bNodeTree *UNUSED(ntree), bNode *node)
+bool nodeDeclarationEnsureOnOutdatedNode(bNodeTree *UNUSED(ntree), bNode *node)
 {
   if (node->declaration != nullptr) {
-    return;
+    return false;
   }
   if (node->typeinfo->declare == nullptr) {
-    return;
+    return false;
   }
   if (node->typeinfo->declaration_is_dynamic) {
     node->declaration = new blender::nodes::NodeDeclaration();
@@ -4040,6 +4061,20 @@ void nodeDeclarationEnsure(bNodeTree *UNUSED(ntree), bNode *node)
     BLI_assert(node->typeinfo->fixed_declaration != nullptr);
     node->declaration = node->typeinfo->fixed_declaration;
   }
+  return true;
+}
+
+/**
+ * If the node implements a `declare` function, this function makes sure that `node->declaration`
+ * is up to date. It is expected that the sockets of the node are up to date already.
+ */
+bool nodeDeclarationEnsure(bNodeTree *ntree, bNode *node)
+{
+  if (nodeDeclarationEnsureOnOutdatedNode(ntree, node)) {
+    nodeSocketDeclarationsUpdate(node);
+    return true;
+  }
+  return false;
 }
 
 /* ************** Node Clipboard *********** */
diff --git a/source/blender/editors/space_node/node_draw.cc b/source/blender/editors/space_node/node_draw.cc
index b6c24a55a78..5fdf816339b 100644
--- a/source/blender/editors/space_node/node_draw.cc
+++ b/source/blender/editors/space_node/node_draw.cc
@@ -1107,21 +1107,8 @@ static void node_socket_draw_nested(const bContext *C,
             C, *data->ntree, *data->node, *data->socket);
 
         std::stringstream output;
-        if (data->node->declaration != nullptr) {
-          ListBase *list;
-          Span<blender::nodes::SocketDeclarationPtr> decl_list;
-
-          if (data->socket->in_out == SOCK_IN) {
-            list = &data->node->inputs;
-            decl_list = data->node->declaration->inputs();
-          }
-          else {
-            list = &data->node->outputs;
-            decl_list = data->node->declaration->outputs();
-          }
-
-          const int socket_index = BLI_findindex(list, data->socket);
-          const blender::nodes::SocketDeclaration &socket_decl = *decl_list[socket_index];
+        if (data->socket->declaration != nullptr) {
+          const blender::nodes::SocketDeclaration &socket_decl = *data->socket->declaration;
           blender::StringRef description = socket_decl.description();
           if (!description.is_empty()) {
             output << TIP_(description.data()) << ".\n\n";
diff --git a/source/blender/makesdna/DNA_node_types.h b/source/blender/makesdna/DNA_node_types.h
index 94e4cc1323a..79287885dec 100644
--- a/source/blender/makesdna/DNA_node_types.h
+++ b/source/blender/makesdna/DNA_node_types.h
@@ -81,6 +81,19 @@ typedef struct bNodeStack {
 #define NS_CR_FIT 4
 #define NS_CR_STRETCH 5
 
+/** Workaround to forward-declare C++ type in C header. */
+#ifdef __cplusplus
+namespace blender::nodes {
+class NodeDeclaration;
+class SocketDeclaration;
+}  // namespace blender::nodes
+using NodeDeclarationHandle = blender::nodes::NodeDeclaration;
+using SocketDeclarationHandle = blender::nodes::SocketDeclaration;
+#else
+typedef struct NodeDeclarationHandle NodeDeclarationHandle;
+typedef struct SocketDeclarationHandle SocketDeclarationHandle;
+#endif
+
 typedef struct bNodeSocket {
   struct bNodeSocket *next, *prev, *new_sock;
 
@@ -153,6 +166,12 @@ typedef struct bNodeSocket {
    * kept for forward compatibility */
   /** Custom data for inputs, only UI writes in this. */
   bNodeStack ns DNA_DEPRECATED;
+
+  /**
+   * References a socket declaration that is owned by `node->declaration`. This is only runtime
+   * data. It has to be updated when the node declaration changes.
+   */
+  const SocketDeclarationHandle *declaration;
 } bNodeSocket;
 
 /* sock->type */
@@ -220,16 +239,6 @@ typedef enum eNodeSocketFlag {
   SOCK_HIDE_LABEL = (1 << 12),
 } eNodeSocketFlag;
 
-/** Workaround to forward-declare C++ type in C header. */
-#ifdef __cplusplus
-namespace blender::nodes {
-class NodeDeclaration;
-}
-using NodeDeclarationHandle = blender::nodes::NodeDeclaration;
-#else
-typedef struct NodeDeclarationHandle NodeDeclarationHandle;
-#endif
-
 /* TODO: Limit data in bNode to what we want to see saved. */
 typedef struct bNode {
   struct bNode *next, *prev, *new_node;
diff --git a/source/blender/nodes/NOD_node_declaration.hh b/source/blender/nodes/NOD_node_declaration.hh
index e9f2996bb30..9b99026d6a7 100644
--- a/source/blender/nodes/NOD_node_declaration.hh
+++ b/source/blender/nodes/NOD_node_declaration.hh
@@ -228,7 +228,6 @@ class NodeDeclaration {
   friend NodeDeclarationBuilder;
 
  public:
-  void build(bNodeTree &ntree, bNode &node) const;
   bool matches(const bNode &node) const;
 
   Span<SocketDeclarationPtr> inputs() const;
diff --git a/source/blender/nodes/intern/node_declaration.cc b/source/blender/nodes/intern/node_declaration.cc
index 8a38b68ec59..e804d10ad75 100644
--- a/source/blender/nodes/intern/node_declaration.cc
+++ b/source/blender/nodes/intern/node_declaration.cc
@@ -20,16 +20,6 @@
 
 namespace blender::nodes {
 
-void NodeDeclaration::build(bNodeTree &ntree, bNode &node) const
-{
-  for (const SocketDeclarationPtr &decl : inputs_) {
-    decl->build(ntree, node, SOCK_IN);
-  }
-  for (const SocketDeclarationPtr &decl : outputs_) {
-    decl->build(ntree, node, SOCK_OUT);
-  }
-}
-
 bool NodeDeclaration::matches(const bNode &node) const
 {
   auto check_sockets = [&](ListBase sockets, Span<SocketDeclarationPtr> socket_decls) {
diff --git a/source/blender/nodes/intern/node_geometry_exec.cc b/source/blender/nodes/intern/node_geometry_exec.cc
index c7a3e795c33..27bc206187d 100644
--- a/source/blender/nodes/intern/node_geometry_exec.cc
+++ b/source/blender/nodes/intern/node_geometry_exec.cc
@@ -39,8 +39,8 @@ void GeoNodeExecParams::error_message_add(const NodeWarningType type, std::strin
 void GeoNodeExecParams::check_input_geometry_set(StringRef identifier,
                                                  const GeometrySet &geometry_set) const
 {
-  const int input_index = provider_->dnode->input_by_identifier(identifier).index();
-  const SocketDeclaration &decl = *provider_->dnode->declaration()->inputs()[input_index];
+  const SocketDeclaration &decl =
+      *provider_->dnode->input_by_identifier(identifier).bsocket()->declaration;
   const decl::Geometry *geo_decl = dynamic_cast<const decl::Geometry *>(&decl);
   if (geo_decl == nullptr) {
     return;
diff --git a/source/blender/nodes/intern/node_socket.cc b/source/blender/nodes/intern/node_socket.cc
index 11356178d87..dce54d58dce 100644
--- a/source/blender/nodes/intern/node_socket.cc
+++ b/source/blender/nodes/intern/node_socket.cc
@@ -269,10 +269,11 @@ void node_verify_sockets(bNodeTree *ntree, bNode *node, bool do_id_user)
     return;
   }
   if (ntype->declare != nullptr) {
-    nodeDeclarationEnsure(ntree, node);
+    nodeDeclarationEnsureOnOutdatedNode(ntree, node);
     if (!node->declaration->matches(*node)) {
       refresh_node(*ntree, *node, *node->declaration, do_id_user);
     }
+    nodeSocketDeclarationsUpdate(node);
     return;
   }
   /* Don't try to match socket lists when there are no templates.