Geometry Nodes: use virtual arrays in internal attribute api

A virtual array is a data structure that is similar to a normal array
in that its elements can be accessed by an index. However, a virtual
array does not have to be a contiguous array internally. Instead, its
elements can be layed out arbitrarily while element access happens
through a virtual function call. However, the virtual array data
structures are designed so that the virtual function call can be avoided
in cases where it could become a bottleneck.

Most commonly, a virtual array is backed by an actual array/span or
is a single value internally, that is the same for every index.
Besides those, there are many more specialized virtual arrays like the
ones that provides vertex positions based on the `MVert` struct or
vertex group weights.

Not all attributes used by geometry nodes are stored in simple contiguous
arrays. To provide uniform access to all kinds of attributes, the attribute
API has to provide virtual array functionality that hides the implementation
details of attributes.

Before this refactor, the attribute API provided its own virtual array
implementation as part of the `ReadAttribute` and `WriteAttribute` types.
That resulted in unnecessary code duplication with the virtual array system.
Even worse, it bound many algorithms used by geometry nodes to the specifics
of the attribute API, even though they could also use different data sources
(such as data from sockets, default values, later results of expressions, ...).

This refactor removes the `ReadAttribute` and `WriteAttribute` types and
replaces them with `GVArray` and `GVMutableArray` respectively. The `GV`
stands for "generic virtual". The "generic" means that the data type contained
in those virtual arrays is only known at run-time. There are the corresponding
statically typed types `VArray<T>` and `VMutableArray<T>` as well.

No regressions are expected from this refactor. It does come with one
improvement for users. The attribute API can convert the data type
on write now. This is especially useful when writing to builtin attributes
like `material_index` with e.g. the Attribute Math node (which usually
just writes to float attributes, while `material_index` is an integer attribute).

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

2 files changed, 143 insertions, 19 deletions

diff --git a/source/blender/nodes/intern/node_geometry_exec.cc b/source/blender/nodes/intern/node_geometry_exec.cc
index a4fb99a988e..2648336f0c0 100644
--- a/source/blender/nodes/intern/node_geometry_exec.cc
+++ b/source/blender/nodes/intern/node_geometry_exec.cc
@@ -22,6 +22,7 @@
 
 #include "NOD_geometry_exec.hh"
 #include "NOD_type_callbacks.hh"
+#include "NOD_type_conversions.hh"
 
 #include "node_geometry_util.hh"
 
@@ -53,22 +54,29 @@ const bNodeSocket *GeoNodeExecParams::find_available_socket(const StringRef name
   return nullptr;
 }
 
