Geometry Nodes: Add initializer for attribute creation

Previously we always had to set attribute values after creating
the attribute. This patch adds an initializer argument to
`attribute_try_create` which can fill it in a few ways, which
are explained in code comments.

This fixes T87597.

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

5 files changed, 188 insertions, 26 deletions

diff --git a/source/blender/blenkernel/BKE_geometry_set.hh b/source/blender/blenkernel/BKE_geometry_set.hh
index e9545cba1f1..027338a5d5c 100644
--- a/source/blender/blenkernel/BKE_geometry_set.hh
+++ b/source/blender/blenkernel/BKE_geometry_set.hh
@@ -70,6 +70,65 @@ using AttributeForeachCallback = blender::FunctionRef<bool(blender::StringRefNul
                                                            const AttributeMetaData &meta_data)>;
 
 /**
+ * Base class for the attribute intializer types described below.
+ */
+struct AttributeInit {
+  enum class Type {
+    Default,
+    VArray,
+    MoveArray,
+  };
+  Type type;
+  AttributeInit(const Type type) : type(type)
+  {
+  }
+};
+
+/**
+ * Create an attribute using the default value for the data type.
+ * The default values may depend on the attribute provider implementation.
+ */
+struct AttributeInitDefault : public AttributeInit {
+  AttributeInitDefault() : AttributeInit(Type::Default)
+  {
+  }
+};
+
+/**
+ * Create an attribute by copying data from an existing virtual array. The virtual array
+ * must have the same type as the newly created attribute.
+ *
+ * Note that this can be used to fill the new attribute with the default
+ */
+struct AttributeInitVArray : public AttributeInit {
+  const blender::fn::GVArray *varray;
+
+  AttributeInitVArray(const blender::fn::GVArray *varray)
+      : AttributeInit(Type::VArray), varray(varray)
+  {
+  }
+};
+
+/**
+ * Create an attribute with a by passing ownership of a pre-allocated contiguous array of data.
+ * Sometimes data is created before a geometry component is available. In that case, it's
+ * preferable to move data directly to the created attribute to avoid a new allocation and a copy.
+ *
+ * Note that this will only have a benefit for attributes that are stored directly as contigious
+ * arrays, so not for some built-in attributes.
+ *
+ * The array must be allocated with MEM_*, since `attribute_try_create` will free the array if it
+ * can't be used directly, and that is generally how Blender expects custom data to be allocated.
+ */
+struct AttributeInitMove : public AttributeInit {
+  void *data = nullptr;
+
+  AttributeInitMove(void *data) : AttributeInit(Type::MoveArray), data(data)
+  {
+  }
+};
+
+/**
  * This is the base class for specialized geometry component types.
  */
 class GeometryComponent {
@@ -137,11 +196,13 @@ class GeometryComponent {
   /* Returns true when the attribute has been created. */
   bool attribute_try_create(const blender::StringRef attribute_name,
                             const AttributeDomain domain,
-                            const CustomDataType data_type);
+                            const CustomDataType data_type,
+                            const AttributeInit &initializer);
 
   /* Try to create the builtin attribute with the given name. No data type or domain has to be
    * provided, because those are fixed for builtin attributes. */
-  bool attribute_try_create_builtin(const blender::StringRef attribute_name);
+  bool attribute_try_create_builtin(const blender::StringRef attribute_name,
+                                    const AttributeInit &initializer);
 
   blender::Set<std::string> attribute_names() const;
   bool attribute_foreach(const AttributeForeachCallback callback) const;
diff --git a/source/blender/blenkernel/intern/attribute_access.cc b/source/blender/blenkernel/intern/attribute_access.cc
index 571763e3719..6c37d34dc9b 100644
--- a/source/blender/blenkernel/intern/attribute_access.cc
+++ b/source/blender/blenkernel/intern/attribute_access.cc
@@ -254,7 +254,43 @@ bool BuiltinCustomDataLayerProvider::try_delete(GeometryComponent &component) co
   return delete_success;
 }
 
