Geometry Nodes: transfer corner and point attributes in Point Distribute node

If the mesh has any corner or point attributes (e.g. vertex weights or
uv maps), those attributes will now be available on the generated points
as well.

Other domains can be supported as well. I just did not implement those yet,
because we don't have a use case for them.

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

1 files changed, 142 insertions, 23 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_point_distribute.cc b/source/blender/nodes/geometry/nodes/node_geo_point_distribute.cc
index 5e9b780db6d..eaf13b94eb9 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_point_distribute.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_point_distribute.cc
@@ -26,6 +26,7 @@
 #include "DNA_meshdata_types.h"
 #include "DNA_pointcloud_types.h"
 
+#include "BKE_attribute_math.hh"
 #include "BKE_bvhutils.h"
 #include "BKE_deform.h"
 #include "BKE_mesh.h"
@@ -217,14 +218,16 @@ BLI_NOINLINE static void eliminate_points_based_on_mask(Span<bool> elimination_m
   }
 }
 
-BLI_NOINLINE static void compute_special_attributes(const Mesh &mesh,
-                                                    Span<float3> bary_coords,
-                                                    Span<int> looptri_indices,
-                                                    MutableSpan<float3> r_normals,
-                                                    MutableSpan<int> r_ids,
-                                                    MutableSpan<float3> r_rotations)
+template<typename T>
+BLI_NOINLINE static void interpolate_attribute_point(const Mesh &mesh,
+                                                     const Span<float3> bary_coords,
+                                                     const Span<int> looptri_indices,
+                                                     const Span<T> data_in,
+                                                     MutableSpan<T> data_out)
 {
+  BLI_assert(data_in.size() == mesh.totvert);
   Span<MLoopTri> looptris = get_mesh_looptris(mesh);
+
   for (const int i : bary_coords.index_range()) {
     const int looptri_index = looptri_indices[i];
     const MLoopTri &looptri = looptris[looptri_index];
@@ -233,20 +236,111 @@ BLI_NOINLINE static void compute_special_attributes(const Mesh &mesh,
     const int v0_index = mesh.mloop[looptri.tri[0]].v;
     const int v1_index = mesh.mloop[looptri.tri[1]].v;
     const int v2_index = mesh.mloop[looptri.tri[2]].v;
-    const float3 v0_pos = mesh.mvert[v0_index].co;
-    const float3 v1_pos = mesh.mvert[v1_index].co;
-    const float3 v2_pos = mesh.mvert[v2_index].co;
 
-    r_ids[i] = (int)(bary_coord.hash()) + looptri_index;
-    normal_tri_v3(r_normals[i], v0_pos, v1_pos, v2_pos);
-    r_rotations[i] = normal_to_euler_rotation(r_normals[i]);
+    const T &v0 = data_in[v0_index];
+    const T &v1 = data_in[v1_index];
+    const T &v2 = data_in[v2_index];
+
+    const T interpolated_value = attribute_math::mix3(bary_coord, v0, v1, v2);
+    data_out[i] = interpolated_value;
   }
 }
 
-BLI_NOINLINE static void add_remaining_point_attributes(const Mesh &mesh,
-                                                        GeometryComponent &component,
-                                                        Span<float3> bary_coords,
-                                                        Span<int> looptri_indices)
+template<typename T>
+BLI_NOINLINE static void interpolate_attribute_corner(const Mesh &mesh,
+                                                      const Span<float3> bary_coords,
+                                                      const Span<int> looptri_indices,
+                                                      const Span<T> data_in,
+                                                      MutableSpan<T> data_out)
+{
+  BLI_assert(data_in.size() == mesh.totloop);
+  Span<MLoopTri> looptris = get_mesh_looptris(mesh);
+
+  for (const int i : bary_coords.index_range()) {
+    const int looptri_index = looptri_indices[i];
+    const MLoopTri &looptri = looptris[looptri_index];
+    const float3 &bary_coord = bary_coords[i];
+
+    const int loop_index_0 = looptri.tri[0];
+    const int loop_index_1 = looptri.tri[1];
+    const int loop_index_2 = looptri.tri[2];
+
+    const T &v0 = data_in[loop_index_0];
+    const T &v1 = data_in[loop_index_1];
+    const T &v2 = data_in[loop_index_2];
+
+    const T interpolated_value = attribute_math::mix3(bary_coord, v0, v1, v2);
+    data_out[i] = interpolated_value;
+  }
+}
+
+BLI_NOINLINE static void interpolate_attribute(const Mesh &mesh,
+                                               Span<float3> bary_coords,
+                                               Span<int> looptri_indices,
+                                               const StringRef attribute_name,
+                                               const ReadAttribute &attribute_in,
+                                               GeometryComponent &component)
+{
+  const CustomDataType data_type = attribute_in.custom_data_type();
+  const AttributeDomain domain = attribute_in.domain();
+  if (!ELEM(domain, ATTR_DOMAIN_POINT, ATTR_DOMAIN_CORNER)) {
+    /* Not supported currently. */
+    return;
+  }
+
+  OutputAttributePtr attribute_out = component.attribute_try_get_for_output(
+      attribute_name, ATTR_DOMAIN_POINT, data_type);
+  if (!attribute_out) {
+    return;
+  }
+
+  attribute_math::convert_to_static_type(data_type, [&](auto dummy) {
+    using T = decltype(dummy);
+
+    Span data_in = attribute_in.get_span<T>();
+    MutableSpan data_out = attribute_out->get_span_for_write_only<T>();
+
+    switch (domain) {
+      case ATTR_DOMAIN_POINT: {
+        interpolate_attribute_point<T>(mesh, bary_coords, looptri_indices, data_in, data_out);
+        break;
+      }
+      case ATTR_DOMAIN_CORNER: {
+        interpolate_attribute_corner<T>(mesh, bary_coords, looptri_indices, data_in, data_out);
+        break;
+      }
+      default: {
+        BLI_assert(false);
+        break;
+      }
+    }
+  });
+  attribute_out.apply_span_and_save();
+}
+
+BLI_NOINLINE static void interpolate_existing_attributes(const MeshComponent &mesh_component,
+                                                         GeometryComponent &component,
+                                                         Span<float3> bary_coords,
+                                                         Span<int> looptri_indices)
+{
+  const Mesh &mesh = *mesh_component.get_for_read();
+
+  Set<std::string> attribute_names = mesh_component.attribute_names();
+  for (StringRefNull attribute_name : attribute_names) {
+    if (ELEM(attribute_name, "position", "normal", "id")) {
+      continue;
+    }
+
+    ReadAttributePtr attribute_in = mesh_component.attribute_try_get_for_read(attribute_name);
+    interpolate_attribute(
+        mesh, bary_coords, looptri_indices, attribute_name, *attribute_in, component);
+  }
+}
+
+BLI_NOINLINE static void compute_special_attributes(const Mesh &mesh,
+                                                    GeometryComponent &component,
+                                                    Span<float3> bary_coords,
+                                                    Span<int> looptri_indices)
 {
   OutputAttributePtr id_attribute = component.attribute_try_get_for_output(
       "id", ATTR_DOMAIN_POINT, CD_PROP_INT32);
@@ -255,18 +349,43 @@ BLI_NOINLINE static void add_remaining_point_attributes(const Mesh &mesh,
   OutputAttributePtr rotation_attribute = component.attribute_try_get_for_output(
       "rotation", ATTR_DOMAIN_POINT, CD_PROP_FLOAT3);
 
-  compute_special_attributes(mesh,
-                             bary_coords,
-                             looptri_indices,
-                             normal_attribute->get_span_for_write_only<float3>(),
-                             id_attribute->get_span_for_write_only<int>(),
-                             rotation_attribute->get_span_for_write_only<float3>());
+  MutableSpan<int> ids = id_attribute->get_span_for_write_only<int>();
+  MutableSpan<float3> normals = normal_attribute->get_span_for_write_only<float3>();
+  MutableSpan<float3> rotations = rotation_attribute->get_span_for_write_only<float3>();
+
+  Span<MLoopTri> looptris = get_mesh_looptris(mesh);
+  for (const int i : bary_coords.index_range()) {
+    const int looptri_index = looptri_indices[i];
+    const MLoopTri &looptri = looptris[looptri_index];
+    const float3 &bary_coord = bary_coords[i];
+
+    const int v0_index = mesh.mloop[looptri.tri[0]].v;
+    const int v1_index = mesh.mloop[looptri.tri[1]].v;
+    const int v2_index = mesh.mloop[looptri.tri[2]].v;
+    const float3 v0_pos = mesh.mvert[v0_index].co;
+    const float3 v1_pos = mesh.mvert[v1_index].co;
+    const float3 v2_pos = mesh.mvert[v2_index].co;
+
+    ids[i] = (int)(bary_coord.hash()) + looptri_index;
+    normal_tri_v3(normals[i], v0_pos, v1_pos, v2_pos);
+    rotations[i] = normal_to_euler_rotation(normals[i]);
+  }
 
   id_attribute.apply_span_and_save();
   normal_attribute.apply_span_and_save();
   rotation_attribute.apply_span_and_save();
 }
 
+BLI_NOINLINE static void add_remaining_point_attributes(const MeshComponent &mesh_component,
+                                                        GeometryComponent &component,
+                                                        Span<float3> bary_coords,
+                                                        Span<int> looptri_indices)
+{
+  interpolate_existing_attributes(mesh_component, component, bary_coords, looptri_indices);
+  compute_special_attributes(
+      *mesh_component.get_for_read(), component, bary_coords, looptri_indices);
+}
+
 static void sample_mesh_surface_with_minimum_distance(const Mesh &mesh,
                                                       const float max_density,
                                                       const float minimum_distance,
@@ -351,7 +470,7 @@ static void geo_node_point_distribute_exec(GeoNodeExecParams params)
       geometry_set_out.get_component_for_write<PointCloudComponent>();
   point_component.replace(pointcloud);
 
-  add_remaining_point_attributes(*mesh_in, point_component, bary_coords, looptri_indices);
+  add_remaining_point_attributes(mesh_component, point_component, bary_coords, looptri_indices);
 
   params.set_output("Geometry", std::move(geometry_set_out));
 }