Geometry Nodes: new Sample UV Surface node

This node allows sampling an attribute on a mesh surface based
on a UV coordinate. Internally, this has to do a "reverse uv lookup",
i.e. the node has to find the polygon that corresponds to the uv
coordinate. Therefore, the uv map of the mesh should not have
overlapping faces.

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

8 files changed, 312 insertions, 0 deletions

diff --git a/release/scripts/startup/bl_ui/node_add_menu_geometry.py b/release/scripts/startup/bl_ui/node_add_menu_geometry.py
index 3544c9a4925..2859b99195c 100644
--- a/release/scripts/startup/bl_ui/node_add_menu_geometry.py
+++ b/release/scripts/startup/bl_ui/node_add_menu_geometry.py
@@ -198,6 +198,7 @@ class NODE_MT_geometry_node_GEO_MESH(Menu):
         node_add_menu.add_node_type(layout, "GeometryNodeMeshToPoints")
         node_add_menu.add_node_type(layout, "GeometryNodeMeshToVolume")
         node_add_menu.add_node_type(layout, "GeometryNodeSampleNearestSurface")
+        node_add_menu.add_node_type(layout, "GeometryNodeSampleUVSurface")
         node_add_menu.add_node_type(layout, "GeometryNodeScaleElements")
         node_add_menu.add_node_type(layout, "GeometryNodeSplitEdges")
         node_add_menu.add_node_type(layout, "GeometryNodeSubdivideMesh")
diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index f79fefcae24..658b6daa88c 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1541,6 +1541,7 @@ struct TexResult;
 #define GEO_NODE_MESH_TOPOLOGY_EDGES_OF_VERTEX 1184
 #define GEO_NODE_MESH_TOPOLOGY_FACE_OF_CORNER 1185
 #define GEO_NODE_MESH_TOPOLOGY_VERTEX_OF_CORNER 1186
+#define GEO_NODE_SAMPLE_UV_SURFACE 1187
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 6413ad4c661..d7f4ec740b3 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -4813,6 +4813,7 @@ static void registerGeometryNodes()
   register_node_type_geo_sample_index();
   register_node_type_geo_sample_nearest_surface();
   register_node_type_geo_sample_nearest();
+  register_node_type_geo_sample_uv_surface();
   register_node_type_geo_scale_elements();
   register_node_type_geo_scale_instances();
   register_node_type_geo_separate_components();
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index 4ec2799186e..ed5a3563fb5 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -10462,6 +10462,18 @@ static void def_geo_sample_nearest(StructRNA *srna)
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_update");
 }
 