-bool BuiltinCustomDataLayerProvider::try_create(GeometryComponent &component) const
+static bool add_custom_data_layer_from_attribute_init(CustomData &custom_data,
+                                                      const CustomDataType data_type,
+                                                      const int domain_size,
+                                                      const AttributeInit &initializer)
+{
+  switch (initializer.type) {
+    case AttributeInit::Type::Default: {
+      void *data = CustomData_add_layer(&custom_data, data_type, CD_DEFAULT, nullptr, domain_size);
+      return data != nullptr;
+    }
+    case AttributeInit::Type::VArray: {
+      void *data = CustomData_add_layer(&custom_data, data_type, CD_DEFAULT, nullptr, domain_size);
+      if (data == nullptr) {
+        return false;
+      }
+      const GVArray *varray = static_cast<const AttributeInitVArray &>(initializer).varray;
+      varray->materialize_to_uninitialized(IndexRange(varray->size()), data);
+      return true;
+    }
+    case AttributeInit::Type::MoveArray: {
+      void *source_data = static_cast<const AttributeInitMove &>(initializer).data;
+      void *data = CustomData_add_layer(
+          &custom_data, data_type, CD_ASSIGN, source_data, domain_size);
+      if (data == nullptr) {
+        MEM_freeN(source_data);
+        return false;
+      }
+      return true;
+    }
+  }
+
+  BLI_assert_unreachable();
+  return false;
+}
+
+bool BuiltinCustomDataLayerProvider::try_create(GeometryComponent &component,
+                                                const AttributeInit &initializer) const
 {
   if (createable_ != Creatable) {
     return false;
@@ -267,10 +303,10 @@ bool BuiltinCustomDataLayerProvider::try_create(GeometryComponent &component) co
     /* Exists already. */
     return false;
   }
+
   const int domain_size = component.attribute_domain_size(domain_);
-  const void *data = CustomData_add_layer(
-      custom_data, stored_type_, CD_DEFAULT, nullptr, domain_size);
-  const bool success = data != nullptr;
+  const bool success = add_custom_data_layer_from_attribute_init(
+      *custom_data, stored_type_, domain_size, initializer);
   if (success) {
     custom_data_access_.update_custom_data_pointers(component);
   }
@@ -372,10 +408,52 @@ bool CustomDataAttributeProvider::try_delete(GeometryComponent &component,
   return false;
 }
 
+static bool add_named_custom_data_layer_from_attribute_init(const StringRef attribute_name,
+                                                            CustomData &custom_data,
+                                                            const CustomDataType data_type,
+                                                            const int domain_size,
+                                                            const AttributeInit &initializer)
+{
+  char attribute_name_c[MAX_NAME];
+  attribute_name.copy(attribute_name_c);
+
+  switch (initializer.type) {
+    case AttributeInit::Type::Default: {
+      void *data = CustomData_add_layer_named(
+          &custom_data, data_type, CD_DEFAULT, nullptr, domain_size, attribute_name_c);
+      return data != nullptr;
+    }
+    case AttributeInit::Type::VArray: {
+      void *data = CustomData_add_layer_named(
+          &custom_data, data_type, CD_DEFAULT, nullptr, domain_size, attribute_name_c);
+      if (data == nullptr) {
+        return false;
+      }
+      const GVArray *varray = static_cast<const AttributeInitVArray &>(initializer).varray;
+      varray->materialize_to_uninitialized(IndexRange(varray->size()), data);
+      return true;
+    }
+    case AttributeInit::Type::MoveArray: {
+      void *source_data = static_cast<const AttributeInitMove &>(initializer).data;
+      void *data = CustomData_add_layer_named(
+          &custom_data, data_type, CD_ASSIGN, source_data, domain_size, attribute_name_c);
+      if (data == nullptr) {
+        MEM_freeN(source_data);
+        return false;
+      }
+      return true;
+    }
+  }
+
+  BLI_assert_unreachable();
+  return false;
+}
+
 bool CustomDataAttributeProvider::try_create(GeometryComponent &component,
                                              const StringRef attribute_name,
                                              const AttributeDomain domain,
-                                             const CustomDataType data_type) const
+                                             const CustomDataType data_type,
+                                             const AttributeInit &initializer) const
 {
   if (domain_ != domain) {
     return false;
@@ -393,10 +471,8 @@ bool CustomDataAttributeProvider::try_create(GeometryComponent &component,
     }
   }
   const int domain_size = component.attribute_domain_size(domain_);
-  char attribute_name_c[MAX_NAME];
-  attribute_name.copy(attribute_name_c);
-  CustomData_add_layer_named(
-      custom_data, data_type, CD_DEFAULT, nullptr, domain_size, attribute_name_c);
+  add_named_custom_data_layer_from_attribute_init(
+      attribute_name, *custom_data, data_type, domain_size, initializer);
   return true;
 }
 
