Geometry Nodes: initial attribute interpolation between domains

This patch adds support for accessing corner attributes on the point domain.
The immediate benefit of this is that now (interpolated) uv coordinates are
available on points without having to use the Point Distribute node.

This is also very useful for parts of T84297, because once we have vertex
colors, those will also be available on points, even though they are stored
per corner.

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

5 files changed, 323 insertions, 5 deletions

diff --git a/source/blender/blenkernel/BKE_attribute_math.hh b/source/blender/blenkernel/BKE_attribute_math.hh
index b8bb2048d9d..606811ca12d 100644
--- a/source/blender/blenkernel/BKE_attribute_math.hh
+++ b/source/blender/blenkernel/BKE_attribute_math.hh
@@ -14,9 +14,11 @@
  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  */
 
+#include "BLI_array.hh"
 #include "BLI_color.hh"
 #include "BLI_float2.hh"
 #include "BLI_float3.hh"
+
 #include "DNA_customdata_types.h"
 
 namespace blender::attribute_math {
@@ -52,7 +54,12 @@ void convert_to_static_type(const CustomDataType data_type, const Func &func)
   }
 }
 
-/* Interpolate between three values. */
+/* -------------------------------------------------------------------- */
+/** \name Mix three values of the same type.
+ *
+ * This is typically used to interpolate values within a triangle.
+ * \{ */
+
 template<typename T> T mix3(const float3 &weights, const T &v0, const T &v1, const T &v2);
 
 template<> inline bool mix3(const float3 &weights, const bool &v0, const bool &v1, const bool &v2)
@@ -91,4 +98,159 @@ inline Color4f mix3(const float3 &weights, const Color4f &v0, const Color4f &v1,
   return result;
 }
 
