Geometry Nodes: Support for dynamic attributes on curve splines

With this patch you will be able to add and remove attributes from curve
data inside of geometry nodes. The following is currently implemented:
* Adding attributes with any data type to splines or spline points.
* Support for working with multiple splines at the same time.
* Interaction with the three builtin point attributes.
* Resampling attributes in the resample node.

The following is not implemented in this patch:
* Joining attributes when joining splines with the join geometry node.
* Domain interpolation between spline and point domains.
* More efficient ways to call attribute operations once per spline.

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

1 files changed, 290 insertions, 6 deletions

diff --git a/source/blender/blenkernel/intern/geometry_component_curve.cc b/source/blender/blenkernel/intern/geometry_component_curve.cc
index d6c7cae2727..73c9dae92bc 100644
--- a/source/blender/blenkernel/intern/geometry_component_curve.cc
+++ b/source/blender/blenkernel/intern/geometry_component_curve.cc
@@ -22,6 +22,12 @@
 
 #include "attribute_access_intern.hh"
 
+using blender::fn::GMutableSpan;
+using blender::fn::GSpan;
+using blender::fn::GVArray_For_GSpan;
+using blender::fn::GVArray_GSpan;
+using blender::fn::GVMutableArray_For_GMutableSpan;
+
 /* -------------------------------------------------------------------- */
 /** \name Geometry Component Implementation
  * \{ */
@@ -445,6 +451,20 @@ template<typename T> class VMutableArray_For_SplinePoints final : public VMutabl
   }
 };
 
+template<typename T> GVArrayPtr point_data_gvarray(Array<Span<T>> spans, Array<int> offsets)
+{
+  return std::make_unique<fn::GVArray_For_EmbeddedVArray<T, VArray_For_SplinePoints<T>>>(
+      offsets.last(), std::move(spans), std::move(offsets));
+}
+
+template<typename T>
+GVMutableArrayPtr point_data_gvarray(Array<MutableSpan<T>> spans, Array<int> offsets)
+{
+  return std::make_unique<
+      fn::GVMutableArray_For_EmbeddedVMutableArray<T, VMutableArray_For_SplinePoints<T>>>(
+      offsets.last(), std::move(spans), std::move(offsets));
+}
+
 /**
  * Virtual array implementation specifically for control point positions. This is only needed for
  * Bezier splines, where adjusting the position also requires adjusting handle positions depending
@@ -581,8 +601,7 @@ template<typename T> class BuiltinPointAttributeProvider : public BuiltinAttribu
       spans[i] = get_span_(*splines[i]);
     }
 
-    return std::make_unique<fn::GVArray_For_EmbeddedVArray<T, VArray_For_SplinePoints<T>>>(
-        offsets.last(), std::move(spans), std::move(offsets));
+    return point_data_gvarray(spans, offsets);
   }
 
   GVMutableArrayPtr try_get_for_write(GeometryComponent &component) const override
@@ -607,9 +626,7 @@ template<typename T> class BuiltinPointAttributeProvider : public BuiltinAttribu
       }
     }
 
-    return std::make_unique<
-        fn::GVMutableArray_For_EmbeddedVMutableArray<T, VMutableArray_For_SplinePoints<T>>>(
-        offsets.last(), std::move(spans), std::move(offsets));
+    return point_data_gvarray(spans, offsets);
   }
 
   bool try_delete(GeometryComponent &UNUSED(component)) const final
@@ -683,6 +700,256 @@ class PositionAttributeProvider final : public BuiltinPointAttributeProvider<flo
 /** \} */
 
 /* -------------------------------------------------------------------- */
