Geometry Nodes: Input data type utility function

This commit adds a simple utility function for getting the data type of an
attribute or its "constant" socket counterparts. No functional changes.

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

1 files changed, 50 insertions, 9 deletions

diff --git a/source/blender/nodes/intern/node_geometry_exec.cc b/source/blender/nodes/intern/node_geometry_exec.cc
index a6d9115f01f..6a22adb25a7 100644
--- a/source/blender/nodes/intern/node_geometry_exec.cc
+++ b/source/blender/nodes/intern/node_geometry_exec.cc
@@ -19,23 +19,31 @@
 
 namespace blender::nodes {
 
-ReadAttributePtr GeoNodeExecParams::get_input_attribute(const StringRef name,
-                                                        const GeometryComponent &component,
-                                                        const AttributeDomain domain,
-                                                        const CustomDataType type,
-                                                        const void *default_value) const
+const bNodeSocket *GeoNodeExecParams::find_available_socket(const StringRef name) const
 {
-  const bNodeSocket *found_socket = nullptr;
   LISTBASE_FOREACH (const bNodeSocket *, socket, &node_.inputs) {
     if ((socket->flag & SOCK_UNAVAIL) != 0) {
       continue;
     }
     if (name == socket->name) {
-      found_socket = socket;
-      break;
+      return socket;
     }
   }
-  BLI_assert(found_socket != nullptr);
+
+  return nullptr;
+}
+
+ReadAttributePtr GeoNodeExecParams::get_input_attribute(const StringRef name,
+                                                        const GeometryComponent &component,
+                                                        const AttributeDomain domain,
+                                                        const CustomDataType type,
+                                                        const void *default_value) const
+{
+  const bNodeSocket *found_socket = this->find_available_socket(name);
+  BLI_assert(found_socket != nullptr); /* There should always be available socket for the name. */
+  if (found_socket == nullptr) {
+    return component.attribute_get_constant_for_read(domain, type, default_value);
+  }
 
   if (found_socket->type == SOCK_STRING) {
     const std::string name = this->get_input<std::string>(found_socket->identifier);
@@ -60,6 +68,39 @@ ReadAttributePtr GeoNodeExecParams::get_input_attribute(const StringRef name,
   return component.attribute_get_constant_for_read(domain, type, default_value);
 }
 
+CustomDataType GeoNodeExecParams::get_input_attribute_data_type(
+    const StringRef name,
+    const GeometryComponent &component,
+    const CustomDataType default_type) const
+{
+  const bNodeSocket *found_socket = this->find_available_socket(name);
+  BLI_assert(found_socket != nullptr); /* There should always be available socket for the name. */
+  if (found_socket == nullptr) {
+    return default_type;
+  }
+
+  if (found_socket->type == SOCK_STRING) {
+    const std::string name = this->get_input<std::string>(found_socket->identifier);
+    ReadAttributePtr attribute = component.attribute_try_get_for_read(name);
+    if (!attribute) {
+      return default_type;
+    }
+    return attribute->custom_data_type();
+  }
+  if (found_socket->type == SOCK_FLOAT) {
+    return CD_PROP_FLOAT;
+  }
+  if (found_socket->type == SOCK_VECTOR) {
+    return CD_PROP_FLOAT3;
+  }
+  if (found_socket->type == SOCK_RGBA) {
+    return CD_PROP_COLOR;
+  }
+
+  BLI_assert(false);
+  return default_type;
+}
+
 void GeoNodeExecParams::check_extract_input(StringRef identifier,
                                             const CPPType *requested_type) const
 {