initial sample mesh surface

This is a field based alternative for the attribute transfer node.
The main goal here is to experiment with how attribute transfer
should be done with fields.

10 files changed, 293 insertions, 72 deletions

diff --git a/source/blender/blenkernel/BKE_field.hh b/source/blender/blenkernel/BKE_field.hh
index 0c675ec9a89..801060cc070 100644
--- a/source/blender/blenkernel/BKE_field.hh
+++ b/source/blender/blenkernel/BKE_field.hh
@@ -350,14 +350,16 @@ template<typename KeyT> class GVArrayInputField : public Field {
 class MultiFunctionField : public Field {
  private:
   Vector<FieldPtr> input_fields_;
-  const MultiFunction *fn_;
+  optional_ptr<const MultiFunction> fn_;
   const int output_param_index_;
 
  public:
   MultiFunctionField(Vector<FieldPtr> input_fields,
-                     const MultiFunction &fn,
+                     optional_ptr<const MultiFunction> fn,
                      const int output_param_index)
-      : input_fields_(std::move(input_fields)), fn_(&fn), output_param_index_(output_param_index)
+      : input_fields_(std::move(input_fields)),
+        fn_(std::move(fn)),
+        output_param_index_(output_param_index)
   {
   }
 
diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index 346b5b578d4..7b3321bf70f 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1487,6 +1487,7 @@ int ntreeTexExecTree(struct bNodeTree *ntree,
 #define GEO_NODE_EXTRUDE_AND_MOVE 1082
 #define GEO_NODE_POSITION 1083
 #define GEO_NODE_SET_POSITION 1084
+#define GEO_NODE_SAMPLE_MESH_SURFACE 1085
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index ae3cff7c79c..33b227ae73d 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -5189,6 +5189,7 @@ static void registerGeometryNodes()
   register_node_type_geo_points_to_volume();
   register_node_type_geo_position();
   register_node_type_geo_raycast();
+  register_node_type_geo_sample_mesh_surface();
   register_node_type_geo_sample_texture();
   register_node_type_geo_select_by_handle_type();
   register_node_type_geo_select_by_material();
diff --git a/source/blender/nodes/CMakeLists.txt b/source/blender/nodes/CMakeLists.txt
index 5adcd7e006c..9ac8dff68ab 100644
--- a/source/blender/nodes/CMakeLists.txt
+++ b/source/blender/nodes/CMakeLists.txt
@@ -217,6 +217,7 @@ set(SRC
   geometry/nodes/node_geo_points_to_volume.cc
   geometry/nodes/node_geo_position.cc
   geometry/nodes/node_geo_raycast.cc
+  geometry/nodes/node_geo_sample_mesh_surface.cc
   geometry/nodes/node_geo_select_by_material.cc
   geometry/nodes/node_geo_separate_components.cc
   geometry/nodes/node_geo_set_position.cc
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index 82fe217df40..b6390825625 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -102,6 +102,7 @@ void register_node_type_geo_point_translate(void);
 void register_node_type_geo_points_to_volume(void);
 void register_node_type_geo_position(void);
 void register_node_type_geo_raycast(void);
+void register_node_type_geo_sample_mesh_surface(void);
 void register_node_type_geo_sample_texture(void);
 void register_node_type_geo_select_by_handle_type(void);
 void register_node_type_geo_select_by_material(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index 94230fc74d6..b29cb1389e0 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -345,6 +345,7 @@ DefNode(GeometryNode, GEO_NODE_POINT_TRANSLATE, def_geo_point_translate, "POINT_
 DefNode(GeometryNode, GEO_NODE_POINTS_TO_VOLUME, def_geo_points_to_volume, "POINTS_TO_VOLUME", PointsToVolume, "Points to Volume", "")
 DefNode(GeometryNode, GEO_NODE_POSITION, 0, "POSITION", Position, "Position", "")
 DefNode(GeometryNode, GEO_NODE_RAYCAST, def_geo_raycast, "RAYCAST", Raycast, "Raycast", "")
+DefNode(GeometryNode, GEO_NODE_SAMPLE_MESH_SURFACE, 0, "SAMPLE_MESH_SURFACE", SampleMeshSurface, "Sample Mesh Surface", "")
 DefNode(GeometryNode, GEO_NODE_SELECT_BY_MATERIAL, 0, "SELECT_BY_MATERIAL", SelectByMaterial, "Select by Material", "")
 DefNode(GeometryNode, GEO_NODE_SET_POSITION, 0, "SET_POSITION", SetPosition, "Set Position", "")
 DefNode(GeometryNode, GEO_NODE_SEPARATE_COMPONENTS, 0, "SEPARATE_COMPONENTS", SeparateComponents, "Separate Components", "")
diff --git a/source/blender/nodes/geometry/node_geometry_util.cc b/source/blender/nodes/geometry/node_geometry_util.cc
index 4f8842c1f17..f03b1e966e4 100644
--- a/source/blender/nodes/geometry/node_geometry_util.cc
+++ b/source/blender/nodes/geometry/node_geometry_util.cc
@@ -220,6 +220,64 @@ void prepare_field_inputs(bke::FieldInputs &field_inputs,
   }
 }
 
+template<typename T>
+void fill_attribute_impl(GeometryComponent &component,
+                         OutputAttribute &attribute,
+                         const bke::Field &field)
+{
+  const AttributeDomain domain = attribute.domain();
+  const int domain_size = attribute->size();
+  bke::FieldInputs field_inputs = field.prepare_inputs();
+  Vector<std::unique_ptr<bke::FieldInputValue>> input_values;
+  prepare_field_inputs(field_inputs, component, domain, input_values);
+  bke::FieldOutput field_output = field.evaluate(IndexMask(domain_size), field_inputs);
+  for (const int i : IndexRange(domain_size)) {
+    T value;
+    field_output.varray_ref().get(i, &value);
+    attribute->set_by_copy(i, &value);
+  }
+}
+
+void try_freeze_field_on_geometry(GeometryComponent &component,
+                                  const AnonymousCustomDataLayerID &layer_id,
+                                  AttributeDomain domain,
+                                  const bke::Field &field)
+{
+  const CustomDataType data_type = bke::cpp_type_to_custom_data_type(field.output_type());
+  OutputAttribute attribute = component.attribute_try_get_anonymous_for_output(
+      layer_id, domain, data_type);
+  if (!attribute) {
+    return;
+  }
+
+  switch (data_type) {
+    case CD_PROP_FLOAT: {
+      fill_attribute_impl<float>(component, attribute, field);
+      break;
+    }
+    case CD_PROP_FLOAT3: {
+      fill_attribute_impl<float3>(component, attribute, field);
+      break;
+    }
+    case CD_PROP_COLOR: {
+      fill_attribute_impl<ColorGeometry4f>(component, attribute, field);
+      break;
+    }
+    case CD_PROP_BOOL: {
+      fill_attribute_impl<bool>(component, attribute, field);
+      break;
+    }
+    case CD_PROP_INT32: {
+      fill_attribute_impl<int>(component, attribute, field);
+      break;
+    }
+    default:
+      break;
+  }
+
+  attribute.save();
+}
+
 }  // namespace blender::nodes
 
 bool geo_node_poll_default(bNodeType *UNUSED(ntype),
diff --git a/source/blender/nodes/geometry/node_geometry_util.hh b/source/blender/nodes/geometry/node_geometry_util.hh
index 2aa7276578a..7ffacee88c4 100644
--- a/source/blender/nodes/geometry/node_geometry_util.hh
+++ b/source/blender/nodes/geometry/node_geometry_util.hh
@@ -107,4 +107,9 @@ void prepare_field_inputs(bke::FieldInputs &field_inputs,
                           const AttributeDomain domain,
                           Vector<std::unique_ptr<bke::FieldInputValue>> &r_values);
 
+void try_freeze_field_on_geometry(GeometryComponent &component,
+                                  const AnonymousCustomDataLayerID &layer_id,
+                                  AttributeDomain domain,
+                                  const bke::Field &field);
+
 }  // namespace blender::nodes
diff --git a/source/blender/nodes/geometry/nodes/node_geo_attribute_freeze.cc b/source/blender/nodes/geometry/nodes/node_geo_attribute_freeze.cc
index 4f4468d8860..8404f884373 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_attribute_freeze.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_freeze.cc
@@ -17,6 +17,8 @@
 #include "UI_interface.h"
 #include "UI_resources.h"
 
+#include "BKE_attribute_math.hh"
+
 #include "node_geometry_util.hh"
 
 static bNodeSocketTemplate geo_node_attribute_freeze_in[] = {
@@ -133,74 +135,6 @@ static void set_output_field(GeoNodeExecParams &params,
   }
 }
 
-template<typename T>
-void fill_attribute_impl(GeometryComponent &component,
-                         OutputAttribute &attribute,
-                         const GeoNodeExecParams &params,
-                         const StringRef input_name)
-{
-  const AttributeDomain domain = attribute.domain();
-  const int domain_size = attribute->size();
-  bke::FieldRef<T> value_field = params.get_input_field<T>(input_name);
-  bke::FieldInputs field_inputs = value_field->prepare_inputs();
-  Vector<std::unique_ptr<bke::FieldInputValue>> input_values;
-  prepare_field_inputs(field_inputs, component, domain, input_values);
-  bke::FieldOutput field_output = value_field->evaluate(IndexMask(domain_size), field_inputs);
-  for (const int i : IndexRange(domain_size)) {
-    T value;
-    field_output.varray_ref().get(i, &value);
-    attribute->set_by_copy(i, &value);
-  }
-}
-
-static void fill_attribute_data(GeometryComponent &component,
-                                const GeoNodeExecParams &params,
-                                OutputAttribute &attribute,
-                                const CustomDataType data_type)
-{
-  switch (data_type) {
-    case CD_PROP_FLOAT: {
-      fill_attribute_impl<float>(component, attribute, params, "Value_001");
-      break;
-    }
-    case CD_PROP_FLOAT3: {
-      fill_attribute_impl<float3>(component, attribute, params, "Value");
-      break;
-    }
-    case CD_PROP_COLOR: {
-      fill_attribute_impl<ColorGeometry4f>(component, attribute, params, "Value_002");
-      break;
-    }
-    case CD_PROP_BOOL: {
-      fill_attribute_impl<bool>(component, attribute, params, "Value_003");
-      break;
-    }
-    case CD_PROP_INT32: {
-      fill_attribute_impl<int>(component, attribute, params, "Value_004");
-      break;
-    }
-    default:
-      break;
-  }
-}
-
-static void fill_anonymous(GeometryComponent &component,
-                           const GeoNodeExecParams &params,
-                           const AnonymousCustomDataLayerID &layer_id,
-                           const AttributeDomain domain,
-                           const CustomDataType data_type)
-{
-  OutputAttribute attribute = component.attribute_try_get_anonymous_for_output(
-      layer_id, domain, data_type);
-  if (!attribute) {
-    return;
-  }
-
-  fill_attribute_data(component, params, attribute, data_type);
-
-  attribute.save();
-}
-
 static void geo_node_attribute_freeze_exec(GeoNodeExecParams params)
 {
   GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
@@ -214,11 +148,33 @@ static void geo_node_attribute_freeze_exec(GeoNodeExecParams params)
 
   AnonymousCustomDataLayerID *id = CustomData_anonymous_id_new("Attribute Freeze");
 
+  FieldPtr field;
+  switch (data_type) {
+    case CD_PROP_FLOAT:
+      field = params.get_input_field<float>("Value_001").field();
+      break;
+    case CD_PROP_FLOAT3:
+      field = params.get_input_field<float3>("Value").field();
+      break;
+    case CD_PROP_COLOR:
+      field = params.get_input_field<ColorGeometry4f>("Value_002").field();
+      break;
+    case CD_PROP_BOOL:
+      field = params.get_input_field<bool>("Value_003").field();
+      break;
+    case CD_PROP_INT32:
+      field = params.get_input_field<int>("Value_004").field();
+      break;
+    default:
+      break;
+  }
+
   static const Array<GeometryComponentType> types = {
       GEO_COMPONENT_TYPE_MESH, GEO_COMPONENT_TYPE_POINT_CLOUD, GEO_COMPONENT_TYPE_CURVE};
   for (const GeometryComponentType type : types) {
     if (geometry_set.has(type)) {
-      fill_anonymous(geometry_set.get_component_for_write(type), params, *id, domain, data_type);
+      GeometryComponent &component = geometry_set.get_component_for_write(type);
+      try_freeze_field_on_geometry(component, *id, domain, *field);
     }
   }
 
diff --git a/source/blender/nodes/geometry/nodes/node_geo_sample_mesh_surface.cc b/source/blender/nodes/geometry/nodes/node_geo_sample_mesh_surface.cc
new file mode 100644
index 00000000000..73d6871251c
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_sample_mesh_surface.cc
@@ -0,0 +1,195 @@
+/*
+ * 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 "UI_interface.h"
+#include "UI_resources.h"
+
+#include "BLI_kdopbvh.h"
+
+#include "DNA_mesh_types.h"
+#include "DNA_meshdata_types.h"
+
+#include "BKE_attribute_math.hh"
+#include "BKE_bvhutils.h"
+#include "BKE_customdata.h"
+#include "BKE_mesh_runtime.h"
+#include "BKE_mesh_sample.hh"
+
+#include "node_geometry_util.hh"
+
+static bNodeSocketTemplate geo_node_sample_mesh_surface_in[] = {
+    {SOCK_GEOMETRY, N_("Mesh")},
+    {SOCK_VECTOR,
+     N_("Position"),
+     0.0f,
+     0.0f,
+     0.0f,
+     1.0f,
+     -FLT_MAX,
+     FLT_MAX,
+     PROP_TRANSLATION,
+     SOCK_HIDE_VALUE | SOCK_FIELD},
+    {SOCK_RGBA, N_("Custom"), 1, 1, 1, 1, 0, 1, PROP_NONE, SOCK_FIELD},
+    {-1, ""},
+};
+
+static bNodeSocketTemplate geo_node_sample_mesh_surface_out[] = {
+    {SOCK_VECTOR, N_("Position")},
+    {SOCK_VECTOR, N_("Normal")},
+    {SOCK_FLOAT, N_("Distance")},
+    {SOCK_RGBA, N_("Custom")},
+    {-1, ""},
+};
+
+namespace blender::nodes {
+
+class SampleMeshSurfaceFunction : public fn::MultiFunction {
+ private:
+  GeometrySet geometry_set_;
+  AnonymousCustomDataLayerID *attribute_id_;
+
+ public:
+  SampleMeshSurfaceFunction(GeometrySet geometry_set, AnonymousCustomDataLayerID *attribute_id)
+      : geometry_set_(std::move(geometry_set)), attribute_id_(attribute_id)
+  {
+    static fn::MFSignature signature = create_signature();
+    this->set_signature(&signature);
+    CustomData_anonymous_id_strong_increment(attribute_id_);
+  }
+
+  ~SampleMeshSurfaceFunction() override
+  {
+    CustomData_anonymous_id_strong_decrement(attribute_id_);
+  }
+
+  static blender::fn::MFSignature create_signature()
+  {
+    blender::fn::MFSignatureBuilder signature{"Sample Mesh Surface"};
+    signature.single_input<float3>("Position");
+    signature.single_output<float3>("Position");
+    signature.single_output<float3>("Normal");
+    signature.single_output<float>("Distance");
+    signature.single_output<ColorGeometry4f>("Attribute");
+    return signature.build();
+  }
+
+  void call(IndexMask mask, fn::MFParams params, fn::MFContext UNUSED(context)) const override
+  {
+    const VArray<float3> &src_positions = params.readonly_single_input<float3>(0, "Position");
+    MutableSpan<float3> sampled_positions = params.uninitialized_single_output<float3>(1,
+                                                                                       "Position");
+    MutableSpan<float3> sampled_normals = params.uninitialized_single_output<float3>(2, "Normal");
+    MutableSpan<float> sampled_distances = params.uninitialized_single_output<float>(3,
+                                                                                     "Distance");
+    MutableSpan<ColorGeometry4f> sampled_attribute =
+        params.uninitialized_single_output<ColorGeometry4f>(4, "Attribute");
+
+    if (!geometry_set_.has_mesh()) {
+      return this->return_default(mask, params);
+    }
+
+    const MeshComponent *mesh_component = geometry_set_.get_component_for_read<MeshComponent>();
+    const Mesh *mesh = mesh_component->get_for_read();
+
+    GVArrayPtr attribute_ptr = mesh_component->attribute_try_get_anonymous_for_read(
+        *attribute_id_, ATTR_DOMAIN_CORNER, CD_PROP_COLOR, nullptr);
+    if (!attribute_ptr) {
+      return this->return_default(mask, params);
+    }
+    GVArray_Typed<ColorGeometry4f> attribute{*attribute_ptr};
+
+    const MLoopTri *looptris = BKE_mesh_runtime_looptri_ensure(mesh);
+
+    BVHTreeFromMesh tree_data;
+    BKE_bvhtree_from_mesh_get(&tree_data, mesh, BVHTREE_FROM_LOOPTRI, 2);
+
+    for (const int i : mask) {
+      BVHTreeNearest nearest;
+      nearest.dist_sq = FLT_MAX;
+      const float3 src_position = src_positions[i];
+      BLI_bvhtree_find_nearest(
+          tree_data.tree, src_position, &nearest, tree_data.nearest_callback, &tree_data);
+      sampled_positions[i] = nearest.co;
+      sampled_normals[i] = nearest.no;
+      sampled_distances[i] = sqrtf(nearest.dist_sq);
+
+      const MLoopTri &looptri = looptris[nearest.index];
+
+      float3 v1 = mesh->mvert[mesh->mloop[looptri.tri[0]].v].co;
+      float3 v2 = mesh->mvert[mesh->mloop[looptri.tri[1]].v].co;
+      float3 v3 = mesh->mvert[mesh->mloop[looptri.tri[2]].v].co;
+
+      ColorGeometry4f col1 = attribute[looptri.tri[0]];
+      ColorGeometry4f col2 = attribute[looptri.tri[1]];
+      ColorGeometry4f col3 = attribute[looptri.tri[2]];
+
+      float3 bary_coords;
+      interp_weights_tri_v3(bary_coords, v1, v2, v3, nearest.co);
+      ColorGeometry4f final_col = attribute_math::mix3(bary_coords, col1, col2, col3);
+      sampled_attribute[i] = final_col;
+    }
+  }
+
+  void return_default(IndexMask mask, fn::MFParams &params) const
+  {
+    params.uninitialized_single_output<float3>(1, "Position").fill_indices(mask, float3{0, 0, 0});
+    params.uninitialized_single_output<float3>(2, "Normal").fill_indices(mask, float3{0, 0, 0});
+    params.uninitialized_single_output<float>(3, "Distance").fill_indices(mask, 0.0f);
+    params.uninitialized_single_output<ColorGeometry4f>(4, "Attribute")
+        .fill_indices(mask, ColorGeometry4f{0, 0, 0, 1});
+  }
+};
+
+static void geo_node_sample_mesh_surface_exec(GeoNodeExecParams params)
+{
+  GeometrySet geometry_set = params.extract_input<GeometrySet>("Mesh");
+  geometry_set = geometry_set_realize_instances(geometry_set);
+
+  FieldPtr position_field = params.get_input_field<float3>("Position").field();
+  bke::FieldRef<ColorGeometry4f> attribute_field = params.get_input_field<ColorGeometry4f>(
+      "Custom");
+
+  AnonymousCustomDataLayerID *layer_id = CustomData_anonymous_id_new("Sample Mesh Surface");
+  MeshComponent &mesh_component = geometry_set.get_component_for_write<MeshComponent>();
+  try_freeze_field_on_geometry(
+      mesh_component, *layer_id, ATTR_DOMAIN_CORNER, *attribute_field.field());
+
+  auto make_output_field = [&](int out_param_index) -> FieldPtr {
+    auto fn = std::make_unique<SampleMeshSurfaceFunction>(geometry_set, layer_id);
+    return new bke::MultiFunctionField(Vector<FieldPtr>{position_field},
+                                       optional_ptr<const fn::MultiFunction>{std::move(fn)},
+                                       out_param_index);
+  };
+
+  params.set_output("Position", bke::FieldRef<float3>(make_output_field(1)));
+  params.set_output("Normal", bke::FieldRef<float3>(make_output_field(2)));
+  params.set_output("Distance", bke::FieldRef<float>(make_output_field(3)));
+  params.set_output("Custom", bke::FieldRef<ColorGeometry4f>(make_output_field(4)));
+}
+
+}  // namespace blender::nodes
+
+void register_node_type_geo_sample_mesh_surface()
+{
+  static bNodeType ntype;
+
+  geo_node_type_base(
+      &ntype, GEO_NODE_SAMPLE_MESH_SURFACE, "Sample Mesh Surface", NODE_CLASS_ATTRIBUTE, 0);
+  node_type_socket_templates(
+      &ntype, geo_node_sample_mesh_surface_in, geo_node_sample_mesh_surface_out);
+  ntype.geometry_node_execute = blender::nodes::geo_node_sample_mesh_surface_exec;
+  nodeRegisterType(&ntype);
+}