Geometry Nodes: support reading from spans of WriteAttribute

Previously, the span returned by `WriteAttribute`s might not contain the
current value of the attribute for performance reasons. To avoid some
bugs, the span now always contains the old values (they might have to
be copied over from the internal storage, dependending on how the
attribute is stored).
The old behavior is still available with the `get_span_for_write_only`
method. The span that it returns might not contain the current
attribute values. Therefore, it should only be used when you want
to overwrite an attribute without looking at the old values.

source/blender/blenkernel/BKE_attribute_access.hh

@@ -169,8 +169,10 @@ class WriteAttribute {
   }
 
   /* Get a span that new attribute values can be written into. When all values have been changed,
-   * #apply_span has to be called. The span might not contain the original attribute values. */
+   * #apply_span has to be called. */
   fn::GMutableSpan get_span();
+  /* The span returned by this method might not contain the current attribute values. */
+  fn::GMutableSpan get_span_for_write_only();
   /* Write the changes to the span into the actual attribute, if they aren't already. */
   void apply_span();
 
@@ -178,7 +180,7 @@ class WriteAttribute {
   virtual void get_internal(const int64_t index, void *r_value) const = 0;
   virtual void set_internal(const int64_t index, const void *value) = 0;
 
-  virtual void initialize_span();
+  virtual void initialize_span(const bool write_only);
   virtual void apply_span_if_necessary();
 };
 
@@ -250,12 +252,16 @@ template<typename T> class TypedWriteAttribute {
   }
 
   /* Get a span that new values can be written into. Once all values have been updated #apply_span
-   * has to be called. The span might *not* contain the initial attribute values, so one should
-   * generally only write to the span. */
+   * has to be called. */
   MutableSpan<T> get_span()
   {
     return attribute_->get_span().typed<T>();
   }
+  /* The span returned by this method might not contain the current attribute values. */
+  MutableSpan<T> get_span_for_write_only()
+  {
+    return attribute_->get_span_for_write_only().typed<T>();
+  }
 
   /* Write back all changes to the actual attribute, if necessary. */
   void apply_span()

source/blender/blenkernel/intern/attribute_access.cc

@@ -97,26 +97,45 @@ WriteAttribute::~WriteAttribute()
 
 /**
  * Get a mutable span that can be modified. When all modifications to the attribute are done,
- * #apply_span_if_necessary should be called.
- */
+ * #apply_span should be called. */
 fn::GMutableSpan WriteAttribute::get_span()
 {
   if (size_ == 0) {
     return fn::GMutableSpan(cpp_type_);
   }
   if (array_buffer_ == nullptr) {
-    this->initialize_span();
+    this->initialize_span(false);
   }
   array_should_be_applied_ = true;
   return fn::GMutableSpan(cpp_type_, array_buffer_, size_);
 }
 
-void WriteAttribute::initialize_span()
+fn::GMutableSpan WriteAttribute::get_span_for_write_only()
 {
-  array_buffer_ = MEM_mallocN_aligned(cpp_type_.size() * size_, cpp_type_.alignment(), __func__);
+  if (size_ == 0) {
+    return fn::GMutableSpan(cpp_type_);
+  }
+  if (array_buffer_ == nullptr) {
+    this->initialize_span(true);
+  }
+  array_should_be_applied_ = true;
+  return fn::GMutableSpan(cpp_type_, array_buffer_, size_);
+}
+
+void WriteAttribute::initialize_span(const bool write_only)
+{
+  const int element_size = cpp_type_.size();
+  array_buffer_ = MEM_mallocN_aligned(element_size * size_, cpp_type_.alignment(), __func__);
   array_is_temporary_ = true;
-  /* This does nothing for trivial types, but is necessary for general correctness. */
-  cpp_type_.construct_default_n(array_buffer_, size_);
+  if (write_only) {
+    /* This does nothing for trivial types, but is necessary for general correctness. */
+    cpp_type_.construct_default_n(array_buffer_, size_);
+  }
+  else {
+    for (const int i : IndexRange(size_)) {
+      this->get(i, POINTER_OFFSET(array_buffer_, i * element_size));
+    }
+  }
 }
 
 void WriteAttribute::apply_span()