-ReadAttributePtr GeoNodeExecParams::get_input_attribute(const StringRef name,
-                                                        const GeometryComponent &component,
-                                                        const AttributeDomain domain,
-                                                        const CustomDataType type,
-                                                        const void *default_value) const
+GVArrayPtr GeoNodeExecParams::get_input_attribute(const StringRef name,
+                                                  const GeometryComponent &component,
+                                                  const AttributeDomain domain,
+                                                  const CustomDataType type,
+                                                  const void *default_value) const
 {
   const bNodeSocket *found_socket = this->find_available_socket(name);
   BLI_assert(found_socket != nullptr); /* There should always be available socket for the name. */
+  const CPPType *cpp_type = bke::custom_data_type_to_cpp_type(type);
+  const int64_t domain_size = component.attribute_domain_size(domain);
+
+  if (default_value == nullptr) {
+    default_value = cpp_type->default_value();
+  }
+
   if (found_socket == nullptr) {
-    return component.attribute_get_constant_for_read(domain, type, default_value);
+    return std::make_unique<fn::GVArray_For_SingleValue>(*cpp_type, domain_size, default_value);
   }
 
   if (found_socket->type == SOCK_STRING) {
     const std::string name = this->get_input<std::string>(found_socket->identifier);
     /* Try getting the attribute without the default value. */
-    ReadAttributePtr attribute = component.attribute_try_get_for_read(name, domain, type);
+    GVArrayPtr attribute = component.attribute_try_get_for_read(name, domain, type);
     if (attribute) {
       return attribute;
     }
@@ -80,25 +88,29 @@ ReadAttributePtr GeoNodeExecParams::get_input_attribute(const StringRef name,
       this->error_message_add(NodeWarningType::Error,
                               TIP_("No attribute with name \"") + name + "\"");
     }
-    return component.attribute_get_constant_for_read(domain, type, default_value);
+    return std::make_unique<fn::GVArray_For_SingleValue>(*cpp_type, domain_size, default_value);
   }
+  const DataTypeConversions &conversions = get_implicit_type_conversions();
   if (found_socket->type == SOCK_FLOAT) {
     const float value = this->get_input<float>(found_socket->identifier);
-    return component.attribute_get_constant_for_read_converted(
-        domain, CD_PROP_FLOAT, type, &value);
+    BUFFER_FOR_CPP_TYPE_VALUE(*cpp_type, buffer);
+    conversions.convert_to_uninitialized(CPPType::get<float>(), *cpp_type, &value, buffer);
+    return std::make_unique<fn::GVArray_For_SingleValue>(*cpp_type, domain_size, buffer);
   }
   if (found_socket->type == SOCK_VECTOR) {
     const float3 value = this->get_input<float3>(found_socket->identifier);
-    return component.attribute_get_constant_for_read_converted(
-        domain, CD_PROP_FLOAT3, type, &value);
+    BUFFER_FOR_CPP_TYPE_VALUE(*cpp_type, buffer);
+    conversions.convert_to_uninitialized(CPPType::get<float3>(), *cpp_type, &value, buffer);
+    return std::make_unique<fn::GVArray_For_SingleValue>(*cpp_type, domain_size, buffer);
   }
   if (found_socket->type == SOCK_RGBA) {
     const Color4f value = this->get_input<Color4f>(found_socket->identifier);
-    return component.attribute_get_constant_for_read_converted(
-        domain, CD_PROP_COLOR, type, &value);
+    BUFFER_FOR_CPP_TYPE_VALUE(*cpp_type, buffer);
+    conversions.convert_to_uninitialized(CPPType::get<Color4f>(), *cpp_type, &value, buffer);
+    return std::make_unique<fn::GVArray_For_SingleValue>(*cpp_type, domain_size, buffer);
   }
   BLI_assert(false);
-  return component.attribute_get_constant_for_read(domain, type, default_value);
+  return std::make_unique<fn::GVArray_For_SingleValue>(*cpp_type, domain_size, default_value);
 }
 
 CustomDataType GeoNodeExecParams::get_input_attribute_data_type(
@@ -114,11 +126,11 @@ CustomDataType GeoNodeExecParams::get_input_attribute_data_type(
 
   if (found_socket->type == SOCK_STRING) {
     const std::string name = this->get_input<std::string>(found_socket->identifier);
-    ReadAttributePtr attribute = component.attribute_try_get_for_read(name);
+    ReadAttributeLookup attribute = component.attribute_try_get_for_read(name);
     if (!attribute) {
       return default_type;
     }
-    return attribute->custom_data_type();
+    return bke::cpp_type_to_custom_data_type(attribute.varray->type());
   }
   if (found_socket->type == SOCK_FLOAT) {
     return CD_PROP_FLOAT;
@@ -157,9 +169,9 @@ AttributeDomain GeoNodeExecParams::get_highest_priority_input_domain(
 
     if (found_socket->type == SOCK_STRING) {
       const std::string name = this->get_input<std::string>(found_socket->identifier);
-      ReadAttributePtr attribute = component.attribute_try_get_for_read(name);
+      ReadAttributeLookup attribute = component.attribute_try_get_for_read(name);
       if (attribute) {
-        input_domains.append(attribute->domain());
+        input_domains.append(attribute.domain);
       }
     }
   }
diff --git a/source/blender/nodes/intern/type_conversions.cc b/source/blender/nodes/intern/type_conversions.cc
index 1c1b7c7feb5..63f7b8a9ee8 100644
--- a/source/blender/nodes/intern/type_conversions.cc
+++ b/source/blender/nodes/intern/type_conversions.cc
@@ -24,6 +24,9 @@
 
 namespace blender::nodes {
 
+using fn::GVArrayPtr;
+using fn::GVMutableArray;
+using fn::GVMutableArrayPtr;
 using fn::MFDataType;
 
 template<typename From, typename To, To (*ConversionF)(const From &)>
@@ -227,6 +230,11 @@ void DataTypeConversions::convert_to_uninitialized(const CPPType &from_type,
                                                    const void *from_value,
                                                    void *to_value) const
 {
+  if (from_type == to_type) {
+    from_type.copy_to_uninitialized(from_value, to_value);
+    return;
+  }
+
   const ConversionFunctions *functions = this->get_conversion_functions(
       MFDataType::ForSingle(from_type), MFDataType::ForSingle(to_type));
   BLI_assert(functions != nullptr);
@@ -234,4 +242,108 @@ void DataTypeConversions::convert_to_uninitialized(const CPPType &from_type,
   functions->convert_single_to_uninitialized(from_value, to_value);
 }
 
+class GVArray_For_ConvertedGVArray : public GVArray {
+ private:
+  GVArrayPtr varray_;
+  const CPPType &from_type_;
+  ConversionFunctions old_to_new_conversions_;
+
+ public:
+  GVArray_For_ConvertedGVArray(GVArrayPtr varray,
+                               const CPPType &to_type,
+                               const DataTypeConversions &conversions)
+      : GVArray(to_type, varray->size()), varray_(std::move(varray)), from_type_(varray_->type())
+  {
+    old_to_new_conversions_ = *conversions.get_conversion_functions(from_type_, to_type);
+  }
+
+ private:
+  void get_impl(const int64_t index, void *r_value) const override
+  {
+    BUFFER_FOR_CPP_TYPE_VALUE(from_type_, buffer);
+    varray_->get(index, buffer);
+    old_to_new_conversions_.convert_single_to_initialized(buffer, r_value);
+    from_type_.destruct(buffer);
+  }
+
+  void get_to_uninitialized_impl(const int64_t index, void *r_value) const override
+  {
+    BUFFER_FOR_CPP_TYPE_VALUE(from_type_, buffer);
+    varray_->get(index, buffer);
+    old_to_new_conversions_.convert_single_to_uninitialized(buffer, r_value);
+    from_type_.destruct(buffer);
+  }
+};
+
+class GVMutableArray_For_ConvertedGVMutableArray : public GVMutableArray {
+ private:
+  GVMutableArrayPtr varray_;
+  const CPPType &from_type_;
+  ConversionFunctions old_to_new_conversions_;
+  ConversionFunctions new_to_old_conversions_;
+
+ public:
+  GVMutableArray_For_ConvertedGVMutableArray(GVMutableArrayPtr varray,
+                                             const CPPType &to_type,
+                                             const DataTypeConversions &conversions)
+      : GVMutableArray(to_type, varray->size()),
+        varray_(std::move(varray)),
+        from_type_(varray_->type())
+  {
+    old_to_new_conversions_ = *conversions.get_conversion_functions(from_type_, to_type);
+    new_to_old_conversions_ = *conversions.get_conversion_functions(to_type, from_type_);
+  }
+
+ private:
+  void get_impl(const int64_t index, void *r_value) const override
+  {
+    BUFFER_FOR_CPP_TYPE_VALUE(from_type_, buffer);
+    varray_->get(index, buffer);
+    old_to_new_conversions_.convert_single_to_initialized(buffer, r_value);
+    from_type_.destruct(buffer);
+  }
+
+  void get_to_uninitialized_impl(const int64_t index, void *r_value) const override
+  {
+    BUFFER_FOR_CPP_TYPE_VALUE(from_type_, buffer);
+    varray_->get(index, buffer);
+    old_to_new_conversions_.convert_single_to_uninitialized(buffer, r_value);
+    from_type_.destruct(buffer);
+  }
+
+  void set_by_move_impl(const int64_t index, void *value) override
+  {
+    BUFFER_FOR_CPP_TYPE_VALUE(from_type_, buffer);
+    new_to_old_conversions_.convert_single_to_uninitialized(value, buffer);
+    varray_->set_by_relocate(index, buffer);
+  }
+};
+
+fn::GVArrayPtr DataTypeConversions::try_convert(fn::GVArrayPtr varray,
+                                                const CPPType &to_type) const
+{
+  const CPPType &from_type = varray->type();
+  if (from_type == to_type) {
+    return varray;
+  }
+  if (!this->is_convertible(from_type, to_type)) {
+    return {};
+  }
+  return std::make_unique<GVArray_For_ConvertedGVArray>(std::move(varray), to_type, *this);
+}
+
+fn::GVMutableArrayPtr DataTypeConversions::try_convert(fn::GVMutableArrayPtr varray,
+                                                       const CPPType &to_type) const
+{
+  const CPPType &from_type = varray->type();
+  if (from_type == to_type) {
+    return varray;
+  }
+  if (!this->is_convertible(from_type, to_type)) {
+    return {};
+  }
+  return std::make_unique<GVMutableArray_For_ConvertedGVMutableArray>(
+      std::move(varray), to_type, *this);
+}
+
 }  // namespace blender::nodes