1 files changed, 32 insertions, 0 deletions
diff --git a/source/blender/nodes/geometry/node_geometry_tree.cc b/source/blender/nodes/geometry/node_geometry_tree.cc
index 2c4c7da64cc..f4cd00b88ed 100644
--- a/source/blender/nodes/geometry/node_geometry_tree.cc
+++ b/source/blender/nodes/geometry/node_geometry_tree.cc
@@ -67,6 +67,8 @@ static void geometry_node_tree_get_from_context(const bContext *C,
 
 static void geometry_node_tree_update(bNodeTree *ntree)
 {
+  ntreeSetOutput(ntree);
+
   /* Needed to give correct types to reroutes. */
   ntree_update_reroute_nodes(ntree);
 }
@@ -82,6 +84,34 @@ static void foreach_nodeclass(Scene *UNUSED(scene), void *calldata, bNodeClassCa
   func(calldata, NODE_CLASS_LAYOUT, N_("Layout"));
 }
 
+static bool geometry_node_tree_validate_link(bNodeTree *UNUSED(ntree), bNodeLink *link)
+{
+  /* Geometry, string, object, material, texture and collection sockets can only be connected to
+   * themselves. The other types can be converted between each other. */
+  if (ELEM(link->fromsock->type, SOCK_FLOAT, SOCK_VECTOR, SOCK_RGBA, SOCK_BOOLEAN, SOCK_INT) &&
+      ELEM(link->tosock->type, SOCK_FLOAT, SOCK_VECTOR, SOCK_RGBA, SOCK_BOOLEAN, SOCK_INT)) {
+    return true;
+  }
+  return (link->tosock->type == link->fromsock->type);
+}
+
+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,
+              SOCK_TEXTURE,
+              SOCK_MATERIAL);
+}
+
 void register_node_tree_type_geo(void)
 {
   bNodeTreeType *tt = ntreeType_Geometry = static_cast<bNodeTreeType *>(
@@ -95,6 +125,8 @@ 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;
+  tt->validate_link = geometry_node_tree_validate_link;
 
   ntreeTypeAdd(tt);
 }