@@ -219,7 +238,7 @@ template<typename T> class ArrayWriteAttribute final : public WriteAttribute {
     data_[index] = *reinterpret_cast<const T *>(value);
   }
 
-  void initialize_span() override
+  void initialize_span(const bool UNUSED(write_only)) override
   {
     array_buffer_ = data_.data();
     array_is_temporary_ = false;

source/blender/nodes/geometry/nodes/node_geo_attribute_color_ramp.cc

@@ -56,7 +56,7 @@ static void execute_on_component(const GeoNodeExecParams &params, GeometryCompon
       input_name, result_domain, 0.0f);
 
   Span<float> data_in = attribute_in.get_span();
-  MutableSpan<Color4f> data_out = attribute_out.get_span();
+  MutableSpan<Color4f> data_out = attribute_out.get_span_for_write_only();
 
   ColorBand *color_ramp = &node_storage->color_ramp;
   for (const int i : data_in.index_range()) {

source/blender/nodes/geometry/nodes/node_geo_attribute_compare.cc

@@ -263,7 +263,7 @@ static void attribute_compare_calc(GeometryComponent &component, const GeoNodeEx
   }
 
   BooleanWriteAttribute attribute_result_bool = std::move(attribute_result);
-  MutableSpan<bool> result_span = attribute_result_bool.get_span();
+  MutableSpan<bool> result_span = attribute_result_bool.get_span_for_write_only();
 
   /* Use specific types for correct equality operations, but for other operations we use implicit
    * conversions and float comparison. In other words, the comparison is not element-wise. */

source/blender/nodes/geometry/nodes/node_geo_attribute_fill.cc

@@ -78,7 +78,7 @@ static void fill_attribute(GeometryComponent &component, const GeoNodeExecParams
     case CD_PROP_FLOAT: {
       FloatWriteAttribute float_attribute = std::move(attribute);
       const float value = params.get_input<float>("Value_001");
-      MutableSpan<float> attribute_span = float_attribute.get_span();
+      MutableSpan<float> attribute_span = float_attribute.get_span_for_write_only();
       attribute_span.fill(value);
       float_attribute.apply_span();
       break;
@@ -86,7 +86,7 @@ static void fill_attribute(GeometryComponent &component, const GeoNodeExecParams
     case CD_PROP_FLOAT3: {
       Float3WriteAttribute float3_attribute = std::move(attribute);
       const float3 value = params.get_input<float3>("Value");
-      MutableSpan<float3> attribute_span = float3_attribute.get_span();
+      MutableSpan<float3> attribute_span = float3_attribute.get_span_for_write_only();
       attribute_span.fill(value);
       float3_attribute.apply_span();
       break;
@@ -94,7 +94,7 @@ static void fill_attribute(GeometryComponent &component, const GeoNodeExecParams
     case CD_PROP_COLOR: {
       Color4fWriteAttribute color4f_attribute = std::move(attribute);
       const Color4f value = params.get_input<Color4f>("Value_002");
-      MutableSpan<Color4f> attribute_span = color4f_attribute.get_span();
+      MutableSpan<Color4f> attribute_span = color4f_attribute.get_span_for_write_only();
       attribute_span.fill(value);
       color4f_attribute.apply_span();
       break;
@@ -102,7 +102,7 @@ static void fill_attribute(GeometryComponent &component, const GeoNodeExecParams
     case CD_PROP_BOOL: {
       BooleanWriteAttribute boolean_attribute = std::move(attribute);
       const bool value = params.get_input<bool>("Value_003");
-      MutableSpan<bool> attribute_span = boolean_attribute.get_span();
+      MutableSpan<bool> attribute_span = boolean_attribute.get_span_for_write_only();
       attribute_span.fill(value);
       boolean_attribute.apply_span();
       break;

source/blender/nodes/geometry/nodes/node_geo_attribute_math.cc

@@ -75,7 +75,7 @@ static void do_math_operation(const FloatReadAttribute &input_a,
 
   Span<float> span_a = input_a.get_span();
   Span<float> span_b = input_b.get_span();
-  MutableSpan<float> span_result = result.get_span();
+  MutableSpan<float> span_result = result.get_span_for_write_only();
 
   bool success = try_dispatch_float_math_fl_fl_to_fl(
       operation, [&](auto math_function, const FloatMathOperationInfo &UNUSED(info)) {

source/blender/nodes/geometry/nodes/node_geo_attribute_vector_math.cc

@@ -149,7 +149,7 @@ static void do_math_operation_fl3_fl3_to_fl3(const Float3ReadAttribute &input_a,
 
   Span<float3> span_a = input_a.get_span();
   Span<float3> span_b = input_b.get_span();
-  MutableSpan<float3> span_result = result.get_span();
+  MutableSpan<float3> span_result = result.get_span_for_write_only();
 
   bool success = try_dispatch_float_math_fl3_fl3_to_fl3(
       operation, [&](auto math_function, const FloatMathOperationInfo &UNUSED(info)) {
@@ -179,7 +179,7 @@ static void do_math_operation_fl3_fl3_fl3_to_fl3(const Float3ReadAttribute &inpu
   Span<float3> span_a = input_a.get_span();
   Span<float3> span_b = input_b.get_span();
   Span<float3> span_c = input_c.get_span();
-  MutableSpan<float3> span_result = result.get_span();
+  MutableSpan<float3> span_result = result.get_span_for_write_only();
 
   bool success = try_dispatch_float_math_fl3_fl3_fl3_to_fl3(
       operation, [&](auto math_function, const FloatMathOperationInfo &UNUSED(info)) {
@@ -208,7 +208,7 @@ static void do_math_operation_fl3_fl3_to_fl(const Float3ReadAttribute &input_a,
 
   Span<float3> span_a = input_a.get_span();
   Span<float3> span_b = input_b.get_span();
-  MutableSpan<float> span_result = result.get_span();
+  MutableSpan<float> span_result = result.get_span_for_write_only();
 
   bool success = try_dispatch_float_math_fl3_fl3_to_fl(
       operation, [&](auto math_function, const FloatMathOperationInfo &UNUSED(info)) {
@@ -236,7 +236,7 @@ static void do_math_operation_fl3_fl_to_fl3(const Float3ReadAttribute &input_a,
 
   Span<float3> span_a = input_a.get_span();
   Span<float> span_b = input_b.get_span();
-  MutableSpan<float3> span_result = result.get_span();
+  MutableSpan<float3> span_result = result.get_span_for_write_only();
 
   bool success = try_dispatch_float_math_fl3_fl_to_fl3(
       operation, [&](auto math_function, const FloatMathOperationInfo &UNUSED(info)) {
@@ -262,7 +262,7 @@ static void do_math_operation_fl3_to_fl3(const Float3ReadAttribute &input_a,
   const int size = input_a.size();
 
   Span<float3> span_a = input_a.get_span();
-  MutableSpan<float3> span_result = result.get_span();
+  MutableSpan<float3> span_result = result.get_span_for_write_only();
 
   bool success = try_dispatch_float_math_fl3_to_fl3(
       operation, [&](auto math_function, const FloatMathOperationInfo &UNUSED(info)) {
@@ -287,7 +287,7 @@ static void do_math_operation_fl3_to_fl(const Float3ReadAttribute &input_a,
   const int size = input_a.size();
 
   Span<float3> span_a = input_a.get_span();
-  MutableSpan<float> span_result = result.get_span();
+  MutableSpan<float> span_result = result.get_span_for_write_only();
 
   bool success = try_dispatch_float_math_fl3_to_fl(
       operation, [&](auto math_function, const FloatMathOperationInfo &UNUSED(info)) {

source/blender/nodes/geometry/nodes/node_geo_join_geometry.cc

@@ -182,7 +182,7 @@ static void join_attributes(Span<const GeometryComponent *> src_components,
         write_attribute->domain() != domain) {
       continue;
     }
-    fn::GMutableSpan dst_span = write_attribute->get_span();
+    fn::GMutableSpan dst_span = write_attribute->get_span_for_write_only();
     fill_new_attribute(src_components, attribute_name, data_type, domain, dst_span);
     write_attribute->apply_span();
   }