Geometry Nodes: avoid data copy in store named attribute node

1 files changed, 25 insertions, 9 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_store_named_attribute.cc b/source/blender/nodes/geometry/nodes/node_geo_store_named_attribute.cc
index 5b0816fa3e3..1d1446ce1bd 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_store_named_attribute.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_store_named_attribute.cc
@@ -91,18 +91,34 @@ static void try_capture_field_on_geometry(GeometryComponent &component,
   const int domain_size = component.attribute_domain_size(domain);
   const IndexMask mask{IndexMask(domain_size)};
 
-  const CustomDataType data_type = bke::cpp_type_to_custom_data_type(field.cpp_type());
+  const CPPType &type = field.cpp_type();
+  const CustomDataType data_type = bke::cpp_type_to_custom_data_type(type);
+
+  /* Could avoid allocating a new buffer if:
+   * - We are writing to an attribute that exists already.
+   * - The field does not depend on that attribute (we can't easily check for that yet). */
+  void *buffer = MEM_mallocN(type.size() * domain_size, __func__);
 
-  /* Don't use #add_with_destination because the field might depend on an attribute
-   * with that name, and changing it as part of evaluation might affect the result. */
   fn::FieldEvaluator evaluator{field_context, &mask};
-  evaluator.add(field);
+  evaluator.add_with_destination(field, GMutableSpan{type, buffer, domain_size});
   evaluator.evaluate();
-  const GVArray &result = evaluator.get_evaluated(0);
-  OutputAttribute attribute = component.attribute_try_get_for_output_only(name, domain, data_type);
-  if (attribute) {
-    result.materialize(attribute.as_span().data());
-    attribute.save();
+
+  component.attribute_try_delete(name);
+  if (component.attribute_exists(name)) {
+    WriteAttributeLookup write_attribute = component.attribute_try_get_for_write(name);
+    if (write_attribute && write_attribute.domain == domain &&
+        write_attribute.varray.type() == type) {
+      write_attribute.varray.set_all(buffer);
+      write_attribute.tag_modified_fn();
+    }
+    else {
+      /* Cannot change type of built-in attribute. */
+    }
+    type.destruct_n(buffer, domain_size);
+    MEM_freeN(buffer);
+  }
+  else {
+    component.attribute_try_create(name, domain, data_type, AttributeInitMove{buffer});
   }
 }