+/** \} */
+
+/* -------------------------------------------------------------------- */
+/** \name Mix a dynamic amount of values with weights for many elements.
+ *
+ * This section provides an abstraction for "mixers". The abstraction encapsulates details about
+ * how different types should be mixed. Usually #DefaultMixer<T> should be used to get a mixer for
+ * a specific type.
+ * \{ */
+
+template<typename T> class SimpleMixer {
+ private:
+  MutableSpan<T> buffer_;
+  T default_value_;
+  Array<float> total_weights_;
+
+ public:
+  /**
+   * \param buffer: Span where the interpolated values should be stored.
+   * \param default_value: Output value for an element that has not been affected by a #mix_in.
+   */
+  SimpleMixer(MutableSpan<T> buffer, T default_value = {})
+      : buffer_(buffer), default_value_(default_value), total_weights_(buffer.size(), 0.0f)
+  {
+    BLI_STATIC_ASSERT(std::is_trivial_v<T>, "");
+    memset(buffer_.data(), 0, sizeof(T) * buffer_.size());
+  }
+
+  /**
+   * Mix a #value into the element with the given #index.
+   */
+  void mix_in(const int64_t index, const T &value, const float weight = 1.0f)
+  {
+    BLI_assert(weight >= 0.0f);
+    buffer_[index] += value * weight;
+    total_weights_[index] += weight;
+  }
+
+  /**
+   * Has to be called before the buffer provided in the constructor is used.
+   */
+  void finalize()
+  {
+    for (const int64_t i : buffer_.index_range()) {
+      const float weight = total_weights_[i];
+      if (weight > 0.0f) {
+        buffer_[i] *= 1.0f / weight;
+      }
+      else {
+        buffer_[i] = default_value_;
+      }
+    }
+  }
+};
+
+/** This mixer accumulates values in a type that is different from the one that is mixed. Some
+ * types cannot encode the floating point weights in their values (e.g. int and bool). */
+template<typename T, typename AccumulationT, T (*ConvertToT)(const AccumulationT &value)>
+class SimpleMixerWithAccumulationType {
+ private:
+  struct Item {
+    /* Store both values together, because they are accessed together. */
+    AccumulationT value = {0};
+    float weight = 0.0f;
+  };
+
+  MutableSpan<T> buffer_;
+  T default_value_;
+  Array<Item> accumulation_buffer_;
+
+ public:
+  SimpleMixerWithAccumulationType(MutableSpan<T> buffer, T default_value = {})
+      : buffer_(buffer), default_value_(default_value), accumulation_buffer_(buffer.size())
+  {
+  }
+
+  void mix_in(const int64_t index, const T &value, const float weight = 1.0f)
+  {
+    const AccumulationT converted_value = static_cast<AccumulationT>(value);
+    Item &item = accumulation_buffer_[index];
+    item.value += converted_value * weight;
+    item.weight += weight;
+  }
+
+  void finalize()
+  {
+    for (const int64_t i : buffer_.index_range()) {
+      const Item &item = accumulation_buffer_[i];
+      if (item.weight > 0.0f) {
+        const float weight_inv = 1.0f / item.weight;
+        const T converted_value = ConvertToT(item.value * weight_inv);
+        buffer_[i] = converted_value;
+      }
+      else {
+        buffer_[i] = default_value_;
+      }
+    }
+  }
+};
+
+class Color4fMixer {
+ private:
+  MutableSpan<Color4f> buffer_;
+  Color4f default_color_;
+  Array<float> total_weights_;
+
+ public:
+  Color4fMixer(MutableSpan<Color4f> buffer, Color4f default_color = {0, 0, 0, 1});
+  void mix_in(const int64_t index, const Color4f &color, const float weight = 1.0f);
+  void finalize();
+};
+
+template<typename T> struct DefaultMixerStruct {
+  /* Use void by default. This can be check for in `if constexpr` statements. */
+  using type = void;
+};
+template<> struct DefaultMixerStruct<float> {
+  using type = SimpleMixer<float>;
+};
+template<> struct DefaultMixerStruct<float2> {
+  using type = SimpleMixer<float2>;
+};
+template<> struct DefaultMixerStruct<float3> {
+  using type = SimpleMixer<float3>;
+};
+template<> struct DefaultMixerStruct<Color4f> {
+  /* Use a special mixer for colors. Color4f can't be added/multiplied, because this is not
+   * something one should usually do with colors.  */
+  using type = Color4fMixer;
+};
+template<> struct DefaultMixerStruct<int> {
+  static int double_to_int(const double &value)
+  {
+    return static_cast<int>(value);
+  }
+  /* Store interpolated ints in a double temporarily, so that weights are handled correctly. It
+   * uses double instead of float so that it is accurate for all 32 bit integers. */
+  using type = SimpleMixerWithAccumulationType<int, double, double_to_int>;
+};
+template<> struct DefaultMixerStruct<bool> {
+  static bool float_to_bool(const float &value)
+  {
+    return value >= 0.5f;
+  }
+  /* Store interpolated bools in a float temporary. Otherwise information provided by weights is
+   * easily rounded away. */
+  using type = SimpleMixerWithAccumulationType<bool, float, float_to_bool>;
+};
+
+/* Utility to get a good default mixer for a given type. This is `void` when there is no default
+ * mixer for the given type. */
+template<typename T> using DefaultMixer = typename DefaultMixerStruct<T>::type;
+
+/** \} */
+
 }  // namespace blender::attribute_math
diff --git a/source/blender/blenkernel/BKE_geometry_set.hh b/source/blender/blenkernel/BKE_geometry_set.hh
index 9a871574f6f..6a987ee51e0 100644
--- a/source/blender/blenkernel/BKE_geometry_set.hh
+++ b/source/blender/blenkernel/BKE_geometry_set.hh
@@ -174,7 +174,7 @@ class GeometryComponent {
    * interpolate from one domain to another.
    * Returns null if the interpolation is not implemented. */
   virtual blender::bke::ReadAttributePtr attribute_try_adapt_domain(
-      blender::bke::ReadAttributePtr attribute, const AttributeDomain domain) const;
+      blender::bke::ReadAttributePtr attribute, const AttributeDomain new_domain) const;
 
   /* Returns true when the attribute has been deleted. */
   bool attribute_try_delete(const blender::StringRef attribute_name);
