Geometry Nodes: Fields version of attribute proximity node

Add a fields-aware implementation of the attribute proximity node.
The Source position is an implicit position field, but can be
connected with a position input node with alterations before use.
The target input and mode function the same as the original node.

Patch by Johnny Matthews with edits from Hans Goudey (@HooglyBoogly).

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

10 files changed, 294 insertions, 6 deletions

diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index 8eff70af364..76a3e75d285 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1506,6 +1506,7 @@ int ntreeTexExecTree(struct bNodeTree *ntree,
 #define GEO_NODE_MESH_TO_POINTS 1093
 #define GEO_NODE_POINTS_TO_VERTICES 1094
 #define GEO_NODE_CURVE_REVERSE 1095
+#define GEO_NODE_PROXIMITY 1096
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 416dae4c41e..a3a82bee8dc 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -5707,6 +5707,7 @@ static void registerGeometryNodes()
 {
   register_node_type_geo_group();
 
+  register_node_type_geo_legacy_attribute_proximity();
   register_node_type_geo_legacy_attribute_randomize();
   register_node_type_geo_legacy_material_assign();
   register_node_type_geo_legacy_select_by_material();
@@ -5724,7 +5725,6 @@ static void registerGeometryNodes()
   register_node_type_geo_attribute_map_range();
   register_node_type_geo_attribute_math();
   register_node_type_geo_attribute_mix();
-  register_node_type_geo_attribute_proximity();
   register_node_type_geo_attribute_remove();
   register_node_type_geo_attribute_separate_xyz();
   register_node_type_geo_attribute_statistic();
@@ -5790,6 +5790,7 @@ static void registerGeometryNodes()
   register_node_type_geo_point_translate();
   register_node_type_geo_points_to_vertices();
   register_node_type_geo_points_to_volume();
+  register_node_type_geo_proximity();
   register_node_type_geo_raycast();
   register_node_type_geo_realize_instances();
   register_node_type_geo_sample_texture();
diff --git a/source/blender/makesdna/DNA_node_types.h b/source/blender/makesdna/DNA_node_types.h
index 538a80ce3e1..ffa6cb307b8 100644
--- a/source/blender/makesdna/DNA_node_types.h
+++ b/source/blender/makesdna/DNA_node_types.h
@@ -1362,6 +1362,11 @@ typedef struct NodeGeometryAttributeProximity {
   uint8_t target_geometry_element;
 } NodeGeometryAttributeProximity;
 
+typedef struct NodeGeometryProximity {
+  /* GeometryNodeProximityTargetType. */
+  uint8_t target_element;
+} NodeGeometryProximity;
+
 typedef struct NodeGeometryVolumeToMesh {
   /* VolumeToMeshResolutionMode */
   uint8_t resolution_mode;
@@ -1946,6 +1951,12 @@ typedef enum GeometryNodeAttributeProximityTargetType {
   GEO_NODE_PROXIMITY_TARGET_FACES = 2,
 } GeometryNodeAttributeProximityTargetType;
 
+typedef enum GeometryNodeProximityTargetType {
+  GEO_NODE_PROX_TARGET_POINTS = 0,
+  GEO_NODE_PROX_TARGET_EDGES = 1,
+  GEO_NODE_PROX_TARGET_FACES = 2,
+} GeometryNodeProximityTargetType;
+
 typedef enum GeometryNodeBooleanOperation {
   GEO_NODE_BOOLEAN_INTERSECT = 0,
   GEO_NODE_BOOLEAN_UNION = 1,
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index 556b8e94fe6..b5a42babc81 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -9972,7 +9972,7 @@ static void def_geo_collection_info(StructRNA *srna)
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_update");
 }
 
-static void def_geo_attribute_proximity(StructRNA *srna)
+static void def_geo_legacy_attribute_proximity(StructRNA *srna)
 {
   static const EnumPropertyItem target_geometry_element[] = {
       {GEO_NODE_PROXIMITY_TARGET_POINTS,
@@ -10005,6 +10005,39 @@ static void def_geo_attribute_proximity(StructRNA *srna)
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
 }
 
+static void def_geo_proximity(StructRNA *srna)
+{
+  static const EnumPropertyItem target_element_items[] = {
+      {GEO_NODE_PROX_TARGET_POINTS,
+       "POINTS",
+       ICON_NONE,
+       "Points",
+       "Calculate the proximity to the target's points (faster than the other modes)"},
+      {GEO_NODE_PROX_TARGET_EDGES,
+       "EDGES",
+       ICON_NONE,
+       "Edges",
+       "Calculate the proximity to the target's edges"},
+      {GEO_NODE_PROX_TARGET_FACES,
+       "FACES",
+       ICON_NONE,
+       "Faces",
+       "Calculate the proximity to the target's faces"},
+      {0, NULL, 0, NULL, NULL},
+  };
+
+  PropertyRNA *prop;
+
+  RNA_def_struct_sdna_from(srna, "NodeGeometryProximity", "storage");
+
+  prop = RNA_def_property(srna, "target_element", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_items(prop, target_element_items);
+  RNA_def_property_enum_default(prop, GEO_NODE_PROX_TARGET_FACES);
+  RNA_def_property_ui_text(
+      prop, "Target Geometry", "Element of the target geometry to calculate the distance from");
+  RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
+}
+
 static void def_geo_volume_to_mesh(StructRNA *srna)
 {
   PropertyRNA *prop;
diff --git a/source/blender/modifiers/intern/MOD_nodes_evaluator.cc b/source/blender/modifiers/intern/MOD_nodes_evaluator.cc
index b9a9437d761..592e6180f80 100644
--- a/source/blender/modifiers/intern/MOD_nodes_evaluator.cc
+++ b/source/blender/modifiers/intern/MOD_nodes_evaluator.cc
@@ -329,7 +329,11 @@ static void get_socket_value(const SocketRef &socket, void *r_value)
   if (bsocket.flag & SOCK_HIDE_VALUE) {
     const bNode &bnode = *socket.bnode();
     if (bsocket.type == SOCK_VECTOR) {
-      if (ELEM(bnode.type, GEO_NODE_SET_POSITION, SH_NODE_TEX_NOISE, GEO_NODE_MESH_TO_POINTS)) {
+      if (ELEM(bnode.type,
+               GEO_NODE_SET_POSITION,
+               SH_NODE_TEX_NOISE,
+               GEO_NODE_MESH_TO_POINTS,
+               GEO_NODE_PROXIMITY)) {
         new (r_value) Field<float3>(
             std::make_shared<bke::AttributeFieldInput>("position", CPPType::get<float3>()));
         return;
diff --git a/source/blender/nodes/CMakeLists.txt b/source/blender/nodes/CMakeLists.txt
index 49e76ffd1ad..416fcd8f9fd 100644
--- a/source/blender/nodes/CMakeLists.txt
+++ b/source/blender/nodes/CMakeLists.txt
@@ -235,6 +235,7 @@ set(SRC
   geometry/nodes/node_geo_mesh_to_points.cc
   geometry/nodes/node_geo_object_info.cc
   geometry/nodes/node_geo_points_to_vertices.cc
+  geometry/nodes/node_geo_proximity.cc
   geometry/nodes/node_geo_realize_instances.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 7415563803b..af8661d9b8d 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -29,6 +29,7 @@ void register_node_tree_type_geo(void);
 void register_node_type_geo_group(void);
 void register_node_type_geo_custom_group(bNodeType *ntype);
 
+void register_node_type_geo_legacy_attribute_proximity(void);
 void register_node_type_geo_legacy_attribute_randomize(void);
 void register_node_type_geo_legacy_material_assign(void);
 void register_node_type_geo_legacy_select_by_material(void);
@@ -46,7 +47,6 @@ void register_node_type_geo_attribute_fill(void);
 void register_node_type_geo_attribute_map_range(void);
 void register_node_type_geo_attribute_math(void);
 void register_node_type_geo_attribute_mix(void);
-void register_node_type_geo_attribute_proximity(void);
 void register_node_type_geo_attribute_remove(void);
 void register_node_type_geo_attribute_separate_xyz(void);
 void register_node_type_geo_attribute_statistic(void);
@@ -112,6 +112,7 @@ void register_node_type_geo_point_separate(void);
 void register_node_type_geo_point_translate(void);
 void register_node_type_geo_points_to_vertices(void);
 void register_node_type_geo_points_to_volume(void);
+void register_node_type_geo_proximity(void);
 void register_node_type_geo_raycast(void);
 void register_node_type_geo_realize_instances(void);
 void register_node_type_geo_sample_texture(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index ce270572c0f..3ee3faf6122 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -286,7 +286,7 @@ DefNode(GeometryNode, GEO_NODE_LEGACY_ATTRIBUTE_FILL, def_geo_attribute_fill, "L
 DefNode(GeometryNode, GEO_NODE_LEGACY_ATTRIBUTE_MAP_RANGE, def_geo_attribute_map_range, "LEGACY_ATTRIBUTE_MAP_RANGE", LegacyAttributeMapRange, "Attribute Map Range", "")
 DefNode(GeometryNode, GEO_NODE_LEGACY_ATTRIBUTE_MATH, def_geo_attribute_math, "LEGACY_ATTRIBUTE_MATH", LegacyAttributeMath, "Attribute Math", "")
 DefNode(GeometryNode, GEO_NODE_LEGACY_ATTRIBUTE_MIX, def_geo_attribute_mix, "LEGACY_ATTRIBUTE_MIX", LegacyAttributeMix, "Attribute Mix", "")
-DefNode(GeometryNode, GEO_NODE_LEGACY_ATTRIBUTE_PROXIMITY, def_geo_attribute_proximity, "LEGACY_ATTRIBUTE_PROXIMITY", LegacyAttributeProximity, "Attribute Proximity", "")
+DefNode(GeometryNode, GEO_NODE_LEGACY_ATTRIBUTE_PROXIMITY, def_geo_legacy_attribute_proximity, "LEGACY_ATTRIBUTE_PROXIMITY", LegacyAttributeProximity, "Attribute Proximity", "")
 DefNode(GeometryNode, GEO_NODE_LEGACY_ATTRIBUTE_RANDOMIZE, def_geo_attribute_randomize, "LEGACY_ATTRIBUTE_RANDOMIZE", LegacyAttributeRandomize, "Attribute Randomize", "")
 DefNode(GeometryNode, GEO_NODE_LEGACY_ATTRIBUTE_SAMPLE_TEXTURE, 0, "LEGACY_ATTRIBUTE_SAMPLE_TEXTURE", LegacyAttributeSampleTexture, "Attribute Sample Texture", "")
 DefNode(GeometryNode, GEO_NODE_LEGACY_ATTRIBUTE_SEPARATE_XYZ, def_geo_attribute_separate_xyz, "LEGACY_ATTRIBUTE_SEPARATE_XYZ", LegacyAttributeSeparateXYZ, "Attribute Separate XYZ", "")
@@ -362,6 +362,7 @@ DefNode(GeometryNode, GEO_NODE_MESH_SUBDIVIDE, 0, "MESH_SUBDIVIDE", MeshSubdivid
 DefNode(GeometryNode, GEO_NODE_MESH_TO_POINTS, def_geo_mesh_to_points, "MESH_TO_POINTS", MeshToPoints, "Mesh to Points", "")
 DefNode(GeometryNode, GEO_NODE_OBJECT_INFO, def_geo_object_info, "OBJECT_INFO", ObjectInfo, "Object Info", "")
 DefNode(GeometryNode, GEO_NODE_POINTS_TO_VERTICES, 0, "POINTS_TO_VERTICES", PointsToVertices, "Points to Vertices", "")
+DefNode(GeometryNode, GEO_NODE_PROXIMITY, def_geo_proximity, "PROXIMITY", Proximity, "Geometry Proximity", "")
 DefNode(GeometryNode, GEO_NODE_REALIZE_INSTANCES, 0, "REALIZE_INSTANCES", RealizeInstances, "Realize Instances", "")
 DefNode(GeometryNode, GEO_NODE_SEPARATE_COMPONENTS, 0, "SEPARATE_COMPONENTS", SeparateComponents, "Separate Components", "")
 DefNode(GeometryNode, GEO_NODE_SET_POSITION, 0, "SET_POSITION", SetPosition, "Set Position", "")
diff --git a/source/blender/nodes/geometry/nodes/legacy/node_geo_attribute_proximity.cc b/source/blender/nodes/geometry/nodes/legacy/node_geo_attribute_proximity.cc
index 0cf411343cf..6120118f611 100644
--- a/source/blender/nodes/geometry/nodes/legacy/node_geo_attribute_proximity.cc
+++ b/source/blender/nodes/geometry/nodes/legacy/node_geo_attribute_proximity.cc
@@ -232,7 +232,7 @@ static void geo_node_attribute_proximity_exec(GeoNodeExecParams params)
 
 }  // namespace blender::nodes
 
-void register_node_type_geo_attribute_proximity()
+void register_node_type_geo_legacy_attribute_proximity()
 {
   static bNodeType ntype;
 
diff --git a/source/blender/nodes/geometry/nodes/node_geo_proximity.cc b/source/blender/nodes/geometry/nodes/node_geo_proximity.cc
new file mode 100644
index 00000000000..2b1de5fbf95
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_proximity.cc
@@ -0,0 +1,235 @@
+/*
+ * 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 "BLI_task.hh"
+#include "BLI_timeit.hh"
+
+#include "DNA_mesh_types.h"
+
+#include "BKE_bvhutils.h"
+#include "BKE_geometry_set.hh"
+
+#include "UI_interface.h"
+#include "UI_resources.h"
+
+#include "node_geometry_util.hh"
+
+namespace blender::nodes {
+
+static void geo_node_proximity_declare(NodeDeclarationBuilder &b)
+{
+  b.add_input<decl::Vector>("Source Position").implicit_field();
+  b.add_input<decl::Geometry>("Target");
+  b.add_output<decl::Vector>("Position").dependent_field();
+  b.add_output<decl::Float>("Distance").dependent_field();
+}
+
+static void geo_node_proximity_layout(uiLayout *layout, bContext *UNUSED(C), PointerRNA *ptr)
+{
+  uiItemR(layout, ptr, "target_element", 0, "", ICON_NONE);
+}
+
+static void geo_proximity_init(bNodeTree *UNUSED(ntree), bNode *node)
+{
+  NodeGeometryProximity *node_storage = (NodeGeometryProximity *)MEM_callocN(
+      sizeof(NodeGeometryProximity), __func__);
+  node_storage->target_element = GEO_NODE_PROX_TARGET_FACES;
+  node->storage = node_storage;
+}
+
+static void calculate_mesh_proximity(const VArray<float3> &positions,
+                                     const IndexMask mask,
+                                     const Mesh &mesh,
+                                     const GeometryNodeProximityTargetType type,
+                                     const MutableSpan<float> r_distances,
+                                     const MutableSpan<float3> r_locations)
+{
+  BVHTreeFromMesh bvh_data;
+  switch (type) {
+    case GEO_NODE_PROX_TARGET_POINTS:
+      BKE_bvhtree_from_mesh_get(&bvh_data, &mesh, BVHTREE_FROM_VERTS, 2);
+      break;
+    case GEO_NODE_PROX_TARGET_EDGES:
+      BKE_bvhtree_from_mesh_get(&bvh_data, &mesh, BVHTREE_FROM_EDGES, 2);
+      break;
+    case GEO_NODE_PROX_TARGET_FACES:
+      BKE_bvhtree_from_mesh_get(&bvh_data, &mesh, BVHTREE_FROM_LOOPTRI, 2);
+      break;
+  }
+
+  if (bvh_data.tree == nullptr) {
+    return;
+  }
+
+  threading::parallel_for(mask.index_range(), 512, [&](IndexRange range) {
+    BVHTreeNearest nearest;
+    copy_v3_fl(nearest.co, FLT_MAX);
+    nearest.index = -1;
+
+    for (int i : range) {
+      const int index = mask[i];
+      /* Use the distance to the last found point as upper bound to speedup the bvh lookup. */
+      nearest.dist_sq = float3::distance_squared(nearest.co, positions[index]);
+
+      BLI_bvhtree_find_nearest(
+          bvh_data.tree, positions[index], &nearest, bvh_data.nearest_callback, &bvh_data);
+
+      if (nearest.dist_sq < r_distances[index]) {
+        r_distances[index] = nearest.dist_sq;
+        if (!r_locations.is_empty()) {
+          r_locations[index] = nearest.co;
+        }
+      }
+    }
+  });
+
+  free_bvhtree_from_mesh(&bvh_data);
+}
+
+static void calculate_pointcloud_proximity(const VArray<float3> &positions,
+                                           const IndexMask mask,
+                                           const PointCloud &pointcloud,
+                                           const MutableSpan<float> r_distances,
+                                           const MutableSpan<float3> r_locations)
+{
+  BVHTreeFromPointCloud bvh_data;
+  BKE_bvhtree_from_pointcloud_get(&bvh_data, &pointcloud, 2);
+  if (bvh_data.tree == nullptr) {
+    return;
+  }
+
+  threading::parallel_for(mask.index_range(), 512, [&](IndexRange range) {
+    BVHTreeNearest nearest;
+    copy_v3_fl(nearest.co, FLT_MAX);
+    nearest.index = -1;
+
+    for (int i : range) {
+      const int index = mask[i];
+      /* Use the distance to the closest point in the mesh to speedup the pointcloud bvh lookup.
+       * This is ok because we only need to find the closest point in the pointcloud if it's
+       * closer than the mesh. */
+      nearest.dist_sq = r_distances[index];
+
+      BLI_bvhtree_find_nearest(
+          bvh_data.tree, positions[index], &nearest, bvh_data.nearest_callback, &bvh_data);
+
+      if (nearest.dist_sq < r_distances[index]) {
+        r_distances[index] = nearest.dist_sq;
+        if (!r_locations.is_empty()) {
+          r_locations[index] = nearest.co;
+        }
+      }
+    }
+  });
+
+  free_bvhtree_from_pointcloud(&bvh_data);
+}
+
+class ProximityFunction : public fn::MultiFunction {
+ private:
+  GeometrySet target_;
+  GeometryNodeProximityTargetType type_;
+
+ public:
+  ProximityFunction(GeometrySet target, GeometryNodeProximityTargetType type)
+      : target_(std::move(target)), type_(type)
+  {
+    static fn::MFSignature signature = create_signature();
+    this->set_signature(&signature);
+  }
+
+  static fn::MFSignature create_signature()
+  {
+    blender::fn::MFSignatureBuilder signature{"Geometry Proximity"};
+    signature.single_input<float3>("Source Position");
+    signature.single_output<float3>("Position");
+    signature.single_output<float>("Distance");
+    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,
+                                                                               "Source Position");
+    MutableSpan<float3> positions = params.uninitialized_single_output_if_required<float3>(
+        1, "Position");
+    /* Make sure there is a distance array, used for finding the smaller distance when there are
+     * multiple components. Theoretically it would be possible to avoid using the distance array
+     * when there is only one component. However, this only adds an allocation and a single float
+     * comparison per vertex, so it's likely not worth it. */
+    MutableSpan<float> distances = params.uninitialized_single_output<float>(2, "Distance");
+
+    distances.fill(FLT_MAX);
+
+    if (target_.has_mesh()) {
+      calculate_mesh_proximity(
+          src_positions, mask, *target_.get_mesh_for_read(), type_, distances, positions);
+    }
+
+    if (target_.has_pointcloud() && type_ == GEO_NODE_PROX_TARGET_POINTS) {
+      calculate_pointcloud_proximity(
+          src_positions, mask, *target_.get_pointcloud_for_read(), distances, positions);
+    }
+
+    if (params.single_output_is_required(2, "Distance")) {
+      threading::parallel_for(mask.index_range(), 2048, [&](IndexRange range) {
+        for (const int i : range) {
+          const int j = mask[i];
+          distances[j] = std::sqrt(distances[j]);
+        }
+      });
+    }
+  }
+};
+
+static void geo_node_proximity_exec(GeoNodeExecParams params)
+{
+  GeometrySet geometry_set_target = params.extract_input<GeometrySet>("Target");
+
+  if (!geometry_set_target.has_mesh() && !geometry_set_target.has_pointcloud()) {
+    params.set_output("Position", fn::make_constant_field<float3>({0.0f, 0.0f, 0.0f}));
+    params.set_output("Distance", fn::make_constant_field<float>({0.0f}));
+    return;
+  }
+
+  const NodeGeometryProximity &storage = *(const NodeGeometryProximity *)params.node().storage;
+  Field<float3> position_field = params.extract_input<Field<float3>>("Source Position");
+
+  auto proximity_fn = std::make_unique<ProximityFunction>(
+      std::move(geometry_set_target),
+      static_cast<GeometryNodeProximityTargetType>(storage.target_element));
+  auto proximity_op = std::make_shared<FieldOperation>(
+      FieldOperation(std::move(proximity_fn), {std::move(position_field)}));
+
+  params.set_output("Position", Field<float3>(proximity_op, 0));
+  params.set_output("Distance", Field<float>(proximity_op, 1));
+}
+
+}  // namespace blender::nodes
+
+void register_node_type_geo_proximity()
+{
+  static bNodeType ntype;
+
+  geo_node_type_base(&ntype, GEO_NODE_PROXIMITY, "Geometry Proximity", NODE_CLASS_GEOMETRY, 0);
+  node_type_init(&ntype, blender::nodes::geo_proximity_init);
+  node_type_storage(
+      &ntype, "NodeGeometryProximity", node_free_standard_storage, node_copy_standard_storage);
+  ntype.declare = blender::nodes::geo_node_proximity_declare;
+  ntype.geometry_node_execute = blender::nodes::geo_node_proximity_exec;
+  ntype.draw_buttons = blender::nodes::geo_node_proximity_layout;
+  nodeRegisterType(&ntype);
+}