Geometry Nodes: reduce code duplication with new GeometyrFieldInput

Most of our field inputs are currently specific to geometry. This patch introduces
a new `GeometryFieldInput` that reduces the overhead of adding new geometry
field input.

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

2 files changed, 56 insertions, 52 deletions

diff --git a/source/blender/blenkernel/BKE_geometry_set.hh b/source/blender/blenkernel/BKE_geometry_set.hh
index 84564ca5114..3fb8da0889b 100644
--- a/source/blender/blenkernel/BKE_geometry_set.hh
+++ b/source/blender/blenkernel/BKE_geometry_set.hh
@@ -744,13 +744,26 @@ class GeometryComponentFieldContext : public fn::FieldContext {
   }
 };
 
-class AttributeFieldInput : public fn::FieldInput {
+class GeometryFieldInput : public fn::FieldInput {
+ public:
+  using fn::FieldInput::FieldInput;
+
+  GVArray get_varray_for_context(const fn::FieldContext &context,
+                                 IndexMask mask,
+                                 ResourceScope &scope) const override;
+
+  virtual GVArray get_varray_for_context(const GeometryComponent &component,
+                                         const AttributeDomain domain,
+                                         IndexMask mask) const = 0;
+};
+
+class AttributeFieldInput : public GeometryFieldInput {
  private:
   std::string name_;
 
  public:
   AttributeFieldInput(std::string name, const CPPType &type)
-      : fn::FieldInput(type, name), name_(std::move(name))
+      : GeometryFieldInput(type, name), name_(std::move(name))
   {
     category_ = Category::NamedAttribute;
   }
@@ -767,9 +780,9 @@ class AttributeFieldInput : public fn::FieldInput {
     return name_;
   }
 
-  GVArray get_varray_for_context(const fn::FieldContext &context,
-                                 IndexMask mask,
-                                 ResourceScope &scope) const override;
+  GVArray get_varray_for_context(const GeometryComponent &component,
+                                 const AttributeDomain domain,
+                                 IndexMask mask) const override;
 
   std::string socket_inspection_name() const override;
 
@@ -777,16 +790,16 @@ class AttributeFieldInput : public fn::FieldInput {
   bool is_equal_to(const fn::FieldNode &other) const override;
 };
 
-class IDAttributeFieldInput : public fn::FieldInput {
+class IDAttributeFieldInput : public GeometryFieldInput {
  public:
-  IDAttributeFieldInput() : fn::FieldInput(CPPType::get<int>())
+  IDAttributeFieldInput() : GeometryFieldInput(CPPType::get<int>())
   {
     category_ = Category::Generated;
   }
 
-  GVArray get_varray_for_context(const fn::FieldContext &context,
-                                 IndexMask mask,
-                                 ResourceScope &scope) const override;
+  GVArray get_varray_for_context(const GeometryComponent &component,
+                                 const AttributeDomain domain,
+                                 IndexMask mask) const override;
 
   std::string socket_inspection_name() const override;
 
@@ -794,7 +807,7 @@ class IDAttributeFieldInput : public fn::FieldInput {
   bool is_equal_to(const fn::FieldNode &other) const override;
 };
 
-class AnonymousAttributeFieldInput : public fn::FieldInput {
+class AnonymousAttributeFieldInput : public GeometryFieldInput {
  private:
   /**
    * A strong reference is required to make sure that the referenced attribute is not removed
@@ -807,7 +820,7 @@ class AnonymousAttributeFieldInput : public fn::FieldInput {
   AnonymousAttributeFieldInput(StrongAnonymousAttributeID anonymous_id,
                                const CPPType &type,
                                std::string producer_name)
-      : fn::FieldInput(type, anonymous_id.debug_name()),
+      : GeometryFieldInput(type, anonymous_id.debug_name()),
         anonymous_id_(std::move(anonymous_id)),
         producer_name_(producer_name)
   {
@@ -823,9 +836,9 @@ class AnonymousAttributeFieldInput : public fn::FieldInput {
     return fn::Field<T>{field_input};
   }
 
-  GVArray get_varray_for_context(const fn::FieldContext &context,
-                                 IndexMask mask,
-                                 ResourceScope &scope) const override;
+  GVArray get_varray_for_context(const GeometryComponent &component,
+                                 const AttributeDomain domain,
+                                 IndexMask mask) const override;
 
   std::string socket_inspection_name() const override;
 
diff --git a/source/blender/blenkernel/intern/attribute_access.cc b/source/blender/blenkernel/intern/attribute_access.cc
index ef2505692ec..fe32279c431 100644
--- a/source/blender/blenkernel/intern/attribute_access.cc
+++ b/source/blender/blenkernel/intern/attribute_access.cc
@@ -1421,23 +1421,27 @@ OutputAttribute GeometryComponent::attribute_try_get_for_output_only(
 
 namespace blender::bke {
 
-GVArray AttributeFieldInput::get_varray_for_context(const fn::FieldContext &context,
-                                                    IndexMask UNUSED(mask),
-                                                    ResourceScope &UNUSED(scope)) const
+GVArray GeometryFieldInput::get_varray_for_context(const fn::FieldContext &context,
+                                                   IndexMask mask,
+                                                   ResourceScope &UNUSED(scope)) const
 {
   if (const GeometryComponentFieldContext *geometry_context =
           dynamic_cast<const GeometryComponentFieldContext *>(&context)) {
     const GeometryComponent &component = geometry_context->geometry_component();
     const AttributeDomain domain = geometry_context->domain();
-    const CustomDataType data_type = cpp_type_to_custom_data_type(*type_);
-    GVArray attribute = component.attribute_try_get_for_read(name_, domain, data_type);
-    if (attribute) {
-      return attribute;
-    }
+    return this->get_varray_for_context(component, domain, mask);
   }
   return {};
 }
 
+GVArray AttributeFieldInput::get_varray_for_context(const GeometryComponent &component,
+                                                    const AttributeDomain domain,
+                                                    IndexMask UNUSED(mask)) const
+{
+  const CustomDataType data_type = cpp_type_to_custom_data_type(*type_);
+  return component.attribute_try_get_for_read(name_, domain, data_type);
+}
+
 std::string AttributeFieldInput::socket_inspection_name() const
 {
   std::stringstream ss;
@@ -1469,25 +1473,20 @@ static StringRef get_random_id_attribute_name(const AttributeDomain domain)
   }
 }
 
-GVArray IDAttributeFieldInput::get_varray_for_context(const fn::FieldContext &context,
-                                                      IndexMask mask,
-                                                      ResourceScope &scope) const
+GVArray IDAttributeFieldInput::get_varray_for_context(const GeometryComponent &component,
+                                                      const AttributeDomain domain,
+                                                      IndexMask mask) const
 {
-  if (const GeometryComponentFieldContext *geometry_context =
-          dynamic_cast<const GeometryComponentFieldContext *>(&context)) {
-    const GeometryComponent &component = geometry_context->geometry_component();
-    const AttributeDomain domain = geometry_context->domain();
-    const StringRef name = get_random_id_attribute_name(domain);
-    GVArray attribute = component.attribute_try_get_for_read(name, domain, CD_PROP_INT32);
-    if (attribute) {
-      BLI_assert(attribute.size() == component.attribute_domain_size(domain));
-      return attribute;
-    }
 
-    /* Use the index as the fallback if no random ID attribute exists. */
-    return fn::IndexFieldInput::get_index_varray(mask, scope);
+  const StringRef name = get_random_id_attribute_name(domain);
+  GVArray attribute = component.attribute_try_get_for_read(name, domain, CD_PROP_INT32);
+  if (attribute) {
+    BLI_assert(attribute.size() == component.attribute_domain_size(domain));
+    return attribute;
   }
-  return {};
+
+  /* Use the index as the fallback if no random ID attribute exists. */
+  return fn::IndexFieldInput::get_index_varray(mask);
 }
 
 std::string IDAttributeFieldInput::socket_inspection_name() const
@@ -1507,20 +1506,12 @@ bool IDAttributeFieldInput::is_equal_to(const fn::FieldNode &other) const
   return dynamic_cast<const IDAttributeFieldInput *>(&other) != nullptr;
 }
 
-GVArray AnonymousAttributeFieldInput::get_varray_for_context(const fn::FieldContext &context,
-                                                             IndexMask UNUSED(mask),
-                                                             ResourceScope &UNUSED(scope)) const
+GVArray AnonymousAttributeFieldInput::get_varray_for_context(const GeometryComponent &component,
+                                                             const AttributeDomain domain,
+                                                             IndexMask UNUSED(mask)) const
 {
-  if (const GeometryComponentFieldContext *geometry_context =
-          dynamic_cast<const GeometryComponentFieldContext *>(&context)) {
-    const GeometryComponent &component = geometry_context->geometry_component();
-    const AttributeDomain domain = geometry_context->domain();
-    const CustomDataType data_type = cpp_type_to_custom_data_type(*type_);
-    GVArray attribute = component.attribute_try_get_for_read(
-        anonymous_id_.get(), domain, data_type);
-    return attribute;
-  }
-  return {};
+  const CustomDataType data_type = cpp_type_to_custom_data_type(*type_);
+  return component.attribute_try_get_for_read(anonymous_id_.get(), domain, data_type);
 }
 
 std::string AnonymousAttributeFieldInput::socket_inspection_name() const