Geometry Nodes: Add "Location" output to Attribute Proximity node

This patch adds an output field to the Attribute Proximity node and
renames the existing string socket from "Result" to "Distance".
  - The "Distance" output contains distance to the closest position
    on the Target geometry.
  - The new "Location" output contains the coordinates of the closest
    position on the Target geometry.

A basic use case for this data is a simple shrinkwrap operation.

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

3 files changed, 90 insertions, 43 deletions

diff --git a/source/blender/blenkernel/BKE_blender_version.h b/source/blender/blenkernel/BKE_blender_version.h
index 7db6980c91a..17eb6e19292 100644
--- a/source/blender/blenkernel/BKE_blender_version.h
+++ b/source/blender/blenkernel/BKE_blender_version.h
@@ -39,7 +39,7 @@ extern "C" {
 
 /* Blender file format version. */
 #define BLENDER_FILE_VERSION BLENDER_VERSION
-#define BLENDER_FILE_SUBVERSION 8
+#define BLENDER_FILE_SUBVERSION 9
 
 /* Minimum Blender version that supports reading file written with the current
  * version. Older Blender versions will test this and show a warning if the file
diff --git a/source/blender/blenloader/intern/versioning_290.c b/source/blender/blenloader/intern/versioning_290.c
index 02de0b0ff08..b05d5ae7d26 100644
--- a/source/blender/blenloader/intern/versioning_290.c
+++ b/source/blender/blenloader/intern/versioning_290.c
@@ -1722,16 +1722,7 @@ void blo_do_versions_290(FileData *fd, Library *UNUSED(lib), Main *bmain)
     FOREACH_NODETREE_END;
   }
 
-  /**
-   * Versioning code until next subversion bump goes here.
-   *
-   * \note Be sure to check when bumping the version:
-   * - "versioning_userdef.c", #blo_do_versions_userdef
-   * - "versioning_userdef.c", #do_versions_theme
-   *
-   * \note Keep this message at the bottom of the function.
-   */
-  {
+  if (!MAIN_VERSION_ATLEAST(bmain, 293, 9)) {
     if (!DNA_struct_elem_find(fd->filesdna, "SceneEEVEE", "float", "bokeh_overblur")) {
       LISTBASE_FOREACH (Scene *, scene, &bmain->scenes) {
         scene->eevee.bokeh_neighbor_max = 10.0f;
@@ -1753,6 +1744,24 @@ void blo_do_versions_290(FileData *fd, Library *UNUSED(lib), Main *bmain)
       }
     }
 
+    FOREACH_NODETREE_BEGIN (bmain, ntree, id) {
+      if (ntree->type == NTREE_GEOMETRY) {
+        version_node_socket_name(ntree, GEO_NODE_ATTRIBUTE_PROXIMITY, "Result", "Distance");
+      }
+    }
+    FOREACH_NODETREE_END;
+  }
+
+  /**
+   * Versioning code until next subversion bump goes here.
+   *
+   * \note Be sure to check when bumping the version:
+   * - "versioning_userdef.c", #blo_do_versions_userdef
+   * - "versioning_userdef.c", #do_versions_theme
+   *
+   * \note Keep this message at the bottom of the function.
+   */
+  {
     /* Keep this block, even when empty. */
   }
 }
diff --git a/source/blender/nodes/geometry/nodes/node_geo_attribute_proximity.cc b/source/blender/nodes/geometry/nodes/node_geo_attribute_proximity.cc
index 51f208ed07d..2ab76540bdf 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_attribute_proximity.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_proximity.cc
@@ -31,7 +31,8 @@
 static bNodeSocketTemplate geo_node_attribute_proximity_in[] = {
     {SOCK_GEOMETRY, N_("Geometry")},
     {SOCK_GEOMETRY, N_("Target")},
-    {SOCK_STRING, N_("Result")},
+    {SOCK_STRING, N_("Distance")},
+    {SOCK_STRING, N_("Location")},
     {-1, ""},
 };
 
@@ -60,50 +61,66 @@ static void geo_attribute_proximity_init(bNodeTree *UNUSED(ntree), bNode *node)
 namespace blender::nodes {
 
 static void proximity_calc(MutableSpan<float> distance_span,
+                           MutableSpan<float3> location_span,
                            Span<float3> positions,
                            BVHTreeFromMesh &tree_data_mesh,
                            BVHTreeFromPointCloud &tree_data_pointcloud,
                            const bool bvh_mesh_success,
-                           const bool bvh_pointcloud_success)
+                           const bool bvh_pointcloud_success,
+                           const bool store_distances,
+                           const bool store_locations)
 {
-  /* The pointcloud loop uses the values already in the span,
-   * which is only set if the mesh BVH is used (because it's first). */
-  if (!bvh_mesh_success) {
-    distance_span.fill(FLT_MAX);
-  }
-
   IndexRange range = positions.index_range();
   parallel_for(range, 512, [&](IndexRange range) {
-    BVHTreeNearest nearest;
+    BVHTreeNearest nearest_from_mesh;
+    BVHTreeNearest nearest_from_pointcloud;
+
+    copy_v3_fl(nearest_from_mesh.co, FLT_MAX);
+    copy_v3_fl(nearest_from_pointcloud.co, FLT_MAX);
+
+    nearest_from_mesh.index = -1;
+    nearest_from_pointcloud.index = -1;
 
-    if (bvh_mesh_success) {
-      copy_v3_fl(nearest.co, FLT_MAX);
-      nearest.index = -1;
+    for (int i : range) {
+      /* Use the distance to the last found point as upper bound to speedup the bvh lookup. */
+      nearest_from_mesh.dist_sq = len_squared_v3v3(nearest_from_mesh.co, positions[i]);
 
-      for (int i : range) {
-        nearest.dist_sq = len_squared_v3v3(nearest.co, positions[i]);
+      if (bvh_mesh_success) {
         BLI_bvhtree_find_nearest(tree_data_mesh.tree,
                                  positions[i],
-                                 &nearest,
+                                 &nearest_from_mesh,
                                  tree_data_mesh.nearest_callback,
                                  &tree_data_mesh);
-        distance_span[i] = sqrtf(nearest.dist_sq);
       }
-    }
 
-    if (bvh_pointcloud_success) {
-      copy_v3_fl(nearest.co, FLT_MAX);
-      nearest.index = -1;
+      /* 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_from_pointcloud.dist_sq = nearest_from_mesh.dist_sq;
 
-      for (int i : range) {
-        /* Use the distance to the last found point as upper bound to speedup the bvh lookup. */
-        nearest.dist_sq = len_squared_v3v3(nearest.co, positions[i]);
+      if (bvh_pointcloud_success) {
         BLI_bvhtree_find_nearest(tree_data_pointcloud.tree,
                                  positions[i],
-                                 &nearest,
+                                 &nearest_from_pointcloud,
                                  tree_data_pointcloud.nearest_callback,
                                  &tree_data_pointcloud);
-        distance_span[i] = std::min(distance_span[i], sqrtf(nearest.dist_sq));
+      }
+
+      if (nearest_from_pointcloud.dist_sq < nearest_from_mesh.dist_sq) {
+        if (store_distances) {
+          distance_span[i] = sqrtf(nearest_from_pointcloud.dist_sq);
+        }
+        if (store_locations) {
+          location_span[i] = nearest_from_pointcloud.co;
+        }
+      }
+      else {
+        if (store_distances) {
+          distance_span[i] = sqrtf(nearest_from_mesh.dist_sq);
+        }
+        if (store_locations) {
+          location_span[i] = nearest_from_mesh.co;
+        }
       }
     }
   });
@@ -151,14 +168,18 @@ static void attribute_calc_proximity(GeometryComponent &component,
   /* This node works on the "point" domain, since that is where positions are stored. */
   const AttributeDomain result_domain = ATTR_DOMAIN_POINT;
 
-  const std::string result_attribute_name = params.get_input<std::string>("Result");
+  const std::string distance_attribute_name = params.get_input<std::string>("Distance");
   OutputAttributePtr distance_attribute = component.attribute_try_get_for_output(
-      result_attribute_name, result_domain, CD_PROP_FLOAT);
+      distance_attribute_name, result_domain, CD_PROP_FLOAT);
+
+  const std::string location_attribute_name = params.get_input<std::string>("Location");
+  OutputAttributePtr location_attribute = component.attribute_try_get_for_output(
+      location_attribute_name, result_domain, CD_PROP_FLOAT3);
 
   ReadAttributePtr position_attribute = component.attribute_try_get_for_read("position");
   BLI_assert(position_attribute->custom_data_type() == CD_PROP_FLOAT3);
 
-  if (!distance_attribute || !position_attribute) {
+  if (!position_attribute || (!distance_attribute && !location_attribute)) {
     return;
   }
 
@@ -183,12 +204,24 @@ static void attribute_calc_proximity(GeometryComponent &component,
                                                  tree_data_pointcloud);
   }
 
-  proximity_calc(distance_attribute->get_span_for_write_only<float>(),
-                 position_attribute->get_span<float3>(),
+  Span<float3> position_span = position_attribute->get_span<float3>();
+
+  MutableSpan<float> distance_span = distance_attribute ?
+                                         distance_attribute->get_span_for_write_only<float>() :
+                                         MutableSpan<float>();
+  MutableSpan<float3> location_span = location_attribute ?
+                                          location_attribute->get_span_for_write_only<float3>() :
+                                          MutableSpan<float3>();
+
+  proximity_calc(distance_span,
+                 location_span,
+                 position_span,
                  tree_data_mesh,
                  tree_data_pointcloud,
                  bvh_mesh_success,
-                 bvh_pointcloud_success);
+                 bvh_pointcloud_success,
+                 distance_attribute,  /* Boolean. */
+                 location_attribute); /* Boolean. */
 
   if (bvh_mesh_success) {
     free_bvhtree_from_mesh(&tree_data_mesh);
@@ -197,7 +230,12 @@ static void attribute_calc_proximity(GeometryComponent &component,
     free_bvhtree_from_pointcloud(&tree_data_pointcloud);
   }
 
-  distance_attribute.apply_span_and_save();
+  if (distance_attribute) {
+    distance_attribute.apply_span_and_save();
+  }
+  if (location_attribute) {
+    location_attribute.apply_span_and_save();
+  }
 }
 
 static void geo_node_attribute_proximity_exec(GeoNodeExecParams params)