Geometry Nodes: simplify supporting different input socket types for attributes

This is a non-functional change. The functionality introduced in this commit
is not used in master yet. It is used by nodes that are being developed in
other branches though.

4 files changed, 88 insertions, 0 deletions

diff --git a/source/blender/nodes/NOD_geometry_exec.hh b/source/blender/nodes/NOD_geometry_exec.hh
index fde576d7429..a7df4bc3e1b 100644
--- a/source/blender/nodes/NOD_geometry_exec.hh
+++ b/source/blender/nodes/NOD_geometry_exec.hh
@@ -148,6 +148,26 @@ class GeoNodeExecParams {
     return self_object_;
   }
 
+  /**
+   * Creates a read-only attribute based on node inputs. The method automatically detects which
+   * input with the given name is available.
+   */
+  ReadAttributePtr get_input_attribute(const StringRef name,
+                                       const GeometryComponent &component,
+                                       const AttributeDomain domain,
+                                       const CustomDataType type,
+                                       const void *default_value) const;
+
+  template<typename T>
+  bke::TypedReadAttribute<T> get_input_attribute(const StringRef name,
+                                                 const GeometryComponent &component,
+                                                 const AttributeDomain domain,
+                                                 const T &default_value) const
+  {
+    const CustomDataType type = bke::cpp_type_to_custom_data_type(CPPType::get<T>());
+    return this->get_input_attribute(name, component, domain, type, &default_value);
+  }
+
  private:
   /* Utilities for detecting common errors at when using this class. */
   void check_extract_input(StringRef identifier, const CPPType *requested_type = nullptr) const;
diff --git a/source/blender/nodes/geometry/node_geometry_util.cc b/source/blender/nodes/geometry/node_geometry_util.cc
index 41bdb1cfff0..34c7d224f03 100644
--- a/source/blender/nodes/geometry/node_geometry_util.cc
+++ b/source/blender/nodes/geometry/node_geometry_util.cc
@@ -17,6 +17,27 @@
 #include "node_geometry_util.hh"
 #include "node_util.h"
 
+namespace blender::nodes {
+
+void update_attribute_input_socket_availabilities(bNode &node,
+                                                  const StringRef name,
+                                                  const GeometryNodeAttributeInputMode mode)
+{
+  const GeometryNodeAttributeInputMode mode_ = (GeometryNodeAttributeInputMode)mode;
+  LISTBASE_FOREACH (bNodeSocket *, socket, &node.inputs) {
+    if (name == socket->name) {
+      const bool is_available =
+          ((socket->type == SOCK_STRING && mode_ == GEO_NODE_ATTRIBUTE_INPUT_ATTRIBUTE) ||
+           (socket->type == SOCK_FLOAT && mode_ == GEO_NODE_ATTRIBUTE_INPUT_FLOAT) ||
+           (socket->type == SOCK_VECTOR && mode_ == GEO_NODE_ATTRIBUTE_INPUT_VECTOR) ||
+           (socket->type == SOCK_RGBA && mode_ == GEO_NODE_ATTRIBUTE_INPUT_COLOR));
+      nodeSetSocketAvailability(socket, is_available);
+    }
+  }
+}
+
+}  // namespace blender::nodes
+
 bool geo_node_poll_default(bNodeType *UNUSED(ntype), bNodeTree *ntree)
 {
   return STREQ(ntree->idname, "GeometryNodeTree");
diff --git a/source/blender/nodes/geometry/node_geometry_util.hh b/source/blender/nodes/geometry/node_geometry_util.hh
index bb26763642b..ec389961615 100644
--- a/source/blender/nodes/geometry/node_geometry_util.hh
+++ b/source/blender/nodes/geometry/node_geometry_util.hh
@@ -37,3 +37,9 @@
 void geo_node_type_base(
     struct bNodeType *ntype, int type, const char *name, short nclass, short flag);
 bool geo_node_poll_default(struct bNodeType *ntype, struct bNodeTree *ntree);
+
+namespace blender::nodes {
+void update_attribute_input_socket_availabilities(bNode &node,
+                                                  const StringRef name,
+                                                  const GeometryNodeAttributeInputMode mode);
+}
diff --git a/source/blender/nodes/intern/node_geometry_exec.cc b/source/blender/nodes/intern/node_geometry_exec.cc
index 50292cb8cfb..a6d9115f01f 100644
--- a/source/blender/nodes/intern/node_geometry_exec.cc
+++ b/source/blender/nodes/intern/node_geometry_exec.cc
@@ -19,6 +19,47 @@
 
 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 *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;
+    }
+  }
+  BLI_assert(found_socket != nullptr);
+
+  if (found_socket->type == SOCK_STRING) {
+    const std::string name = this->get_input<std::string>(found_socket->identifier);
+    return component.attribute_get_for_read(name, domain, type, default_value);
+  }
+  if (found_socket->type == SOCK_FLOAT) {
+    const float value = this->get_input<float>(found_socket->identifier);
+    return component.attribute_get_constant_for_read_converted(
+        domain, CD_PROP_FLOAT, type, &value);
+  }
+  if (found_socket->type == SOCK_VECTOR) {
+    const float3 value = this->get_input<float3>(found_socket->identifier);
+    return component.attribute_get_constant_for_read_converted(
+        domain, CD_PROP_FLOAT3, type, &value);
+  }
+  if (found_socket->type == SOCK_RGBA) {
+    const Color4f value = this->get_input<Color4f>(found_socket->identifier);
+    return component.attribute_get_constant_for_read_converted(
+        domain, CD_PROP_COLOR, type, &value);
+  }
+  BLI_assert(false);
+  return component.attribute_get_constant_for_read(domain, type, default_value);
+}
+
 void GeoNodeExecParams::check_extract_input(StringRef identifier,
                                             const CPPType *requested_type) const
 {