@@ -368,6 +368,8 @@ class MeshComponent : public GeometryComponent {
   Mesh *get_for_write();
 
   int attribute_domain_size(const AttributeDomain domain) const final;
+  blender::bke::ReadAttributePtr attribute_try_adapt_domain(
+      blender::bke::ReadAttributePtr attribute, const AttributeDomain new_domain) const final;
 
   bool is_empty() const final;
 
diff --git a/source/blender/blenkernel/CMakeLists.txt b/source/blender/blenkernel/CMakeLists.txt
index 3f22612652c..6b6d2b45d02 100644
--- a/source/blender/blenkernel/CMakeLists.txt
+++ b/source/blender/blenkernel/CMakeLists.txt
@@ -81,6 +81,7 @@ set(SRC
   intern/asset.cc
   intern/attribute.c
   intern/attribute_access.cc
+  intern/attribute_math.cc
   intern/autoexec.c
   intern/blender.c
   intern/blender_copybuffer.c
diff --git a/source/blender/blenkernel/intern/attribute_access.cc b/source/blender/blenkernel/intern/attribute_access.cc
index 772309349ff..cc833e094c8 100644
--- a/source/blender/blenkernel/intern/attribute_access.cc
+++ b/source/blender/blenkernel/intern/attribute_access.cc
@@ -17,6 +17,7 @@
 #include <utility>
 
 #include "BKE_attribute_access.hh"
+#include "BKE_attribute_math.hh"
 #include "BKE_customdata.h"
 #include "BKE_deform.h"
 #include "BKE_geometry_set.hh"
@@ -323,6 +324,29 @@ template<typename T> class ArrayReadAttribute final : public ReadAttribute {
   }
 };
 
+template<typename T> class OwnedArrayReadAttribute final : public ReadAttribute {
+ private:
+  Array<T> data_;
+
+ public:
+  OwnedArrayReadAttribute(AttributeDomain domain, Array<T> data)
+      : ReadAttribute(domain, CPPType::get<T>(), data.size()), data_(std::move(data))
+  {
+  }
+
+  void get_internal(const int64_t index, void *r_value) const override
+  {
+    new (r_value) T(data_[index]);
+  }
+
+  void initialize_span() const override
+  {
+    /* The data will not be modified, so this const_cast is fine. */
+    array_buffer_ = const_cast<T *>(data_.data());
+    array_is_temporary_ = false;
+  }
+};
+
 template<typename StructT,
          typename ElemT,
          ElemT (*GetFunc)(const StructT &),
@@ -1344,10 +1368,10 @@ ReadAttributePtr GeometryComponent::attribute_try_get_for_read(
   return {};
 }
 
-ReadAttributePtr GeometryComponent::attribute_try_adapt_domain(ReadAttributePtr attribute,
-                                                               const AttributeDomain domain) const
+ReadAttributePtr GeometryComponent::attribute_try_adapt_domain(
+    ReadAttributePtr attribute, const AttributeDomain new_domain) const
 {
-  if (attribute && attribute->domain() == domain) {
+  if (attribute && attribute->domain() == new_domain) {
     return attribute;
   }
   return {};
@@ -1765,4 +1789,75 @@ int MeshComponent::attribute_domain_size(const AttributeDomain domain) const
   return 0;
 }
 
+namespace blender::bke {
+
+template<typename T>
+void adapt_mesh_domain_corner_to_point_impl(const Mesh &mesh,
+                                            const TypedReadAttribute<T> &attribute,
+                                            MutableSpan<T> r_values)
+{
+  BLI_assert(r_values.size() == mesh.totvert);
+  attribute_math::DefaultMixer<T> mixer(r_values);
+
+  for (const int loop_index : IndexRange(mesh.totloop)) {
+    const T value = attribute[loop_index];
+    const MLoop &loop = mesh.mloop[loop_index];
+    const int point_index = loop.v;
+    mixer.mix_in(point_index, value);
+  }
+  mixer.finalize();
+}
+
+static ReadAttributePtr adapt_mesh_domain_corner_to_point(const Mesh &mesh,
+                                                          ReadAttributePtr attribute)
+{
+  ReadAttributePtr new_attribute;
+  const CustomDataType data_type = attribute->custom_data_type();
+  attribute_math::convert_to_static_type(data_type, [&](auto dummy) {
+    using T = decltype(dummy);
+    if constexpr (!std::is_void_v<attribute_math::DefaultMixer<T>>) {
+      /* We compute all interpolated values at once, because for this interpolation, one has to
+       * iterate over all loops anyway. */
+      Array<T> values(mesh.totvert);
+      adapt_mesh_domain_corner_to_point_impl<T>(mesh, *attribute, values);
+      new_attribute = std::make_unique<OwnedArrayReadAttribute<T>>(ATTR_DOMAIN_POINT,
+                                                                   std::move(values));
+    }
+  });
+  return new_attribute;
+}
+
+}  // namespace blender::bke
+
+ReadAttributePtr MeshComponent::attribute_try_adapt_domain(ReadAttributePtr attribute,
+                                                           const AttributeDomain new_domain) const
+{
+  if (!attribute) {
+    return {};
+  }
+  if (attribute->size() == 0) {
+    return {};
+  }
+  const AttributeDomain old_domain = attribute->domain();
+  if (old_domain == new_domain) {
+    return attribute;
+  }
+
+  switch (old_domain) {
+    case ATTR_DOMAIN_CORNER: {
+      switch (new_domain) {
+        case ATTR_DOMAIN_POINT:
+          return blender::bke::adapt_mesh_domain_corner_to_point(*mesh_, std::move(attribute));
+        default:
+          break;
+      }
+      break;
+    }
+    default:
+      break;
+  }
+
+  return {};
+}
+
 /** \} */
diff --git a/source/blender/blenkernel/intern/attribute_math.cc b/source/blender/blenkernel/intern/attribute_math.cc
new file mode 100644
index 00000000000..4ff3a6ceff5
--- /dev/null
+++ b/source/blender/blenkernel/intern/attribute_math.cc
@@ -0,0 +1,58 @@
+/*
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ */
+
+#include "BKE_attribute_math.hh"
+
+namespace blender::attribute_math {
+
+Color4fMixer::Color4fMixer(MutableSpan<Color4f> output_buffer, Color4f default_color)
+    : buffer_(output_buffer),
+      default_color_(default_color),
+      total_weights_(output_buffer.size(), 0.0f)
+{
+  buffer_.fill(Color4f(0, 0, 0, 0));
+}
+
+void Color4fMixer::mix_in(const int64_t index, const Color4f &color, const float weight)
+{
+  BLI_assert(weight >= 0.0f);
+  Color4f &output_color = buffer_[index];
+  output_color.r += color.r * weight;
+  output_color.g += color.g * weight;
+  output_color.b += color.b * weight;
+  output_color.a += color.a * weight;
+  total_weights_[index] += weight;
+}
+
+void Color4fMixer::finalize()
+{
+  for (const int64_t i : buffer_.index_range()) {
+    const float weight = total_weights_[i];
+    Color4f &output_color = buffer_[i];
+    if (weight > 0.0f) {
+      const float weight_inv = 1.0f / weight;
+      output_color.r *= weight_inv;
+      output_color.g *= weight_inv;
+      output_color.b *= weight_inv;
+      output_color.a *= weight_inv;
+    }
+    else {
+      output_color = default_color_;
+    }
+  }
+}
+
+}  // namespace blender::attribute_math