+/** \name Dynamic Control Point Attributes
+ *
+ * The dynamic control point attribute implementation is very similar to the builtin attribute
+ * implementation-- it uses the same virtual array types. In order to work, this code depends on
+ * the fact that all a curve's splines will have the same attributes and they all have the same
+ * type.
+ * \{ */
+
+class DynamicPointAttributeProvider final : public DynamicAttributesProvider {
+ private:
+  static constexpr uint64_t supported_types_mask = CD_MASK_PROP_FLOAT | CD_MASK_PROP_FLOAT2 |
+                                                   CD_MASK_PROP_FLOAT3 | CD_MASK_PROP_INT32 |
+                                                   CD_MASK_PROP_COLOR | CD_MASK_PROP_BOOL;
+
+ public:
+  ReadAttributeLookup try_get_for_read(const GeometryComponent &component,
+                                       const StringRef attribute_name) const final
+  {
+    const CurveEval *curve = get_curve_from_component_for_read(component);
+    if (curve == nullptr || curve->splines().size() == 0) {
+      return {};
+    }
+
+    Span<SplinePtr> splines = curve->splines();
+    Vector<GSpan> spans; /* GSpan has no default constructor. */
+    spans.reserve(splines.size());
+    std::optional<GSpan> first_span = splines[0]->attributes.get_for_read(attribute_name);
+    if (!first_span) {
+      return {};
+    }
+    spans.append(*first_span);
+    for (const int i : IndexRange(1, splines.size() - 1)) {
+      std::optional<GSpan> span = splines[i]->attributes.get_for_read(attribute_name);
+      if (!span) {
+        /* All splines should have the same set of data layers. It would be possible to recover
+         * here and return partial data instead, but that would add a lot of complexity for a
+         * situation we don't even expect to encounter. */
+        BLI_assert_unreachable();
+        return {};
+      }
+      if (span->type() != spans.last().type()) {
+        /* Data layer types on separate splines do not match. */
+        BLI_assert_unreachable();
+        return {};
+      }
+      spans.append(*span);
+    }
+
+    /* First check for the simpler situation when we can return a simpler span virtual array. */
+    if (spans.size() == 1) {
+      return {std::make_unique<GVArray_For_GSpan>(spans.first()), ATTR_DOMAIN_POINT};
+    }
+
+    ReadAttributeLookup attribute = {};
+    Array<int> offsets = curve->control_point_offsets();
+    attribute_math::convert_to_static_type(spans[0].type(), [&](auto dummy) {
+      using T = decltype(dummy);
+      Array<Span<T>> data(splines.size());
+      for (const int i : splines.index_range()) {
+        data[i] = spans[i].typed<T>();
+        BLI_assert(data[i].data() != nullptr);
+      }
+      attribute = {point_data_gvarray(data, offsets), ATTR_DOMAIN_POINT};
+    });
+    return attribute;
+  }
+
+  /* This function is almost the same as #try_get_for_read, but without const. */
+  WriteAttributeLookup try_get_for_write(GeometryComponent &component,
+                                         const StringRef attribute_name) const final
+  {
+    CurveEval *curve = get_curve_from_component_for_write(component);
+    if (curve == nullptr || curve->splines().size() == 0) {
+      return {};
+    }
+
+    MutableSpan<SplinePtr> splines = curve->splines();
+    Vector<GMutableSpan> spans; /* GMutableSpan has no default constructor. */
+    spans.reserve(splines.size());
+    std::optional<GMutableSpan> first_span = splines[0]->attributes.get_for_write(attribute_name);
+    if (!first_span) {
+      return {};
+    }
+    spans.append(*first_span);
+    for (const int i : IndexRange(1, splines.size() - 1)) {
+      std::optional<GMutableSpan> span = splines[i]->attributes.get_for_write(attribute_name);
+      if (!span) {
+        /* All splines should have the same set of data layers. It would be possible to recover
+         * here and return partial data instead, but that would add a lot of complexity for a
+         * situation we don't even expect to encounter. */
+        BLI_assert_unreachable();
+        return {};
+      }
+      if (span->type() != spans.last().type()) {
+        /* Data layer types on separate splines do not match. */
+        BLI_assert_unreachable();
+        return {};
+      }
+      spans.append(*span);
+    }
+
+    /* First check for the simpler situation when we can return a simpler span virtual array. */
+    if (spans.size() == 1) {
+      return {std::make_unique<GVMutableArray_For_GMutableSpan>(spans.first()), ATTR_DOMAIN_POINT};
+    }
+
+    WriteAttributeLookup attribute = {};
+    Array<int> offsets = curve->control_point_offsets();
+    attribute_math::convert_to_static_type(spans[0].type(), [&](auto dummy) {
+      using T = decltype(dummy);
+      Array<MutableSpan<T>> data(splines.size());
+      for (const int i : splines.index_range()) {
+        data[i] = spans[i].typed<T>();
+        BLI_assert(data[i].data() != nullptr);
+      }
+      attribute = {point_data_gvarray(data, offsets), ATTR_DOMAIN_POINT};
+    });
+    return attribute;
+  }
+
+  bool try_delete(GeometryComponent &component, const StringRef attribute_name) const final
+  {
+    CurveEval *curve = get_curve_from_component_for_write(component);
+    if (curve == nullptr) {
+      return false;
+    }
+
+    bool layer_freed = false;
+    for (SplinePtr &spline : curve->splines()) {
+      spline->attributes.remove(attribute_name);
+    }
+    return layer_freed;
+  }
+
+  static GVArrayPtr varray_from_initializer(const AttributeInit &initializer,
+                                            const CustomDataType data_type,
+                                            const int total_size)
+  {
+    switch (initializer.type) {
+      case AttributeInit::Type::Default:
+        /* This function shouldn't be called in this case, since there
+         * is no need to copy anything to the new custom data array. */
+        BLI_assert_unreachable();
+        return {};
+      case AttributeInit::Type::VArray:
+        return static_cast<const AttributeInitVArray &>(initializer).varray->shallow_copy();
+      case AttributeInit::Type::MoveArray:
+        return std::make_unique<fn::GVArray_For_GSpan>(
+            GSpan(*bke::custom_data_type_to_cpp_type(data_type),
+                  static_cast<const AttributeInitMove &>(initializer).data,
+                  total_size));
+    }
+    BLI_assert_unreachable();
+    return {};
+  }
+
+  bool try_create(GeometryComponent &component,
+                  const StringRef attribute_name,
+                  const AttributeDomain domain,
+                  const CustomDataType data_type,
+                  const AttributeInit &initializer) const final
+  {
+    BLI_assert(this->type_is_supported(data_type));
+    if (domain != ATTR_DOMAIN_POINT) {
+      return false;
+    }
+    CurveEval *curve = get_curve_from_component_for_write(component);
+    if (curve == nullptr || curve->splines().size() == 0) {
+      return false;
+    }
+
+    MutableSpan<SplinePtr> splines = curve->splines();
+
+    /* First check the one case that allows us to avoid copying the input data. */
+    if (splines.size() == 1 && initializer.type == AttributeInit::Type::MoveArray) {
+      void *source_data = static_cast<const AttributeInitMove &>(initializer).data;
+      if (!splines[0]->attributes.create_by_move(attribute_name, data_type, source_data)) {
+        MEM_freeN(source_data);
+        return false;
+      }
+      return true;
+    }
+
+    /* Otherwise just create a custom data layer on each of the splines. */
+    for (const int i : splines.index_range()) {
+      if (!splines[i]->attributes.create(attribute_name, data_type)) {
+        /* If attribute creation fails on one of the splines, we cannot leave the custom data
+         * layers in the previous splines around, so delete them before returning. However,
+         * this is not an expected case. */
+        BLI_assert_unreachable();
+        return false;
+      }
+    }
+
+    /* With a default initializer type, we can keep the values at their initial values. */
+    if (initializer.type == AttributeInit::Type::Default) {
+      return true;
+    }
+
+    WriteAttributeLookup write_attribute = this->try_get_for_write(component, attribute_name);
+    /* We just created the attribute, it should exist. */
+    BLI_assert(write_attribute);
+
+    const int total_size = curve->control_point_offsets().last();
+    GVArrayPtr source_varray = varray_from_initializer(initializer, data_type, total_size);
+    /* TODO: When we can call a variant of #set_all with a virtual array argument,
+     * this theoretically unnecessary materialize step could be removed. */
+    GVArray_GSpan source_varray_span{*source_varray};
+    write_attribute.varray->set_all(source_varray_span.data());
+
+    if (initializer.type == AttributeInit::Type::MoveArray) {
+      MEM_freeN(static_cast<const AttributeInitMove &>(initializer).data);
+    }
+
+    return true;
+  }
+
+  bool foreach_attribute(const GeometryComponent &component,
+                         const AttributeForeachCallback callback) const final
+  {
+    const CurveEval *curve = get_curve_from_component_for_read(component);
+    if (curve == nullptr || curve->splines().size() == 0) {
+      return false;
+    }
+
+    Span<SplinePtr> splines = curve->splines();
+
+    /* In a debug build, check that all corresponding custom data layers have the same type. */
+    curve->assert_valid_point_attributes();
+
+    /* Use the first spline as a representative for all the others. */
+    splines.first()->attributes.foreach_attribute(callback, ATTR_DOMAIN_POINT);
+
+    return true;
+  }
+
+  void foreach_domain(const FunctionRef<void(AttributeDomain)> callback) const final
+  {
+    callback(ATTR_DOMAIN_POINT);
+  }
+
+  bool type_is_supported(CustomDataType data_type) const
+  {
+    return ((1ULL << data_type) & supported_types_mask) != 0;
+  }
+};
+
+/** \} */
+
+/* -------------------------------------------------------------------- */
 /** \name Attribute Provider Declaration
  * \{ */
 
