Functions: Various improvements to the spans and generic data structures

Most of this code is covered by unit tests.

8 files changed, 253 insertions, 202 deletions

diff --git a/source/blender/functions/FN_array_spans.hh b/source/blender/functions/FN_array_spans.hh
index ca818ad6b11..fac2ef42c9d 100644
--- a/source/blender/functions/FN_array_spans.hh
+++ b/source/blender/functions/FN_array_spans.hh
@@ -30,22 +30,18 @@ namespace blender {
 namespace fn {
 
 /**
- * A virtual array span. Every element of this span contains a virtual span. So it behaves like a
- * blender::Span, but might not be backed up by an actual array.
+ * Depending on the use case, the referenced data might have a different structure. More
+ * categories can be added when necessary.
  */
-template<typename T> class VArraySpan {
- private:
-  /**
-   * Depending on the use case, the referenced data might have a different structure. More
-   * categories can be added when necessary.
-   */
-  enum Category {
-    SingleArray,
-    StartsAndSizes,
-  };
+enum class VArraySpanCategory {
+  SingleArray,
+  StartsAndSizes,
+};
 
+template<typename T> class VArraySpanBase {
+ protected:
   uint m_virtual_size;
-  Category m_category;
+  VArraySpanCategory m_category;
 
   union {
     struct {
@@ -59,50 +55,77 @@ template<typename T> class VArraySpan {
   } m_data;
 
  public:
-  VArraySpan()
+  bool is_single_array() const
   {
-    m_virtual_size = 0;
-    m_category = StartsAndSizes;
-    m_data.starts_and_sizes.starts = nullptr;
-    m_data.starts_and_sizes.sizes = nullptr;
+    switch (m_category) {
+      case VArraySpanCategory::SingleArray:
+        return true;
+      case VArraySpanCategory::StartsAndSizes:
+        return m_virtual_size == 1;
+    }
+    BLI_assert(false);
+    return false;
   }
 
-  VArraySpan(Span<T> span, uint virtual_size)
+  bool is_empty() const
   {
-    m_virtual_size = virtual_size;
-    m_category = SingleArray;
-    m_data.single_array.start = span.data();
-    m_data.single_array.size = span.size();
+    return this->m_virtual_size == 0;
   }
 
-  VArraySpan(Span<const T *> starts, Span<uint> sizes)
+  uint size() const
   {
-    BLI_assert(starts.size() == sizes.size());
-    m_virtual_size = starts.size();
-    m_category = StartsAndSizes;
-    m_data.starts_and_sizes.starts = starts.begin();
-    m_data.starts_and_sizes.sizes = sizes.begin();
+    return this->m_virtual_size;
   }
+};
 
-  bool is_empty() const
+/**
+ * A virtual array span. Every element of this span contains a virtual span. So it behaves like
+ * a blender::Span, but might not be backed up by an actual array.
+ */
+template<typename T> class VArraySpan : public VArraySpanBase<T> {
+ private:
+  friend class GVArraySpan;
+
+  VArraySpan(const VArraySpanBase<void> &other)
   {
-    return m_virtual_size == 0;
+    memcpy(this, &other, sizeof(VArraySpanBase<void>));
   }
 
-  uint size() const
+ public:
+  VArraySpan()
+  {
+    this->m_virtual_size = 0;
+    this->m_category = VArraySpanCategory::StartsAndSizes;
+    this->m_data.starts_and_sizes.starts = nullptr;
+    this->m_data.starts_and_sizes.sizes = nullptr;
+  }
+
+  VArraySpan(Span<T> span, uint virtual_size)
   {
-    return m_virtual_size;
+    this->m_virtual_size = virtual_size;
+    this->m_category = VArraySpanCategory::SingleArray;
+    this->m_data.single_array.start = span.data();
+    this->m_data.single_array.size = span.size();
+  }
+
+  VArraySpan(Span<const T *> starts, Span<uint> sizes)
+  {
+    BLI_assert(starts.size() == sizes.size());
+    this->m_virtual_size = starts.size();
+    this->m_category = VArraySpanCategory::StartsAndSizes;
+    this->m_data.starts_and_sizes.starts = starts.begin();
+    this->m_data.starts_and_sizes.sizes = sizes.begin();
   }
 
   VSpan<T> operator[](uint index) const
   {
-    BLI_assert(index < m_virtual_size);
-    switch (m_category) {
-      case SingleArray:
-        return VSpan<T>(Span<T>(m_data.single_array.start, m_data.single_array.size));
-      case StartsAndSizes:
-        return VSpan<T>(
-            Span<T>(m_data.starts_and_sizes.starts[index], m_data.starts_and_sizes.sizes[index]));
+    BLI_assert(index < this->m_virtual_size);
+    switch (this->m_category) {
+      case VArraySpanCategory::SingleArray:
+        return VSpan<T>(Span<T>(this->m_data.single_array.start, this->m_data.single_array.size));
+      case VArraySpanCategory::StartsAndSizes:
+        return VSpan<T>(Span<T>(this->m_data.starts_and_sizes.starts[index],
+                                this->m_data.starts_and_sizes.sizes[index]));
     }
     BLI_assert(false);
     return {};
@@ -113,99 +136,65 @@ template<typename T> class VArraySpan {
  * A generic virtual array span. It's just like a VArraySpan, but the type is only known at
  * run-time.
  */
-class GVArraySpan {
+class GVArraySpan : public VArraySpanBase<void> {
  private:
-  /**
-   * Depending on the use case, the referenced data might have a different structure. More
-   * categories can be added when necessary.
-   */
-  enum Category {
-    SingleArray,
-    StartsAndSizes,
-  };
-
   const CPPType *m_type;
-  uint m_virtual_size;
-  Category m_category;
-
-  union {
-    struct {
-      const void *values;
-      uint size;
-    } single_array;
-    struct {
-      const void *const *starts;
-      const uint *sizes;
-    } starts_and_sizes;
-  } m_data;
 
   GVArraySpan() = default;
 
  public:
   GVArraySpan(const CPPType &type)
   {
-    m_type = &type;
-    m_virtual_size = 0;
-    m_category = StartsAndSizes;
-    m_data.starts_and_sizes.starts = nullptr;
-    m_data.starts_and_sizes.sizes = nullptr;
+    this->m_type = &type;
+    this->m_virtual_size = 0;
+    this->m_category = VArraySpanCategory::StartsAndSizes;
+    this->m_data.starts_and_sizes.starts = nullptr;
+    this->m_data.starts_and_sizes.sizes = nullptr;
   }
 
   GVArraySpan(GSpan array, uint virtual_size)
   {
-    m_type = &array.type();
-    m_virtual_size = virtual_size;
-    m_category = SingleArray;
-    m_data.single_array.values = array.buffer();
-    m_data.single_array.size = array.size();
+    this->m_type = &array.type();
+    this->m_virtual_size = virtual_size;
+    this->m_category = VArraySpanCategory::SingleArray;
+    this->m_data.single_array.start = array.buffer();
+    this->m_data.single_array.size = array.size();
   }
 
   GVArraySpan(const CPPType &type, Span<const void *> starts, Span<uint> sizes)
   {
     BLI_assert(starts.size() == sizes.size());
-    m_type = &type;
-    m_virtual_size = starts.size();
-    m_category = StartsAndSizes;
-    m_data.starts_and_sizes.starts = starts.begin();
-    m_data.starts_and_sizes.sizes = sizes.begin();
-  }
-
-  bool is_empty() const
-  {
-    return m_virtual_size == 0;
+    this->m_type = &type;
+    this->m_virtual_size = starts.size();
+    this->m_category = VArraySpanCategory::StartsAndSizes;
+    this->m_data.starts_and_sizes.starts = (void **)starts.begin();
+    this->m_data.starts_and_sizes.sizes = sizes.begin();
   }
 
-  uint size() const
+  template<typename T> GVArraySpan(VArraySpan<T> other)
   {
-    return m_virtual_size;
+    this->m_type = &CPPType::get<T>();
+    memcpy(this, &other, sizeof(VArraySpanBase<void>));
   }
 
   const CPPType &type() const
   {
-    return *m_type;
+    return *this->m_type;
   }
 
   template<typename T> VArraySpan<T> typed() const
   {
     BLI_assert(CPPType::get<T>() == *m_type);
-    switch (m_category) {
-      case SingleArray:
-        return VArraySpan<T>(
-            Span<T>((const T *)m_data.single_array.values, m_data.single_array.size));
-      case StartsAndSizes:
-        return VArraySpan<T>(
-            Span<const T *>((const T *const *)m_data.starts_and_sizes.starts, m_virtual_size),
-            Span<uint>(m_data.starts_and_sizes.sizes, m_virtual_size));
-    }
+    return VArraySpan<T>(*this);
   }
 
   GVSpan operator[](uint index) const
   {
     BLI_assert(index < m_virtual_size);
     switch (m_category) {
-      case SingleArray:
-        return GVSpan(GSpan(*m_type, m_data.single_array.values, m_data.single_array.size));
-      case StartsAndSizes:
+      case VArraySpanCategory::SingleArray:
+        return GVSpan(GSpan(*m_type, m_data.single_array.start, m_data.single_array.size));
+      case VArraySpanCategory::StartsAndSizes:
         return GVSpan(GSpan(
             *m_type, m_data.starts_and_sizes.starts[index], m_data.starts_and_sizes.sizes[index]));
     }
diff --git a/source/blender/functions/FN_cpp_type.hh b/source/blender/functions/FN_cpp_type.hh
index df5218ed5ba..c6e440251ee 100644
--- a/source/blender/functions/FN_cpp_type.hh
+++ b/source/blender/functions/FN_cpp_type.hh
@@ -725,6 +725,8 @@ static std::unique_ptr<const CPPType> create_cpp_type(StringRef name, const T &d
   const blender::fn::CPPType &CPPType_##IDENTIFIER = *CPPTYPE_##IDENTIFIER##_owner; \
   template<> const blender::fn::CPPType &blender::fn::CPPType::get<TYPE_NAME>() \
   { \
+    /* This can happen when trying to access a CPPType during static storage initialization. */ \
+    BLI_assert(CPPTYPE_##IDENTIFIER##_owner.get() != nullptr); \
     return CPPType_##IDENTIFIER; \
   }
 
diff --git a/source/blender/functions/FN_generic_vector_array.hh b/source/blender/functions/FN_generic_vector_array.hh
index 90c0fbad136..1c8f74f2abe 100644
--- a/source/blender/functions/FN_generic_vector_array.hh
+++ b/source/blender/functions/FN_generic_vector_array.hh
@@ -119,6 +119,23 @@ class GVectorArray : NonCopyable, NonMovable {
     m_lengths[index]++;
   }
 
+  void extend(uint index, GVSpan span)
+  {
+    BLI_assert(m_type == span.type());
+    for (uint i = 0; i < span.size(); i++) {
+      this->append(index, span[i]);
+    }
+  }
+
+  void extend(IndexMask mask, GVArraySpan array_span)
+  {
+    BLI_assert(m_type == array_span.type());
+    BLI_assert(mask.min_array_size() <= array_span.size());
+    for (uint i : mask) {
+      this->extend(i, array_span[i]);
+    }
+  }
+
   GMutableSpan operator[](uint index)
   {
     BLI_assert(index < m_starts.size());
@@ -158,6 +175,16 @@ template<typename T> class GVectorArrayRef {
     m_vector_array->append(index, &value);
   }
 
+  void extend(uint index, Span<T> values)
+  {
+    m_vector_array->extend(index, values);
+  }
+
+  void extend(uint index, VSpan<T> values)
+  {
+    m_vector_array->extend(index, GVSpan(values));
+  }
+
   MutableSpan<T> operator[](uint index)
   {
     BLI_assert(index < m_vector_array->m_starts.size());
diff --git a/source/blender/functions/FN_multi_function_data_type.hh b/source/blender/functions/FN_multi_function_data_type.hh
index 856d1cc7253..1a7b179c6ae 100644
--- a/source/blender/functions/FN_multi_function_data_type.hh
+++ b/source/blender/functions/FN_multi_function_data_type.hh
@@ -46,6 +46,8 @@ class MFDataType {
   }
 
  public:
+  MFDataType() = default;
+
   static MFDataType ForSingle(const CPPType &type)
   {
     return MFDataType(Single, type);
diff --git a/source/blender/functions/FN_multi_function_param_type.hh b/source/blender/functions/FN_multi_function_param_type.hh
index bd31a793b21..d89c13505f9 100644
--- a/source/blender/functions/FN_multi_function_param_type.hh
+++ b/source/blender/functions/FN_multi_function_param_type.hh
@@ -135,6 +135,16 @@ class MFParamType {
     return SingleInput;
   }
 
+  bool is_input_or_mutable() const
+  {
+    return ELEM(m_interface_type, Input, Mutable);
+  }
+
+  bool is_output_or_mutable() const
+  {
+    return ELEM(m_interface_type, Output, Mutable);
+  }
+
   friend bool operator==(const MFParamType &a, const MFParamType &b);
   friend bool operator!=(const MFParamType &a, const MFParamType &b);
 };
diff --git a/source/blender/functions/FN_multi_function_params.hh b/source/blender/functions/FN_multi_function_params.hh
index 2388e3b15e5..6a0eb698250 100644
--- a/source/blender/functions/FN_multi_function_params.hh
+++ b/source/blender/functions/FN_multi_function_params.hh
@@ -199,6 +199,10 @@ class MFParams {
     return m_builder->m_mutable_spans[data_index];
   }
 
+  template<typename T> GVectorArrayRef<T> vector_mutable(uint param_index, StringRef name = "")
+  {
+    return this->vector_mutable(param_index, name).typed<T>();
+  }
   GVectorArray &vector_mutable(uint param_index, StringRef name = "")
   {
     this->assert_correct_param(param_index, name, MFParamType::VectorMutable);
diff --git a/source/blender/functions/FN_multi_function_signature.hh b/source/blender/functions/FN_multi_function_signature.hh
index 64bbd7be5c8..73e7b9bb18d 100644
--- a/source/blender/functions/FN_multi_function_signature.hh
+++ b/source/blender/functions/FN_multi_function_signature.hh
@@ -135,6 +135,10 @@ class MFSignatureBuilder {
   {
     this->mutable_(name, MFDataType::ForSingle(type));
   }
+  template<typename T> void vector_mutable(StringRef name)
+  {
+    this->vector_mutable(name, CPPType::get<T>());
+  }
   void vector_mutable(StringRef name, const CPPType &base_type)
   {
     this->mutable_(name, MFDataType::ForVector(base_type));
diff --git a/source/blender/functions/FN_spans.hh b/source/blender/functions/FN_spans.hh
index c206662231c..d3497c05eb9 100644
--- a/source/blender/functions/FN_spans.hh
+++ b/source/blender/functions/FN_spans.hh
@@ -68,7 +68,7 @@ class GSpan {
   }
 
   template<typename T>
-  GSpan(Span<T> array) : GSpan(CPPType::get<T>(), (const void *)array.begin(), array.size())
+  GSpan(Span<T> array) : GSpan(CPPType::get<T>(), (const void *)array.data(), array.size())
   {
   }
 
@@ -171,21 +171,16 @@ class GMutableSpan {
   }
 };
 
-/**
- * A virtual span. It behaves like a blender::Span<T>, but might not be backed up by an actual
- * array.
- */
-template<typename T> class VSpan {
- private:
-  enum Category {
-    Single,
-    FullArray,
-    FullPointerArray,
-  };
+enum class VSpanCategory {
+  Single,
+  FullArray,
+  FullPointerArray,
+};
 
+template<typename T> struct VSpanBase {
+ protected:
   uint m_virtual_size;
-  Category m_category;
-
+  VSpanCategory m_category;
   union {
     struct {
       const T *data;
@@ -199,18 +194,59 @@ template<typename T> class VSpan {
   } m_data;
 
  public:
+  bool is_single_element() const
+  {
+    switch (m_category) {
+      case VSpanCategory::Single:
+        return true;
+      case VSpanCategory::FullArray:
+        return m_virtual_size == 1;
+      case VSpanCategory::FullPointerArray:
+        return m_virtual_size == 1;
+    }
+    BLI_assert(false);
+    return false;
+  }
+
+  bool is_empty() const
+  {
+    return this->m_virtual_size == 0;
+  }
+
+  uint size() const
+  {
+    return this->m_virtual_size;
+  }
+};
+
+BLI_STATIC_ASSERT((sizeof(VSpanBase<void>) == sizeof(VSpanBase<AlignedBuffer<64, 64>>)),
+                  "should not depend on the size of the type");
+
+/**
+ * A virtual span. It behaves like a blender::Span<T>, but might not be backed up by an actual
+ * array.
+ */
+template<typename T> class VSpan : public VSpanBase<T> {
+  friend class GVSpan;
+
+  VSpan(const VSpanBase<void> &values)
+  {
+    memcpy(this, &values, sizeof(VSpanBase<void>));
+  }
+
+ public:
   VSpan()
   {
-    m_virtual_size = 0;
-    m_category = FullArray;
-    m_data.full_array.data = nullptr;
+    this->m_virtual_size = 0;
+    this->m_category = VSpanCategory::FullArray;
+    this->m_data.full_array.data = nullptr;
   }
 
   VSpan(Span<T> values)
   {
-    m_virtual_size = values.size();
-    m_category = FullArray;
-    m_data.full_array.data = values.begin();
+    this->m_virtual_size = values.size();
+    this->m_category = VSpanCategory::FullArray;
+    this->m_data.full_array.data = values.begin();
   }
 
   VSpan(MutableSpan<T> values) : VSpan(Span<T>(values))
@@ -219,43 +255,33 @@ template<typename T> class VSpan {
 
   VSpan(Span<const T *> values)
   {
-    m_virtual_size = values.size();
-    m_category = FullPointerArray;
-    m_data.full_pointer_array.data = values.begin();
+    this->m_virtual_size = values.size();
+    this->m_category = VSpanCategory::FullPointerArray;
+    this->m_data.full_pointer_array.data = values.begin();
   }
 
   static VSpan FromSingle(const T *value, uint virtual_size)
   {
     VSpan ref;
     ref.m_virtual_size = virtual_size;
-    ref.m_category = Single;
+    ref.m_category = VSpanCategory::Single;
     ref.m_data.single.data = value;
     return ref;
   }
 
   const T &operator[](uint index) const
   {
-    BLI_assert(index < m_virtual_size);
-    switch (m_category) {
-      case Single:
-        return *m_data.single.data;
-      case FullArray:
-        return m_data.full_array.data[index];
-      case FullPointerArray:
-        return *m_data.full_pointer_array.data[index];
+    BLI_assert(index < this->m_virtual_size);
+    switch (this->m_category) {
+      case VSpanCategory::Single:
+        return *this->m_data.single.data;
+      case VSpanCategory::FullArray:
+        return this->m_data.full_array.data[index];
+      case VSpanCategory::FullPointerArray:
+        return *this->m_data.full_pointer_array.data[index];
     }
     BLI_assert(false);
-    return *m_data.single.data;
-  }
-
-  bool is_empty() const
-  {
-    return m_virtual_size == 0;
-  }
-
-  uint size() const
-  {
-    return m_virtual_size;
+    return *this->m_data.single.data;
   }
 };
 
@@ -263,53 +289,39 @@ template<typename T> class VSpan {
  * A generic virtual span. It behaves like a blender::Span<T>, but the type is only known at
  * run-time and it might not be backed up by an actual array.
  */
-class GVSpan {
+class GVSpan : public VSpanBase<void> {
  private:
-  enum Category {
-    Single,
-    FullArray,
-    FullPointerArray,
-  };
-
   const CPPType *m_type;
-  uint m_virtual_size;
-  Category m_category;
-
-  union {
-    struct {
-      const void *data;
-    } single;
-    struct {
-      const void *data;
-    } full_array;
-    struct {
-      const void *const *data;
-    } full_pointer_array;
-  } m_data;
 
   GVSpan() = default;
 
  public:
   GVSpan(const CPPType &type)
   {
-    m_type = &type;
-    m_virtual_size = 0;
-    m_category = FullArray;
-    m_data.full_array.data = nullptr;
+    this->m_type = &type;
+    this->m_virtual_size = 0;
+    this->m_category = VSpanCategory::FullArray;
+    this->m_data.full_array.data = nullptr;
   }
 
   GVSpan(GSpan values)
   {
-    m_type = &values.type();
-    m_virtual_size = values.size();
-    m_category = FullArray;
-    m_data.full_array.data = values.buffer();
+    this->m_type = &values.type();
+    this->m_virtual_size = values.size();
+    this->m_category = VSpanCategory::FullArray;
+    this->m_data.full_array.data = values.buffer();
   }
 
   GVSpan(GMutableSpan values) : GVSpan(GSpan(values))
   {
   }
 
+  template<typename T> GVSpan(const VSpanBase<T> &values)
+  {
+    this->m_type = &CPPType::get<T>();
+    memcpy(this, &values, sizeof(VSpanBase<void>));
+  }
+
   template<typename T> GVSpan(Span<T> values) : GVSpan(GSpan(values))
   {
   }
@@ -323,7 +335,7 @@ class GVSpan {
     GVSpan ref;
     ref.m_type = &type;
     ref.m_virtual_size = virtual_size;
-    ref.m_category = Single;
+    ref.m_category = VSpanCategory::Single;
     ref.m_data.single.data = value;
     return ref;
   }
@@ -333,55 +345,56 @@ class GVSpan {
     GVSpan ref;
     ref.m_type = &type;
     ref.m_virtual_size = size;
-    ref.m_category = FullPointerArray;
+    ref.m_category = VSpanCategory::FullPointerArray;
     ref.m_data.full_pointer_array.data = values;
     return ref;
   }
 
-  bool is_empty() const
-  {
-    return m_virtual_size == 0;
-  }
-
-  uint size() const
-  {
-    return m_virtual_size;
-  }
-
   const CPPType &type() const
   {
-    return *m_type;
+    return *this->m_type;
   }
 
   const void *operator[](uint index) const
   {
-    BLI_assert(index < m_virtual_size);
-    switch (m_category) {
-      case Single:
-        return m_data.single.data;
-      case FullArray:
-        return POINTER_OFFSET(m_data.full_array.data, index * m_type->size());
-      case FullPointerArray:
-        return m_data.full_pointer_array.data[index];
+    BLI_assert(index < this->m_virtual_size);
+    switch (this->m_category) {
+      case VSpanCategory::Single:
+        return this->m_data.single.data;
+      case VSpanCategory::FullArray:
+        return POINTER_OFFSET(this->m_data.full_array.data, index * m_type->size());
+      case VSpanCategory::FullPointerArray:
+        return this->m_data.full_pointer_array.data[index];
     }
     BLI_assert(false);
-    return m_data.single.data;
+    return this->m_data.single.data;
   }
 
   template<typename T> VSpan<T> typed() const
   {
     BLI_assert(CPPType::get<T>() == *m_type);
-    switch (m_category) {
-      case Single:
-        return VSpan<T>::FromSingle((const T *)m_data.single.data, m_virtual_size);
-      case FullArray:
-        return VSpan<T>(Span<T>((const T *)m_data.full_array.data, m_virtual_size));
-      case FullPointerArray:
-        return VSpan<T>(
-            Span<const T *>((const T *const *)m_data.full_pointer_array.data, m_virtual_size));
+    return VSpan<T>(*this);
+  }
+
+  const void *as_single_element() const
+  {
+    BLI_assert(this->is_single_element());
+    return (*this)[0];
+  }
+
+  void materialize_to_uninitialized(void *dst) const
+  {
+    this->materialize_to_uninitialized(IndexRange(m_virtual_size), dst);
+  }
+
+  void materialize_to_uninitialized(IndexMask mask, void *dst) const
+  {
+    BLI_assert(this->size() >= mask.min_array_size());
+
+    uint element_size = m_type->size();
+    for (uint i : mask) {
+      m_type->copy_to_uninitialized((*this)[i], POINTER_OFFSET(dst, element_size * i));
     }
-    BLI_assert(false);
-    return {};
   }
 };