Nodes: Add a callback to check for valid socket type

This adds a callback to bNodeTreeType to check which socket types are
valid for the tree type. Function has been implemented for the normal
tree types, and can be implemented for custom node trees with python,
by adding a `classmethod` to the tree. However, only builtin socket
types are supported.

This is relevant for T87049, but it also has the advantage that it is
now clear which node trees support which sockets. Previously this
was assumed to be known by all developers.

Differential Revision: https://developer.blender.org/D10938

4 files changed, 37 insertions, 0 deletions

diff --git a/source/blender/nodes/composite/node_composite_tree.c b/source/blender/nodes/composite/node_composite_tree.c
index 085b5c463b9..19815d01278 100644
--- a/source/blender/nodes/composite/node_composite_tree.c
+++ b/source/blender/nodes/composite/node_composite_tree.c
@@ -205,6 +205,12 @@ static void composite_node_add_init(bNodeTree *UNUSED(bnodetree), bNode *bnode)
   }
 }
 
+static bool composite_node_tree_socket_type_valid(eNodeSocketDatatype socket_type,
+                                                  bNodeTreeType *UNUSED(ntreetype))
+{
+  return ELEM(socket_type, SOCK_FLOAT, SOCK_VECTOR, SOCK_RGBA);
+}
+
 bNodeTreeType *ntreeType_Composite;
 
 void register_node_tree_type_cmp(void)
@@ -227,6 +233,7 @@ void register_node_tree_type_cmp(void)
   tt->update = update;
   tt->get_from_context = composite_get_from_context;
   tt->node_add_init = composite_node_add_init;
+  tt->valid_socket_type = composite_node_tree_socket_type_valid;
 
   tt->rna_ext.srna = &RNA_CompositorNodeTree;
 
diff --git a/source/blender/nodes/geometry/node_geometry_tree.cc b/source/blender/nodes/geometry/node_geometry_tree.cc
index 6d3b1d55005..afb86ff57b8 100644
--- a/source/blender/nodes/geometry/node_geometry_tree.cc
+++ b/source/blender/nodes/geometry/node_geometry_tree.cc
@@ -84,6 +84,21 @@ static void foreach_nodeclass(Scene *UNUSED(scene), void *calldata, bNodeClassCa
   func(calldata, NODE_CLASS_LAYOUT, N_("Layout"));
 }
 
+static bool geometry_node_tree_socket_type_valid(eNodeSocketDatatype socket_type,
+                                                 bNodeTreeType *UNUSED(ntreetype))
+{
+  return ELEM(socket_type,
+              SOCK_FLOAT,
+              SOCK_VECTOR,
+              SOCK_RGBA,
+              SOCK_BOOLEAN,
+              SOCK_INT,
+              SOCK_STRING,
+              SOCK_OBJECT,
+              SOCK_GEOMETRY,
+              SOCK_COLLECTION);
+}
+
 void register_node_tree_type_geo(void)
 {
   bNodeTreeType *tt = ntreeType_Geometry = static_cast<bNodeTreeType *>(
@@ -97,6 +112,7 @@ void register_node_tree_type_geo(void)
   tt->update = geometry_node_tree_update;
   tt->get_from_context = geometry_node_tree_get_from_context;
   tt->foreach_nodeclass = foreach_nodeclass;
+  tt->valid_socket_type = geometry_node_tree_socket_type_valid;
 
   ntreeTypeAdd(tt);
 }
diff --git a/source/blender/nodes/shader/node_shader_tree.c b/source/blender/nodes/shader/node_shader_tree.c
index 7b52b525541..25e4385de96 100644
--- a/source/blender/nodes/shader/node_shader_tree.c
+++ b/source/blender/nodes/shader/node_shader_tree.c
@@ -184,6 +184,12 @@ static bool shader_validate_link(bNodeTree *UNUSED(ntree), bNodeLink *link)
   return true;
 }
 
+static bool shader_node_tree_socket_type_valid(eNodeSocketDatatype socket_type,
+                                               bNodeTreeType *UNUSED(ntreetype))
+{
+  return ELEM(socket_type, SOCK_FLOAT, SOCK_VECTOR, SOCK_RGBA, SOCK_SHADER);
+}
+
 bNodeTreeType *ntreeType_Shader;
 
 void register_node_tree_type_sh(void)
@@ -205,6 +211,7 @@ void register_node_tree_type_sh(void)
   tt->poll = shader_tree_poll;
   tt->get_from_context = shader_get_from_context;
   tt->validate_link = shader_validate_link;
+  tt->valid_socket_type = shader_node_tree_socket_type_valid;
 
   tt->rna_ext.srna = &RNA_ShaderNodeTree;
 
diff --git a/source/blender/nodes/texture/node_texture_tree.c b/source/blender/nodes/texture/node_texture_tree.c
index 48027dc847b..2ae722e3cd8 100644
--- a/source/blender/nodes/texture/node_texture_tree.c
+++ b/source/blender/nodes/texture/node_texture_tree.c
@@ -152,6 +152,12 @@ static void update(bNodeTree *ntree)
   }
 }
 
+static bool texture_node_tree_socket_type_valid(eNodeSocketDatatype socket_type,
+                                                bNodeTreeType *UNUSED(ntreetype))
+{
+  return ELEM(socket_type, SOCK_FLOAT, SOCK_VECTOR, SOCK_RGBA);
+}
+
 bNodeTreeType *ntreeType_Texture;
 
 void register_node_tree_type_tex(void)
@@ -171,6 +177,7 @@ void register_node_tree_type_tex(void)
   tt->local_sync = local_sync;
   tt->local_merge = local_merge;
   tt->get_from_context = texture_get_from_context;
+  tt->valid_socket_type = texture_node_tree_socket_type_valid;
 
   tt->rna_ext.srna = &RNA_TextureNodeTree;