Fix: Attribute Transfer node does not work with a single index

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

1 files changed, 53 insertions, 63 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_transfer_attribute.cc b/source/blender/nodes/geometry/nodes/node_geo_transfer_attribute.cc
index a889678537f..f1b100c15d0 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_transfer_attribute.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_transfer_attribute.cc
@@ -305,7 +305,7 @@ void copy_with_indices(const VArray<T> &src,
 template<typename T>
 void copy_with_indices_clamped(const VArray<T> &src,
                                const IndexMask mask,
-                               const Span<int> indices,
+                               const VArray<int> &indices,
                                const MutableSpan<T> dst)
 {
   if (src.is_empty()) {
@@ -587,18 +587,11 @@ class NearestTransferFunction : public fn::MultiFunction {
   }
 };
 
-static const GeometryComponent *find_best_match_component(const GeometrySet &geometry,
-                                                          const GeometryComponentType type,
-                                                          const AttributeDomain domain)
+static const GeometryComponent *find_target_component(const GeometrySet &geometry,
+                                                      const AttributeDomain domain)
 {
-  /* Prefer transferring from the same component type, if it exists. */
-  if (component_is_available(geometry, type, domain)) {
-    return geometry.get_component_for_read(type);
-  }
-
-  /* If there is no component of the same type, choose the other component based on a consistent
-   * order, rather than some more complicated heuristic. This is the same order visible in the
-   * spreadsheet and used in the ray-cast node. */
+  /* Choose the other component based on a consistent order, rather than some more complicated
+   * heuristic. This is the same order visible in the spreadsheet and used in the ray-cast node. */
   static const Array<GeometryComponentType> supported_types = {
       GEO_COMPONENT_TYPE_MESH, GEO_COMPONENT_TYPE_POINT_CLOUD, GEO_COMPONENT_TYPE_CURVE};
   for (const GeometryComponentType src_type : supported_types) {
@@ -611,76 +604,71 @@ static const GeometryComponent *find_best_match_component(const GeometrySet &geo
 }
 
 /**
- * Use a #FieldInput because it's necessary to know the field context in order to choose the
- * corresponding component type from the input geometry, and only a #FieldInput receives the
- * evaluation context to provide its data.
- *
  * The index-based transfer theoretically does not need realized data when there is only one
  * instance geometry set in the target. A future optimization could be removing that limitation
  * internally.
  */
-class IndexTransferFieldInput : public FieldInput {
+class IndexTransferFunction : public fn::MultiFunction {
   GeometrySet src_geometry_;
   GField src_field_;
-  Field<int> index_field_;
   AttributeDomain domain_;
 
+  fn::MFSignature signature_;
+
+  std::optional<GeometryComponentFieldContext> geometry_context_;
+  std::unique_ptr<FieldEvaluator> evaluator_;
+  const GVArray *src_data_ = nullptr;
+
  public:
-  IndexTransferFieldInput(GeometrySet geometry,
-                          GField src_field,
-                          Field<int> index_field,
-                          const AttributeDomain domain)
-      : FieldInput(src_field.cpp_type(), "Attribute Transfer node"),
-        src_geometry_(std::move(geometry)),
-        src_field_(std::move(src_field)),
-        index_field_(std::move(index_field)),
-        domain_(domain)
+  IndexTransferFunction(GeometrySet geometry, GField src_field, const AttributeDomain domain)
+      : src_geometry_(std::move(geometry)), src_field_(std::move(src_field)), domain_(domain)
   {
     src_geometry_.ensure_owns_direct_data();
-    category_ = Category::Generated;
+
+    signature_ = this->create_signature();
+    this->set_signature(&signature_);
+
+    this->evaluate_field();
   }
 
-  const GVArray *get_varray_for_context(const FieldContext &context,
-                                        const IndexMask mask,
-                                        ResourceScope &scope) const final
+  fn::MFSignature create_signature()
   {
-    const GeometryComponentFieldContext *geometry_context =
-        dynamic_cast<const GeometryComponentFieldContext *>(&context);
-    if (geometry_context == nullptr) {
-      return nullptr;
-    }
-
-    FieldEvaluator index_evaluator{*geometry_context, &mask};
-    index_evaluator.add(index_field_);
-    index_evaluator.evaluate();
-    const VArray<int> &index_varray = index_evaluator.get_evaluated<int>(0);
-    /* The index virtual array is expected to be a span, since transferring the same index for
-     * every element is not very useful. */
-    VArray_Span<int> indices{index_varray};
+    fn::MFSignatureBuilder signature{"Attribute Transfer Index"};
+    signature.single_input<int>("Index");
+    signature.single_output("Attribute", src_field_.cpp_type());
+    return signature.build();
+  }
 
-    const GeometryComponent *component = find_best_match_component(
-        src_geometry_, geometry_context->geometry_component().type(), domain_);
+  void evaluate_field()
+  {
+    const GeometryComponent *component = find_target_component(src_geometry_, domain_);
     if (component == nullptr) {
-      return nullptr;
+      return;
     }
+    const int domain_size = component->attribute_domain_size(domain_);
+    geometry_context_.emplace(GeometryComponentFieldContext(*component, domain_));
+    evaluator_ = std::make_unique<FieldEvaluator>(*geometry_context_, domain_size);
+    evaluator_->add(src_field_);
+    evaluator_->evaluate();
+    src_data_ = &evaluator_->get_evaluated(0);
+  }
 
-    GeometryComponentFieldContext target_context{*component, domain_};
-    /* A potential improvement is to only copy the necessary values
-     * based on the indices retrieved from the index input field. */
-    FieldEvaluator target_evaluator{target_context, component->attribute_domain_size(domain_)};
-    target_evaluator.add(src_field_);
-    target_evaluator.evaluate();
-    const GVArray &src_data = target_evaluator.get_evaluated(0);
+  void call(IndexMask mask, fn::MFParams params, fn::MFContext UNUSED(context)) const override
+  {
+    const VArray<int> &indices = params.readonly_single_input<int>(0, "Index");
+    GMutableSpan dst = params.uninitialized_single_output(1, "Attribute");
 
-    GArray dst(src_field_.cpp_type(), mask.min_array_size());
+    const CPPType &type = dst.type();
+    if (src_data_ == nullptr) {
+      type.fill_construct_indices(type.default_value(), dst.data(), mask);
+      return;
+    }
 
-    attribute_math::convert_to_static_type(src_data.type(), [&](auto dummy) {
+    attribute_math::convert_to_static_type(type, [&](auto dummy) {
       using T = decltype(dummy);
-      GVArray_Typed<T> src{src_data};
-      copy_with_indices_clamped(*src, mask, indices, dst.as_mutable_span().typed<T>());
+      GVArray_Typed<T> src{*src_data_};
+      copy_with_indices_clamped(*src, mask, indices, dst.typed<T>());
     });
-
-    return &scope.construct<fn::GVArray_For_GArray>(std::move(dst));
   }
 };
 
@@ -792,9 +780,11 @@ static void geo_node_transfer_attribute_exec(GeoNodeExecParams params)
     }
     case GEO_NODE_ATTRIBUTE_TRANSFER_INDEX: {
       Field<int> indices = params.extract_input<Field<int>>("Index");
-      std::shared_ptr<FieldInput> input = std::make_shared<IndexTransferFieldInput>(
-          std::move(geometry), std::move(field), std::move(indices), domain);
-      output_field = GField(std::move(input));
+      auto fn = std::make_unique<IndexTransferFunction>(
+          std::move(geometry), std::move(field), domain);
+      auto op = std::make_shared<FieldOperation>(
+          FieldOperation(std::move(fn), {std::move(indices)}));
+      output_field = GField(std::move(op));
       break;
     }
   }