@@ -704,6 +971,20 @@ static ComponentAttributeProviders create_attribute_providers_for_curve()
                                                make_cyclic_read_attribute,
                                                make_cyclic_write_attribute);
 
+  static CustomDataAccessInfo spline_custom_data_access = {
+      [](GeometryComponent &component) -> CustomData * {
+        CurveEval *curve = get_curve_from_component_for_write(component);
+        return curve ? &curve->attributes.data : nullptr;
+      },
+      [](const GeometryComponent &component) -> const CustomData * {
+        const CurveEval *curve = get_curve_from_component_for_read(component);
+        return curve ? &curve->attributes.data : nullptr;
+      },
+      nullptr};
+
+  static CustomDataAttributeProvider spline_custom_data(ATTR_DOMAIN_CURVE,
+                                                        spline_custom_data_access);
+
   static PositionAttributeProvider position;
 
   static BuiltinPointAttributeProvider<float> radius(
@@ -720,7 +1001,10 @@ static ComponentAttributeProviders create_attribute_providers_for_curve()
       [](Spline &spline) { return spline.tilts(); },
       [](Spline &spline) { spline.mark_cache_invalid(); });
 
-  return ComponentAttributeProviders({&position, &radius, &tilt, &resolution, &cyclic}, {});
+  static DynamicPointAttributeProvider point_custom_data;
+
+  return ComponentAttributeProviders({&position, &radius, &tilt, &resolution, &cyclic},
+                                     {&spline_custom_data, &point_custom_data});
 }
 
 }  // namespace blender::bke