@@ -621,7 +697,8 @@ bool GeometryComponent::attribute_try_delete(const StringRef attribute_name)
 
 bool GeometryComponent::attribute_try_create(const StringRef attribute_name,
                                              const AttributeDomain domain,
-                                             const CustomDataType data_type)
+                                             const CustomDataType data_type,
+                                             const AttributeInit &initializer)
 {
   using namespace blender::bke;
   if (attribute_name.is_empty()) {
@@ -640,18 +717,19 @@ bool GeometryComponent::attribute_try_create(const StringRef attribute_name,
     if (builtin_provider->data_type() != data_type) {
       return false;
     }
-    return builtin_provider->try_create(*this);
+    return builtin_provider->try_create(*this, initializer);
   }
   for (const DynamicAttributesProvider *dynamic_provider :
        providers->dynamic_attribute_providers()) {
-    if (dynamic_provider->try_create(*this, attribute_name, domain, data_type)) {
+    if (dynamic_provider->try_create(*this, attribute_name, domain, data_type, initializer)) {
       return true;
     }
   }
   return false;
 }
 
-bool GeometryComponent::attribute_try_create_builtin(const blender::StringRef attribute_name)
+bool GeometryComponent::attribute_try_create_builtin(const blender::StringRef attribute_name,
+                                                     const AttributeInit &initializer)
 {
   using namespace blender::bke;
   if (attribute_name.is_empty()) {
@@ -666,7 +744,7 @@ bool GeometryComponent::attribute_try_create_builtin(const blender::StringRef at
   if (builtin_provider == nullptr) {
     return false;
   }
-  return builtin_provider->try_create(*this);
+  return builtin_provider->try_create(*this, initializer);
 }
 
 Set<std::string> GeometryComponent::attribute_names() const
@@ -874,7 +952,8 @@ static void save_output_attribute(blender::bke::OutputAttribute &output_attribut
   const CPPType &cpp_type = output_attribute.cpp_type();
 
   component.attribute_try_delete(name);
-  if (!component.attribute_try_create(varray.final_name, domain, data_type)) {
+  if (!component.attribute_try_create(
+          varray.final_name, domain, data_type, AttributeInitDefault())) {
     CLOG_WARN(&LOG,
               "Could not create the '%s' attribute with type '%s'.",
               name.c_str(),
@@ -912,7 +991,15 @@ static blender::bke::OutputAttribute create_output_attribute(
   if (component.attribute_is_builtin(attribute_name)) {
     WriteAttributeLookup attribute = component.attribute_try_get_for_write(attribute_name);
     if (!attribute) {
-      component.attribute_try_create_builtin(attribute_name);
+      if (default_value) {
+        const int64_t domain_size = component.attribute_domain_size(domain);
+        const GVArray_For_SingleValueRef default_varray{*cpp_type, domain_size, default_value};
+        component.attribute_try_create_builtin(attribute_name,
+                                               AttributeInitVArray(&default_varray));
+      }
+      else {
+        component.attribute_try_create_builtin(attribute_name, AttributeInitDefault());
+      }
       attribute = component.attribute_try_get_for_write(attribute_name);
       if (!attribute) {
         /* Builtin attribute does not exist and can't be created. */
@@ -933,9 +1020,19 @@ static blender::bke::OutputAttribute create_output_attribute(
     return OutputAttribute(std::move(varray), domain, {}, ignore_old_values);
   }
 
+  const int domain_size = component.attribute_domain_size(domain);
+
   WriteAttributeLookup attribute = component.attribute_try_get_for_write(attribute_name);
   if (!attribute) {
-    component.attribute_try_create(attribute_name, domain, data_type);
+    if (default_value) {
+      const GVArray_For_SingleValueRef default_varray{*cpp_type, domain_size, default_value};
+      component.attribute_try_create(
+          attribute_name, domain, data_type, AttributeInitVArray(&default_varray));
+    }
+    else {
+      component.attribute_try_create(attribute_name, domain, data_type, AttributeInitDefault());
+    }
+
     attribute = component.attribute_try_get_for_write(attribute_name);
     if (!attribute) {
       /* Can't create the attribute. */
@@ -947,7 +1044,6 @@ static blender::bke::OutputAttribute create_output_attribute(
     return OutputAttribute(std::move(attribute.varray), domain, {}, ignore_old_values);
   }
 
-  const int domain_size = component.attribute_domain_size(domain);
   /* Allocate a new array that lives next to the existing attribute. It will overwrite the existing
    * attribute after processing is done. */
   void *data = MEM_mallocN_aligned(
diff --git a/source/blender/blenkernel/intern/attribute_access_intern.hh b/source/blender/blenkernel/intern/attribute_access_intern.hh
index 7cf585dfbfc..b3a795faa30 100644
--- a/source/blender/blenkernel/intern/attribute_access_intern.hh
+++ b/source/blender/blenkernel/intern/attribute_access_intern.hh
@@ -89,7 +89,8 @@ class BuiltinAttributeProvider {
   virtual GVArrayPtr try_get_for_read(const GeometryComponent &component) const = 0;
   virtual GVMutableArrayPtr try_get_for_write(GeometryComponent &component) const = 0;
   virtual bool try_delete(GeometryComponent &component) const = 0;
-  virtual bool try_create(GeometryComponent &UNUSED(component)) const = 0;
+  virtual bool try_create(GeometryComponent &UNUSED(component),
+                          const AttributeInit &UNUSED(initializer)) const = 0;
   virtual bool exists(const GeometryComponent &component) const = 0;
 
   StringRefNull name() const
@@ -122,7 +123,8 @@ class DynamicAttributesProvider {
   virtual bool try_create(GeometryComponent &UNUSED(component),
                           const StringRef UNUSED(attribute_name),
                           const AttributeDomain UNUSED(domain),
-                          const CustomDataType UNUSED(data_type)) const
+                          const CustomDataType UNUSED(data_type),
+                          const AttributeInit &UNUSED(initializer)) const
   {
     /* Some providers should not create new attributes. */
     return false;
@@ -162,7 +164,8 @@ class CustomDataAttributeProvider final : public DynamicAttributesProvider {
   bool try_create(GeometryComponent &component,
                   const StringRef attribute_name,
                   const AttributeDomain domain,
-                  const CustomDataType data_type) const final;
+                  const CustomDataType data_type,
+                  const AttributeInit &initializer) const final;
 
   bool foreach_attribute(const GeometryComponent &component,
                          const AttributeForeachCallback callback) const final;
@@ -278,7 +281,7 @@ class BuiltinCustomDataLayerProvider final : public BuiltinAttributeProvider {
   GVArrayPtr try_get_for_read(const GeometryComponent &component) const final;
   GVMutableArrayPtr try_get_for_write(GeometryComponent &component) const final;
   bool try_delete(GeometryComponent &component) const final;
-  bool try_create(GeometryComponent &component) const final;
+  bool try_create(GeometryComponent &component, const AttributeInit &initializer) const final;
   bool exists(const GeometryComponent &component) const final;
 };
 
diff --git a/source/blender/blenkernel/intern/geometry_component_mesh.cc b/source/blender/blenkernel/intern/geometry_component_mesh.cc
index e5accd98952..e54c3716660 100644
--- a/source/blender/blenkernel/intern/geometry_component_mesh.cc
+++ b/source/blender/blenkernel/intern/geometry_component_mesh.cc
@@ -1014,7 +1014,8 @@ class NormalAttributeProvider final : public BuiltinAttributeProvider {
     return false;
   }
 
-  bool try_create(GeometryComponent &UNUSED(component)) const final
+  bool try_create(GeometryComponent &UNUSED(component),
+                  const AttributeInit &UNUSED(initializer)) const final
   {
     return false;
   }
diff --git a/source/blender/blenkernel/intern/geometry_set_instances.cc b/source/blender/blenkernel/intern/geometry_set_instances.cc
index 2e9f6daabad..47db5d1f901 100644
--- a/source/blender/blenkernel/intern/geometry_set_instances.cc
+++ b/source/blender/blenkernel/intern/geometry_set_instances.cc
@@ -449,7 +449,8 @@ static void join_attributes(Span<GeometryInstanceGroup> set_groups,
     const CPPType *cpp_type = bke::custom_data_type_to_cpp_type(data_type_output);
     BLI_assert(cpp_type != nullptr);
 
-    result.attribute_try_create(entry.key, domain_output, data_type_output);
+    result.attribute_try_create(
+        entry.key, domain_output, data_type_output, AttributeInitDefault());
     WriteAttributeLookup write_attribute = result.attribute_try_get_for_write(name);
     if (!write_attribute || &write_attribute.varray->type() != cpp_type ||
         write_attribute.domain != domain_output) {