+static void def_geo_sample_uv_surface(StructRNA *srna)
+{
+  PropertyRNA *prop = RNA_def_property(srna, "data_type", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_sdna(prop, NULL, "custom1");
+  RNA_def_property_enum_items(prop, rna_enum_attribute_type_items);
+  RNA_def_property_enum_funcs(
+      prop, NULL, NULL, "rna_GeometryNodeAttributeType_type_with_socket_itemf");
+  RNA_def_property_enum_default(prop, CD_PROP_FLOAT);
+  RNA_def_property_ui_text(prop, "Data Type", "");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
+}
+
 static void def_geo_input_material(StructRNA *srna)
 {
   PropertyRNA *prop;
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index 7916bb6d08c..0c0b33770c8 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -131,6 +131,7 @@ void register_node_type_geo_rotate_instances(void);
 void register_node_type_geo_sample_index(void);
 void register_node_type_geo_sample_nearest_surface(void);
 void register_node_type_geo_sample_nearest(void);
+void register_node_type_geo_sample_uv_surface(void);
 void register_node_type_geo_scale_elements(void);
 void register_node_type_geo_scale_instances(void);
 void register_node_type_geo_select_by_handle_type(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index c4e1d54baef..3171a2db895 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -391,6 +391,7 @@ DefNode(GeometryNode, GEO_NODE_SAMPLE_CURVE, def_geo_curve_sample, "SAMPLE_CURVE
 DefNode(GeometryNode, GEO_NODE_SAMPLE_INDEX, def_geo_sample_index, "SAMPLE_INDEX", SampleIndex, "Sample Index", "Retrieve values from specific geometry elements")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_NEAREST_SURFACE, def_geo_sample_nearest_surface, "SAMPLE_NEAREST_SURFACE", SampleNearestSurface, "Sample Nearest Surface", "Calculate the interpolated value of a mesh attribute on the closest point of its surface")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_NEAREST, def_geo_sample_nearest, "SAMPLE_NEAREST", SampleNearest, "Sample Nearest", "Find the element of a geometry closest to a position")
+DefNode(GeometryNode, GEO_NODE_SAMPLE_UV_SURFACE, def_geo_sample_uv_surface, "SAMPLE_UV_SURFACE", SampleUVSurface, "Sample UV Surface", "Calculate the interpolated values of a mesh attribute at a UV coordinate")
 DefNode(GeometryNode, GEO_NODE_SCALE_ELEMENTS, def_geo_scale_elements, "SCALE_ELEMENTS", ScaleElements, "Scale Elements", "Scale groups of connected edges and faces")
 DefNode(GeometryNode, GEO_NODE_SCALE_INSTANCES, 0, "SCALE_INSTANCES", ScaleInstances, "Scale Instances", "Scale geometry instances in local or global space")
 DefNode(GeometryNode, GEO_NODE_SEPARATE_COMPONENTS, 0, "SEPARATE_COMPONENTS",SeparateComponents, "Separate Components","Split a geometry into a separate output for each type of data in the geometry")
diff --git a/source/blender/nodes/geometry/CMakeLists.txt b/source/blender/nodes/geometry/CMakeLists.txt
index 40b2ae61915..a81bbec071b 100644
--- a/source/blender/nodes/geometry/CMakeLists.txt
+++ b/source/blender/nodes/geometry/CMakeLists.txt
@@ -141,6 +141,7 @@ set(SRC
   nodes/node_geo_sample_index.cc
   nodes/node_geo_sample_nearest.cc
   nodes/node_geo_sample_nearest_surface.cc
+  nodes/node_geo_sample_uv_surface.cc
   nodes/node_geo_scale_elements.cc
   nodes/node_geo_scale_instances.cc
   nodes/node_geo_self_object.cc
diff --git a/source/blender/nodes/geometry/nodes/node_geo_sample_uv_surface.cc b/source/blender/nodes/geometry/nodes/node_geo_sample_uv_surface.cc
new file mode 100644
index 00000000000..a8ea0871449
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_sample_uv_surface.cc
@@ -0,0 +1,294 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+
+#include "BKE_attribute_math.hh"
+#include "BKE_mesh.h"
+#include "BKE_type_conversions.hh"
+
+#include "UI_interface.h"
+#include "UI_resources.h"
+
+#include "GEO_reverse_uv_sampler.hh"
+
+#include "NOD_socket_search_link.hh"
+
+#include "node_geometry_util.hh"
+
+namespace blender::nodes::node_geo_sample_uv_surface_cc {
+
+using geometry::ReverseUVSampler;
+
+static void node_declare(NodeDeclarationBuilder &b)
+{
+  b.add_input<decl::Geometry>(N_("Mesh")).supported_type({GEO_COMPONENT_TYPE_MESH});
+
+  b.add_input<decl::Float>(N_("Value"), "Value_Float").hide_value().supports_field();
+  b.add_input<decl::Int>(N_("Value"), "Value_Int").hide_value().supports_field();
+  b.add_input<decl::Vector>(N_("Value"), "Value_Vector").hide_value().supports_field();
+  b.add_input<decl::Color>(N_("Value"), "Value_Color").hide_value().supports_field();
+  b.add_input<decl::Bool>(N_("Value"), "Value_Bool").hide_value().supports_field();
+
+  b.add_input<decl::Vector>(N_("Source UV Map"))
+      .hide_value()
+      .supports_field()
+      .description(N_("The mesh UV map to sample. Should not have overlapping faces"));
+  b.add_input<decl::Vector>(N_("Sample UV"))
+      .supports_field()
+      .description(N_("The coordinates to sample within the UV map"));
+
+  b.add_output<decl::Float>(N_("Value"), "Value_Float").dependent_field({7});
+  b.add_output<decl::Int>(N_("Value"), "Value_Int").dependent_field({7});
+  b.add_output<decl::Vector>(N_("Value"), "Value_Vector").dependent_field({7});
+  b.add_output<decl::Color>(N_("Value"), "Value_Color").dependent_field({7});
+  b.add_output<decl::Bool>(N_("Value"), "Value_Bool").dependent_field({7});
+
+  b.add_output<decl::Bool>(N_("Is Valid"))
+      .dependent_field({7})
+      .description(N_("Whether the node could find a single face to sample at the UV coordinate"));
+}
+
+static void node_layout(uiLayout *layout, bContext *UNUSED(C), PointerRNA *ptr)
+{
+  uiItemR(layout, ptr, "data_type", 0, "", ICON_NONE);
+}
+
+static void node_init(bNodeTree *UNUSED(tree), bNode *node)
+{
+  node->custom1 = CD_PROP_FLOAT;
+}
+
+static void node_update(bNodeTree *ntree, bNode *node)
+{
+  const eCustomDataType data_type = eCustomDataType(node->custom1);
+
+  bNodeSocket *in_socket_mesh = static_cast<bNodeSocket *>(node->inputs.first);
+  bNodeSocket *in_socket_float = in_socket_mesh->next;
+  bNodeSocket *in_socket_int32 = in_socket_float->next;
+  bNodeSocket *in_socket_vector = in_socket_int32->next;
+  bNodeSocket *in_socket_color4f = in_socket_vector->next;
+  bNodeSocket *in_socket_bool = in_socket_color4f->next;
+
+  nodeSetSocketAvailability(ntree, in_socket_vector, data_type == CD_PROP_FLOAT3);
+  nodeSetSocketAvailability(ntree, in_socket_float, data_type == CD_PROP_FLOAT);
+  nodeSetSocketAvailability(ntree, in_socket_color4f, data_type == CD_PROP_COLOR);
+  nodeSetSocketAvailability(ntree, in_socket_bool, data_type == CD_PROP_BOOL);
+  nodeSetSocketAvailability(ntree, in_socket_int32, data_type == CD_PROP_INT32);
+
+  bNodeSocket *out_socket_float = static_cast<bNodeSocket *>(node->outputs.first);
+  bNodeSocket *out_socket_int32 = out_socket_float->next;
+  bNodeSocket *out_socket_vector = out_socket_int32->next;
+  bNodeSocket *out_socket_color4f = out_socket_vector->next;
+  bNodeSocket *out_socket_bool = out_socket_color4f->next;
+
+  nodeSetSocketAvailability(ntree, out_socket_vector, data_type == CD_PROP_FLOAT3);
+  nodeSetSocketAvailability(ntree, out_socket_float, data_type == CD_PROP_FLOAT);
+  nodeSetSocketAvailability(ntree, out_socket_color4f, data_type == CD_PROP_COLOR);
+  nodeSetSocketAvailability(ntree, out_socket_bool, data_type == CD_PROP_BOOL);
+  nodeSetSocketAvailability(ntree, out_socket_int32, data_type == CD_PROP_INT32);
+}
+
+static void node_gather_link_searches(GatherLinkSearchOpParams &params)
+{
+  const NodeDeclaration &declaration = *params.node_type().fixed_declaration;
+  search_link_ops_for_declarations(params, declaration.inputs().take_back(2));
+  search_link_ops_for_declarations(params, declaration.inputs().take_front(1));
+  search_link_ops_for_declarations(params, declaration.outputs().take_back(1));
+
+  const std::optional<eCustomDataType> type = node_data_type_to_custom_data_type(
+      eNodeSocketDatatype(params.other_socket().type));
+  if (type && *type != CD_PROP_STRING) {
+    /* The input and output sockets have the same name. */
+    params.add_item(IFACE_("Value"), [type](LinkSearchOpParams &params) {
+      bNode &node = params.add_node("GeometryNodeSampleUVSurface");
+      node.custom1 = *type;
+      params.update_and_connect_available_socket(node, "Value");
+    });
+  }
+}
+
+class SampleUVSurfaceFunction : public fn::MultiFunction {
+  GeometrySet source_;
+  Field<float2> src_uv_map_field_;
+  GField src_field_;
+
+  /**
+   * Use the most complex domain for now ensuring no information is lost. In the future, it should
+   * be possible to use the most complex domain required by the field inputs, to simplify sampling
+   * and avoid domain conversions.
+   */
+  eAttrDomain domain_ = ATTR_DOMAIN_CORNER;
+
+  fn::MFSignature signature_;
+
+  std::optional<bke::MeshFieldContext> source_context_;
+  std::unique_ptr<FieldEvaluator> source_evaluator_;
+  const GVArray *source_data_;
+  VArraySpan<float2> source_uv_map_;
+
+  std::optional<ReverseUVSampler> reverse_uv_sampler_;
+
+ public:
+  SampleUVSurfaceFunction(GeometrySet geometry, Field<float2> src_uv_map_field, GField src_field)
+      : source_(std::move(geometry)),
+        src_uv_map_field_(std::move(src_uv_map_field)),
+        src_field_(std::move(src_field))
+  {
+    source_.ensure_owns_direct_data();
+    signature_ = this->create_signature();
+    this->set_signature(&signature_);
+    this->evaluate_source();
+  }
+
+  fn::MFSignature create_signature()
+  {
+    blender::fn::MFSignatureBuilder signature{"Sample UV Surface"};
+    signature.single_input<float2>("Sample UV");
+    signature.single_output("Value", src_field_.cpp_type());
+    signature.single_output<bool>("Is Valid");
+    return signature.build();
+  }
+
+  void call(IndexMask mask, fn::MFParams params, fn::MFContext UNUSED(context)) const override
+  {
+    const VArray<float2> &sample_uvs = params.readonly_single_input<float2>(0, "Sample UV");
+    GMutableSpan dst = params.uninitialized_single_output_if_required(1, "Value");
+    MutableSpan<bool> valid_dst = params.uninitialized_single_output_if_required<bool>(2,
+                                                                                       "Is Valid");
+
+    const CPPType &type = src_field_.cpp_type();
+    attribute_math::convert_to_static_type(type, [&](auto dummy) {
+      using T = decltype(dummy);
+      const VArray<T> src_typed = source_data_->typed<T>();
+      MutableSpan<T> dst_typed = dst.typed<T>();
+      for (const int i : mask) {
+        const float2 sample_uv = sample_uvs[i];
+        const ReverseUVSampler::Result result = reverse_uv_sampler_->sample(sample_uv);
+        const bool valid = result.type == ReverseUVSampler::ResultType::Ok;
+        if (!dst_typed.is_empty()) {
+          if (valid) {
+            dst_typed[i] = attribute_math::mix3(result.bary_weights,
+                                                src_typed[result.looptri->tri[0]],
+                                                src_typed[result.looptri->tri[1]],
+                                                src_typed[result.looptri->tri[2]]);
+          }
+          else {
+            dst_typed[i] = {};
+          }
+        }
+        if (!valid_dst.is_empty()) {
+          valid_dst[i] = valid;
+        }
+      }
+    });
+  }
+
+ private:
+  void evaluate_source()
+  {
+    const Mesh &mesh = *source_.get_mesh_for_read();
+    source_context_.emplace(bke::MeshFieldContext{mesh, domain_});
+    const int domain_size = mesh.attributes().domain_size(domain_);
+    source_evaluator_ = std::make_unique<FieldEvaluator>(*source_context_, domain_size);
+    source_evaluator_->add(src_uv_map_field_);
+    source_evaluator_->add(src_field_);
+    source_evaluator_->evaluate();
+    source_uv_map_ = source_evaluator_->get_evaluated<float2>(0);
+    source_data_ = &source_evaluator_->get_evaluated(1);
+
+    reverse_uv_sampler_.emplace(source_uv_map_, mesh.looptris());
+  }
+};
+
+static GField get_input_attribute_field(GeoNodeExecParams &params, const eCustomDataType data_type)
+{
+  switch (data_type) {
+    case CD_PROP_FLOAT:
+      return params.extract_input<Field<float>>("Value_Float");
+    case CD_PROP_FLOAT3:
+      return params.extract_input<Field<float3>>("Value_Vector");
+    case CD_PROP_COLOR:
+      return params.extract_input<Field<ColorGeometry4f>>("Value_Color");
+    case CD_PROP_BOOL:
+      return params.extract_input<Field<bool>>("Value_Bool");
+    case CD_PROP_INT32:
+      return params.extract_input<Field<int>>("Value_Int");
+    default:
+      BLI_assert_unreachable();
+  }
+  return {};
+}
+
+static void output_attribute_field(GeoNodeExecParams &params, GField field)
+{
+  switch (bke::cpp_type_to_custom_data_type(field.cpp_type())) {
+    case CD_PROP_FLOAT: {
+      params.set_output("Value_Float", Field<float>(field));
+      break;
+    }
+    case CD_PROP_FLOAT3: {
+      params.set_output("Value_Vector", Field<float3>(field));
+      break;
+    }
+    case CD_PROP_COLOR: {
+      params.set_output("Value_Color", Field<ColorGeometry4f>(field));
+      break;
+    }
+    case CD_PROP_BOOL: {
+      params.set_output("Value_Bool", Field<bool>(field));
+      break;
+    }
+    case CD_PROP_INT32: {
+      params.set_output("Value_Int", Field<int>(field));
+      break;
+    }
+    default:
+      break;
+  }
+}
+
+static void node_geo_exec(GeoNodeExecParams params)
+{
+  GeometrySet geometry = params.extract_input<GeometrySet>("Mesh");
+  const eCustomDataType data_type = eCustomDataType(params.node().custom1);
+  const Mesh *mesh = geometry.get_mesh_for_read();
+  if (mesh == nullptr) {
+    params.set_default_remaining_outputs();
+    return;
+  }
+  if (mesh->totpoly == 0 && mesh->totvert != 0) {
+    params.error_message_add(NodeWarningType::Error, TIP_("The source mesh must have faces"));
+    params.set_default_remaining_outputs();
+    return;
+  }
+
+  const CPPType &float2_type = CPPType::get<float2>();
+
+  const bke::DataTypeConversions &conversions = bke::get_implicit_type_conversions();
+  Field<float2> source_uv_map = conversions.try_convert(
+      params.extract_input<Field<float3>>("Source UV Map"), float2_type);
+  GField field = get_input_attribute_field(params, data_type);
+  Field<float2> sample_uvs = conversions.try_convert(
+      params.extract_input<Field<float3>>("Sample UV"), float2_type);
+  auto fn = std::make_shared<SampleUVSurfaceFunction>(
+      std::move(geometry), std::move(source_uv_map), std::move(field));
+  auto op = FieldOperation::Create(std::move(fn), {std::move(sample_uvs)});
+  output_attribute_field(params, GField(op, 0));
+  params.set_output("Is Valid", Field<bool>(op, 1));
+}
+
+}  // namespace blender::nodes::node_geo_sample_uv_surface_cc
+
+void register_node_type_geo_sample_uv_surface()
+{
+  namespace file_ns = blender::nodes::node_geo_sample_uv_surface_cc;
+
+  static bNodeType ntype;
+
+  geo_node_type_base(&ntype, GEO_NODE_SAMPLE_UV_SURFACE, "Sample UV Surface", NODE_CLASS_GEOMETRY);
+  node_type_init(&ntype, file_ns::node_init);
+  node_type_update(&ntype, file_ns::node_update);
+  ntype.declare = file_ns::node_declare;
+  ntype.geometry_node_execute = file_ns::node_geo_exec;
+  ntype.draw_buttons = file_ns::node_layout;
+  ntype.gather_link_search_ops = file_ns::node_gather_link_searches;
+  nodeRegisterType(&ntype);
+}