initial WriteAttribute

4 files changed, 314 insertions, 62 deletions

diff --git a/source/blender/blenkernel/BKE_attribute_accessor.hh b/source/blender/blenkernel/BKE_attribute_accessor.hh
index 3cf2fd4f865..0905f81d472 100644
--- a/source/blender/blenkernel/BKE_attribute_accessor.hh
+++ b/source/blender/blenkernel/BKE_attribute_accessor.hh
@@ -39,7 +39,7 @@ class ReadAttribute {
   {
   }
 
-  virtual ~ReadAttribute() = default;
+  virtual ~ReadAttribute();
 
   AttributeDomain domain() const
   {
@@ -67,7 +67,54 @@ class ReadAttribute {
   virtual void get_internal(const int64_t index, void *r_value) const = 0;
 };
 
+class WriteAttribute {
+ protected:
+  const AttributeDomain domain_;
+  const CPPType &cpp_type_;
+  const int64_t size_;
+
+ public:
+  WriteAttribute(AttributeDomain domain, const CPPType &cpp_type, const int64_t size)
+      : domain_(domain), cpp_type_(cpp_type), size_(size)
+  {
+  }
+
+  virtual ~WriteAttribute();
+
+  AttributeDomain domain() const
+  {
+    return domain_;
+  }
+
+  const CPPType &cpp_type() const
+  {
+    return cpp_type_;
+  }
+
+  int64_t size() const
+  {
+    return size_;
+  }
+
+  void get(const int64_t index, void *r_value) const
+  {
+    BLI_assert(index < size_);
+    this->get_internal(index, r_value);
+  }
+
+  void set(const int64_t index, const void *value)
+  {
+    BLI_assert(index < size_);
+    this->set_internal(index, value);
+  }
+
+ protected:
+  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;
+};
+
 using ReadAttributePtr = std::unique_ptr<ReadAttribute>;
+using WriteAttributePtr = std::unique_ptr<WriteAttribute>;
 
 template<typename T> class TypedReadAttribute {
  private:
@@ -95,11 +142,46 @@ template<typename T> class TypedReadAttribute {
   }
 };
 
+template<typename T> class TypedWriteAttribute {
+ private:
+  WriteAttributePtr attribute_;
+
+ public:
+  TypedWriteAttribute(WriteAttributePtr attribute) : attribute_(std::move(attribute))
+  {
+    BLI_assert(attribute_);
+    BLI_assert(attribute_->cpp_type().is<T>());
+  }
+
+  int64_t size() const
+  {
+    return attribute_->size();
+  }
+
+  T operator[](const int64_t index) const
+  {
+    BLI_assert(index < attribute_->size());
+    T value;
+    value.~T();
+    attribute_->get(index, &value);
+    return value;
+  }
+
+  void set(const int64_t index, const T &value)
+  {
+    attribute_->set(index, &value);
+  }
+};
+
 using FloatReadAttribute = TypedReadAttribute<float>;
 using Float3ReadAttribute = TypedReadAttribute<float3>;
+using FloatWriteAttribute = TypedWriteAttribute<float>;
+using Float3WriteAttribute = TypedWriteAttribute<float3>;
 
 ReadAttributePtr mesh_attribute_get_for_read(const MeshComponent &mesh_component,
                                              const StringRef attribute_name);
+std::optional<WriteAttributePtr> mesh_attribute_get_for_write(MeshComponent &mesh_component,
+                                                              const StringRef attribute_name);
 
 ReadAttributePtr mesh_attribute_adapt_domain(const MeshComponent &mesh_component,
                                              ReadAttributePtr attribute,
diff --git a/source/blender/blenkernel/intern/attribute_accessor.cc b/source/blender/blenkernel/intern/attribute_accessor.cc
index ee1cf2acf36..46618e49826 100644
--- a/source/blender/blenkernel/intern/attribute_accessor.cc
+++ b/source/blender/blenkernel/intern/attribute_accessor.cc
@@ -28,14 +28,17 @@
 
 namespace blender::bke {
 
-class VertexWeightReadAttribute final : public ReadAttribute {
+ReadAttribute::~ReadAttribute() = default;
+WriteAttribute::~WriteAttribute() = default;
+
+class VertexWeightWriteAttribute final : public WriteAttribute {
  private:
-  Span<MDeformVert> dverts_;
-  int dvert_index_;
+  MutableSpan<MDeformVert> dverts_;
+  const int dvert_index_;
 
  public:
-  VertexWeightReadAttribute(const MDeformVert *dverts, const int totvert, const int dvert_index)
-      : ReadAttribute(ATTR_DOMAIN_VERTEX, CPPType::get<float>(), totvert),
+  VertexWeightWriteAttribute(MDeformVert *dverts, const int totvert, const int dvert_index)
+      : WriteAttribute(ATTR_DOMAIN_VERTEX, CPPType::get<float>(), totvert),
         dverts_(dverts, totvert),
         dvert_index_(dvert_index)
   {
@@ -43,9 +46,23 @@ class VertexWeightReadAttribute final : public ReadAttribute {
 
   void get_internal(const int64_t index, void *r_value) const override
   {
-    const MDeformVert &dvert = dverts_[index];
+    this->get_internal(dverts_, dvert_index_, index, r_value);
+  }
+
+  void set_internal(const int64_t index, const void *value) override
+  {
+    MDeformWeight *weight = BKE_defvert_ensure_index(&dverts_[index], dvert_index_);
+    weight->weight = *reinterpret_cast<const float *>(value);
+  }
+
+  static void get_internal(const Span<MDeformVert> dverts,
+                           const int dvert_index,
+                           const int64_t index,
+                           void *r_value)
+  {
+    const MDeformVert &dvert = dverts[index];
     for (const MDeformWeight &weight : Span(dvert.dw, dvert.totweight)) {
-      if (weight.def_nr == dvert_index_) {
+      if (weight.def_nr == dvert_index) {
         *(float *)r_value = weight.weight;
         return;
       }
@@ -54,6 +71,46 @@ class VertexWeightReadAttribute final : public ReadAttribute {
   }
 };
 
+class VertexWeightReadAttribute final : public ReadAttribute {
+ private:
+  const Span<MDeformVert> dverts_;
+  const int dvert_index_;
+
+ public:
+  VertexWeightReadAttribute(const MDeformVert *dverts, const int totvert, const int dvert_index)
+      : ReadAttribute(ATTR_DOMAIN_VERTEX, CPPType::get<float>(), totvert),
+        dverts_(dverts, totvert),
+        dvert_index_(dvert_index)
+  {
+  }
+
+  void get_internal(const int64_t index, void *r_value) const override
+  {
+    VertexWeightWriteAttribute::get_internal(dverts_, dvert_index_, index, r_value);
+  }
+};
+
+template<typename T> class ArrayWriteAttribute final : public WriteAttribute {
+ private:
+  MutableSpan<T> data_;
+
+ public:
+  ArrayWriteAttribute(AttributeDomain domain, MutableSpan<T> data)
+      : WriteAttribute(domain, CPPType::get<T>(), data.size()), data_(data)
+  {
+  }
+
+  void get_internal(const int64_t index, void *r_value) const override
+  {
+    new (r_value) T(data_[index]);
+  }
+
+  void set_internal(const int64_t index, const void *value) override
+  {
+    data_[index] = *reinterpret_cast<const T *>(value);
+  }
+};
+
 template<typename T> class ArrayReadAttribute final : public ReadAttribute {
  private:
   Span<T> data_;
@@ -70,26 +127,58 @@ template<typename T> class ArrayReadAttribute final : public ReadAttribute {
   }
 };
 
-template<typename StructT, typename FuncT>
-class DerivedArrayReadAttribute final : public ReadAttribute {
+template<typename StructT, typename ElemT, typename GetFuncT, typename SetFuncT>
+class DerivedArrayWriteAttribute final : public WriteAttribute {
  private:
-  using ElemT = decltype(std::declval<FuncT>()(std::declval<StructT>()));
+  MutableSpan<StructT> data_;
+  GetFuncT get_function_;
+  SetFuncT set_function_;
+
+ public:
+  DerivedArrayWriteAttribute(AttributeDomain domain,
+                             MutableSpan<StructT> data,
+                             GetFuncT get_function,
+                             SetFuncT set_function)
+      : WriteAttribute(domain, CPPType::get<ElemT>(), data.size()),
+        data_(data),
+        get_function_(std::move(get_function)),
+        set_function_(std::move(set_function))
+  {
+  }
+
+  void get_internal(const int64_t index, void *r_value) const override
+  {
+    const StructT &struct_value = data_[index];
+    const ElemT value = get_function_(struct_value);
+    new (r_value) ElemT(value);
+  }
+
+  void set_internal(const int64_t index, const void *value) override
+  {
+    StructT &struct_value = data_[index];
+    const ElemT &typed_value = *reinterpret_cast<const ElemT *>(value);
+    set_function_(struct_value, typed_value);
+  }
+};
 
+template<typename StructT, typename ElemT, typename GetFuncT>
+class DerivedArrayReadAttribute final : public ReadAttribute {
+ private:
   Span<StructT> data_;
-  FuncT function_;
+  GetFuncT get_function_;
 
  public:
-  DerivedArrayReadAttribute(AttributeDomain domain, Span<StructT> data, FuncT function)
+  DerivedArrayReadAttribute(AttributeDomain domain, Span<StructT> data, GetFuncT get_function)
       : ReadAttribute(domain, CPPType::get<ElemT>(), data.size()),
         data_(data),
-        function_(std::move(function))
+        get_function_(std::move(get_function))
   {
   }
 
   void get_internal(const int64_t index, void *r_value) const override
   {
     const StructT &struct_value = data_[index];
-    const ElemT value = function_(struct_value);
+    const ElemT value = get_function_(struct_value);
     new (r_value) ElemT(value);
   }
 };
@@ -115,10 +204,10 @@ class ConstantReadAttribute final : public ReadAttribute {
   }
 };
 
-static ReadAttributePtr get_custom_data_read_attribute(const CustomData &custom_data,
-                                                       const int size,
-                                                       const StringRef attribute_name,
-                                                       const AttributeDomain domain)
+static ReadAttributePtr mesh_attribute_custom_data_read(const CustomData &custom_data,
+                                                        const int size,
+                                                        const StringRef attribute_name,
+                                                        const AttributeDomain domain)
 {
   for (const CustomDataLayer &layer : Span(custom_data.layers, custom_data.totlayer)) {
     if (layer.name != nullptr && layer.name == attribute_name) {
@@ -144,29 +233,53 @@ static ReadAttributePtr get_custom_data_read_attribute(const CustomData &custom_
   return {};
 }
 
-static ReadAttributePtr get_mesh_read_attribute__corner(const MeshComponent &mesh_component,
-                                                        const StringRef attribute_name)
+static WriteAttributePtr mesh_attribute_custom_data_write(CustomData custom_data,
+                                                          const int size,
+                                                          const StringRef attribute_name,
+                                                          const AttributeDomain domain)
 {
-  const Mesh *mesh = mesh_component.get_for_read();
-  if (mesh == nullptr) {
-    return {};
+  for (const CustomDataLayer &layer : Span(custom_data.layers, custom_data.totlayer)) {
+    if (layer.name != nullptr && layer.name == attribute_name) {
+      switch (layer.type) {
+        case CD_PROP_FLOAT:
+          return std::make_unique<ArrayWriteAttribute<float>>(
+              domain, MutableSpan(static_cast<float *>(layer.data), size));
+        case CD_PROP_FLOAT2:
+          return std::make_unique<ArrayWriteAttribute<float2>>(
+              domain, MutableSpan(static_cast<float2 *>(layer.data), size));
+        case CD_PROP_FLOAT3:
+          return std::make_unique<ArrayWriteAttribute<float3>>(
+              domain, MutableSpan(static_cast<float3 *>(layer.data), size));
+        case CD_PROP_INT32:
+          return std::make_unique<ArrayWriteAttribute<int>>(
+              domain, MutableSpan(static_cast<int *>(layer.data), size));
+        case CD_PROP_COLOR:
+          return std::make_unique<ArrayWriteAttribute<Color4f>>(
+              domain, MutableSpan(static_cast<Color4f *>(layer.data), size));
+      }
+    }
   }
-
-  return get_custom_data_read_attribute(
-      mesh->ldata, mesh->totloop, attribute_name, ATTR_DOMAIN_CORNER);
+  return {};
 }
 
-static ReadAttributePtr get_mesh_read_attribute__vertex(const MeshComponent &mesh_component,
-                                                        const StringRef attribute_name)
+ReadAttributePtr mesh_attribute_get_for_read(const MeshComponent &mesh_component,
+                                             const StringRef attribute_name)
 {
   const Mesh *mesh = mesh_component.get_for_read();
   if (mesh == nullptr) {
     return {};
   }
 
+  ReadAttributePtr corner_attribute = mesh_attribute_custom_data_read(
+      mesh->ldata, mesh->totloop, attribute_name, ATTR_DOMAIN_CORNER);
+  if (corner_attribute) {
+    return corner_attribute;
+  }
+
   if (attribute_name == "Position") {
     auto get_vertex_position = [](const MVert &vert) { return float3(vert.co); };
-    return std::make_unique<DerivedArrayReadAttribute<MVert, decltype(get_vertex_position)>>(
+    return std::make_unique<
+        DerivedArrayReadAttribute<MVert, float3, decltype(get_vertex_position)>>(
         ATTR_DOMAIN_VERTEX, Span(mesh->mvert, mesh->totvert), get_vertex_position);
   }
 
@@ -176,56 +289,76 @@ static ReadAttributePtr get_mesh_read_attribute__vertex(const MeshComponent &mes
         mesh->dvert, mesh->totvert, vertex_group_index);
   }
 
-  return get_custom_data_read_attribute(
+  ReadAttributePtr vertex_attribute = mesh_attribute_custom_data_read(
       mesh->vdata, mesh->totvert, attribute_name, ATTR_DOMAIN_VERTEX);
-}
-
-static ReadAttributePtr get_mesh_read_attribute__edge(const MeshComponent &mesh_component,
-                                                      const StringRef attribute_name)
-{
-  const Mesh *mesh = mesh_component.get_for_read();
-  if (mesh == nullptr) {
-    return {};
+  if (vertex_attribute) {
+    return vertex_attribute;
   }
 
-  return get_custom_data_read_attribute(
+  ReadAttributePtr edge_attribute = mesh_attribute_custom_data_read(
       mesh->edata, mesh->totedge, attribute_name, ATTR_DOMAIN_EDGE);
+  if (edge_attribute) {
+    return edge_attribute;
+  }
+
+  ReadAttributePtr polygon_attribute = mesh_attribute_custom_data_read(
+      mesh->pdata, mesh->totpoly, attribute_name, ATTR_DOMAIN_POLYGON);
+  if (polygon_attribute) {
+    return polygon_attribute;
+  }
+
+  return {};
 }
 
-static ReadAttributePtr get_mesh_read_attribute__polygon(const MeshComponent &mesh_component,
-                                                         const StringRef attribute_name)
+std::optional<WriteAttributePtr> mesh_attribute_get_for_write(MeshComponent &mesh_component,
+                                                              const StringRef attribute_name)
 {
-  const Mesh *mesh = mesh_component.get_for_read();
+  Mesh *mesh = mesh_component.get_for_write();
   if (mesh == nullptr) {
     return {};
   }
 
-  return get_custom_data_read_attribute(
-      mesh->pdata, mesh->totpoly, attribute_name, ATTR_DOMAIN_POLYGON);
-}
+  WriteAttributePtr corner_attribute = mesh_attribute_custom_data_write(
+      mesh->ldata, mesh->totloop, attribute_name, ATTR_DOMAIN_CORNER);
+  if (corner_attribute) {
+    return corner_attribute;
+  }
 
-ReadAttributePtr mesh_attribute_get_for_read(const MeshComponent &mesh_component,
-                                             const StringRef attribute_name)
-{
-  ReadAttributePtr corner_level = get_mesh_read_attribute__corner(mesh_component, attribute_name);
-  if (corner_level) {
-    return corner_level;
+  if (attribute_name == "Position") {
+    auto get_vertex_position = [](const MVert &vert) { return float3(vert.co); };
+    auto set_vertex_position = [](MVert &vert, const float3 &co) { copy_v3_v3(vert.co, co); };
+    return std::make_unique<DerivedArrayWriteAttribute<MVert,
+                                                       float3,
+                                                       decltype(get_vertex_position),
+                                                       decltype(set_vertex_position)>>(
+        ATTR_DOMAIN_VERTEX,
+        MutableSpan(mesh->mvert, mesh->totvert),
+        get_vertex_position,
+        set_vertex_position);
   }
 
-  ReadAttributePtr vertex_level = get_mesh_read_attribute__vertex(mesh_component, attribute_name);
-  if (vertex_level) {
-    return vertex_level;
+  const int vertex_group_index = mesh_component.vertex_group_index(attribute_name);
+  if (vertex_group_index >= 0) {
+    return std::make_unique<VertexWeightWriteAttribute>(
+        mesh->dvert, mesh->totvert, vertex_group_index);
+  }
+
+  WriteAttributePtr vertex_attribute = mesh_attribute_custom_data_write(
+      mesh->vdata, mesh->totvert, attribute_name, ATTR_DOMAIN_VERTEX);
+  if (vertex_attribute) {
+    return vertex_attribute;
   }
 
-  ReadAttributePtr edge_level = get_mesh_read_attribute__edge(mesh_component, attribute_name);
-  if (edge_level) {
-    return edge_level;
+  WriteAttributePtr edge_attribute = mesh_attribute_custom_data_write(
+      mesh->edata, mesh->totedge, attribute_name, ATTR_DOMAIN_EDGE);
+  if (edge_attribute) {
+    return edge_attribute;
   }
 
-  ReadAttributePtr polygon_level = get_mesh_read_attribute__polygon(mesh_component,
-                                                                    attribute_name);
-  if (polygon_level) {
-    return polygon_level;
+  WriteAttributePtr polygon_attribute = mesh_attribute_custom_data_write(
+      mesh->pdata, mesh->totpoly, attribute_name, ATTR_DOMAIN_POLYGON);
+  if (polygon_attribute) {
+    return polygon_attribute;
   }
 
   return {};
diff --git a/source/blender/nodes/NOD_geometry_exec.hh b/source/blender/nodes/NOD_geometry_exec.hh
index 46548bfdc77..e98b7d926e9 100644
--- a/source/blender/nodes/NOD_geometry_exec.hh
+++ b/source/blender/nodes/NOD_geometry_exec.hh
@@ -27,11 +27,15 @@
 namespace blender::nodes {
 
 using bke::Float3ReadAttribute;
+using bke::Float3WriteAttribute;
 using bke::FloatReadAttribute;
+using bke::FloatWriteAttribute;
 using bke::PersistentDataHandleMap;
 using bke::PersistentObjectHandle;
 using bke::ReadAttribute;
 using bke::ReadAttributePtr;
+using bke::WriteAttribute;
+using bke::WriteAttributePtr;
 using fn::CPPType;
 using fn::GMutablePointer;
 using fn::GValueMap;
diff --git a/source/blender/nodes/geometry/nodes/node_geo_random_attribute.cc b/source/blender/nodes/geometry/nodes/node_geo_random_attribute.cc
index 12eda279c5f..ca40ac0b51a 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_random_attribute.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_random_attribute.cc
@@ -16,6 +16,8 @@
 
 #include "node_geometry_util.hh"
 
+#include "BLI_rand.hh"
+
 static bNodeSocketTemplate geo_node_random_attribute_in[] = {
     {SOCK_GEOMETRY, N_("Geometry")},
     {SOCK_STRING, N_("Attribute")},
@@ -39,6 +41,37 @@ static void geo_random_attribute_exec(GeoNodeExecParams params)
   const float3 min_value = params.extract_input<float3>("Min");
   const float3 max_value = params.extract_input<float3>("Max");
   const int seed = params.extract_input<int>("Seed");
+
+  MeshComponent &mesh_component = geometry_set.get_component_for_write<MeshComponent>();
+  std::optional<WriteAttributePtr> attribute_opt = bke::mesh_attribute_get_for_write(
+      mesh_component, attribute_name);
+
+  RandomNumberGenerator rng;
+  rng.seed_random(seed);
+
+  if (attribute_opt.has_value()) {
+    WriteAttributePtr attribute = std::move(*attribute_opt);
+    const int size = attribute->size();
+    if (attribute->cpp_type().is<float>()) {
+      FloatWriteAttribute float_attribute = std::move(attribute);
+      for (const int i : IndexRange(size)) {
+        const float value = rng.get_float() * (max_value.x - min_value.x) + min_value.x;
+        float_attribute.set(i, value);
+      }
+    }
+    else if (attribute->cpp_type().is<float3>()) {
+      Float3WriteAttribute float3_attribute = std::move(attribute);
+      for (const int i : IndexRange(size)) {
+        const float x = rng.get_float();
+        const float y = rng.get_float();
+        const float z = rng.get_float();
+        const float3 value = float3(x, y, z) * (max_value - min_value) + min_value;
+        float3_attribute.set(i, value);
+      }
+    }
+  }
+
+  params.set_output("Geometry", geometry_set);
 }
 
 }  // namespace blender::nodes