Geometry Nodes: Add Attribute Proximity Node

This node calculates a distance from each point to the closest position
on a target geometry, similar to the vertex weight proximity modifier.
Mapping the output distance to a different range can be done with an
attribute math node after this node.

A drop-down changes whether to calculate distances from points,
edges, or faces. In points mode, the node also calculates distances
from point cloud points.

Design task and use cases: T84842

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

4 files changed, 232 insertions, 0 deletions

diff --git a/source/blender/nodes/CMakeLists.txt b/source/blender/nodes/CMakeLists.txt
index aedca4b34fb..76b19685456 100644
--- a/source/blender/nodes/CMakeLists.txt
+++ b/source/blender/nodes/CMakeLists.txt
@@ -147,6 +147,7 @@ set(SRC
   geometry/nodes/node_geo_attribute_math.cc
   geometry/nodes/node_geo_attribute_mix.cc
   geometry/nodes/node_geo_attribute_sample_texture.cc
+  geometry/nodes/node_geo_attribute_proximity.cc
   geometry/nodes/node_geo_attribute_randomize.cc
   geometry/nodes/node_geo_attribute_vector_math.cc
   geometry/nodes/node_geo_boolean.cc
@@ -351,6 +352,10 @@ if(WITH_INTERNATIONAL)
   add_definitions(-DWITH_INTERNATIONAL)
 endif()
 
+if(WITH_TBB)
+  add_definitions(-DWITH_TBB)
+endif()
+
 if(WITH_IMAGE_OPENEXR)
   add_definitions(-DWITH_OPENEXR)
 endif()
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index 9b391ab7981..54a728f5bfe 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -28,6 +28,7 @@ void register_node_type_geo_group(void);
 
 void register_node_type_geo_attribute_fill(void);
 void register_node_type_geo_attribute_vector_math(void);
+void register_node_type_geo_attribute_proximity(void);
 void register_node_type_geo_boolean(void);
 void register_node_type_geo_edge_split(void);
 void register_node_type_geo_transform(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index f9730af0c08..1cf08a419c9 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -293,6 +293,8 @@ DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_SAMPLE_TEXTURE, def_geo_attribute_sampl
 DefNode(GeometryNode, GEO_NODE_POINTS_TO_VOLUME, def_geo_points_to_volume, "POINTS_TO_VOLUME", PointsToVolume, "Points to Volume", "")
 DefNode(GeometryNode, GEO_NODE_COLLECTION_INFO, def_geo_collection_info, "COLLECTION_INFO", CollectionInfo, "Collection Info", "")
 DefNode(GeometryNode, GEO_NODE_IS_VIEWPORT, 0, "IS_VIEWPORT", IsViewport, "Is Viewport", "")
+DefNode(GeometryNode, GEO_NODE_ATTRIBUTE_PROXIMITY, def_geo_attribute_proximity, "ATTRIBUTE_PROXIMITY", AttributeProximity, "Attribute Proximity", "")
+
 
 /* undefine macros */
 #undef DefNode
diff --git a/source/blender/nodes/geometry/nodes/node_geo_attribute_proximity.cc b/source/blender/nodes/geometry/nodes/node_geo_attribute_proximity.cc
new file mode 100644
index 00000000000..14b1c4e3a59
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_proximity.cc
@@ -0,0 +1,224 @@
+/*
+ * 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_kdtree.h"
+#include "BLI_task.hh"
+
+#include "BKE_bvhutils.h"
+#include "BLI_kdopbvh.h"
+
+#include "node_geometry_util.hh"
+
+#include "BLI_timeit.hh"
+#include "DNA_mesh_types.h"
+
+static bNodeSocketTemplate geo_node_attribute_proximity_in[] = {
+    {SOCK_GEOMETRY, N_("Geometry")},
+    {SOCK_GEOMETRY, N_("Target")},
+    {SOCK_STRING, N_("Result")},
+    {-1, ""},
+};
+
+static bNodeSocketTemplate geo_node_attribute_proximity_out[] = {
+    {SOCK_GEOMETRY, N_("Geometry")},
+    {-1, ""},
+};
+
+static void geo_attribute_proximity_init(bNodeTree *UNUSED(ntree), bNode *node)
+{
+  NodeGeometryAttributeProximity *node_storage = (NodeGeometryAttributeProximity *)MEM_callocN(
+      sizeof(NodeGeometryAttributeProximity), __func__);
+
+  node_storage->target_geometry_element =
+      GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_FACES;
+  node->storage = node_storage;
+}
+
+namespace blender::nodes {
+
+static void proximity_calc(MutableSpan<float> distance_span,
+                           Span<float3> positions,
+                           BVHTreeFromMesh &tree_data_mesh,
+                           BVHTreeFromPointCloud &tree_data_pointcloud,
+                           const bool bvh_mesh_success,
+                           const bool bvh_pointcloud_success)
+{
+
+  IndexRange range = positions.index_range();
+  parallel_for(range, 512, [&](IndexRange range) {
+    BVHTreeNearest nearest;
+
+    if (bvh_mesh_success) {
+      copy_v3_fl(nearest.co, FLT_MAX);
+      nearest.index = -1;
+
+      for (int i : range) {
+        nearest.dist_sq = len_squared_v3v3(nearest.co, positions[i]);
+        BLI_bvhtree_find_nearest(tree_data_mesh.tree,
+                                 positions[i],
+                                 &nearest,
+                                 tree_data_mesh.nearest_callback,
+                                 &tree_data_mesh);
+        distance_span[i] = sqrtf(nearest.dist_sq);
+      }
+    }
+
+    /* The next loop(s) use the values already in the span. */
+    if (!bvh_mesh_success) {
+      distance_span.fill(FLT_MAX);
+    }
+
+    if (bvh_pointcloud_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.dist_sq = len_squared_v3v3(nearest.co, positions[i]);
+        BLI_bvhtree_find_nearest(tree_data_pointcloud.tree,
+                                 positions[i],
+                                 &nearest,
+                                 tree_data_pointcloud.nearest_callback,
+                                 &tree_data_pointcloud);
+        distance_span[i] = std::min(distance_span[i], sqrtf(nearest.dist_sq));
+      }
+    }
+  });
+}
+
+static bool bvh_from_mesh(const Mesh *target_mesh,
+                          int target_geometry_element,
+                          BVHTreeFromMesh &r_tree_data_mesh)
+{
+  BVHCacheType bvh_type = BVHTREE_FROM_LOOPTRI;
+  switch (target_geometry_element) {
+    case GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_POINTS:
+      bvh_type = BVHTREE_FROM_VERTS;
+      break;
+    case GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_EDGES:
+      bvh_type = BVHTREE_FROM_EDGES;
+      break;
+    case GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_FACES:
+      bvh_type = BVHTREE_FROM_LOOPTRI;
+      break;
+  }
+
+  /* This only updates a cache and can be considered to be logically const. */
+  BKE_bvhtree_from_mesh_get(&r_tree_data_mesh, const_cast<Mesh *>(target_mesh), bvh_type, 2);
+  if (r_tree_data_mesh.tree == nullptr) {
+    return false;
+  }
+  return true;
+}
+
+static bool bvh_from_pointcloud(const PointCloud *target_pointcloud,
+                                BVHTreeFromPointCloud &r_tree_data_pointcloud)
+{
+  BKE_bvhtree_from_pointcloud_get(&r_tree_data_pointcloud, target_pointcloud, 2);
+  if (r_tree_data_pointcloud.tree == nullptr) {
+    return false;
+  }
+  return true;
+}
+
+static void attribute_calc_proximity(GeometryComponent &component,
+                                     GeometrySet &geometry_set_target,
+                                     GeoNodeExecParams &params)
+{
+  const std::string result_attribute_name = params.get_input<std::string>("Result");
+  OutputAttributePtr distance_attribute = component.attribute_try_get_for_output(
+      result_attribute_name, ATTR_DOMAIN_POINT, CD_PROP_FLOAT);
+
+  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) {
+    return;
+  }
+
+  const bNode &node = params.node();
+  const NodeGeometryAttributeProximity &storage = *(const NodeGeometryAttributeProximity *)
+                                                       node.storage;
+
+  BVHTreeFromMesh tree_data_mesh;
+  BVHTreeFromPointCloud tree_data_pointcloud;
+  bool bvh_mesh_success = false;
+  bool bvh_pointcloud_success = false;
+
+  if (geometry_set_target.has_mesh()) {
+    bvh_mesh_success = bvh_from_mesh(
+        geometry_set_target.get_mesh_for_read(), storage.target_geometry_element, tree_data_mesh);
+  }
+
+  if (geometry_set_target.has_pointcloud() &&
+      storage.target_geometry_element ==
+          GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_POINTS) {
+    bvh_pointcloud_success = bvh_from_pointcloud(geometry_set_target.get_pointcloud_for_read(),
+                                                 tree_data_pointcloud);
+  }
+
+  proximity_calc(distance_attribute->get_span_for_write_only<float>(),
+                 position_attribute->get_span<float3>(),
+                 tree_data_mesh,
+                 tree_data_pointcloud,
+                 bvh_mesh_success,
+                 bvh_pointcloud_success);
+
+  if (bvh_mesh_success) {
+    free_bvhtree_from_mesh(&tree_data_mesh);
+  }
+  if (bvh_pointcloud_success) {
+    free_bvhtree_from_pointcloud(&tree_data_pointcloud);
+  }
+
+  distance_attribute.apply_span_and_save();
+}
+
+static void geo_node_attribute_proximity_exec(GeoNodeExecParams params)
+{
+  GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
+  GeometrySet geometry_set_target = params.extract_input<GeometrySet>("Target");
+
+  if (geometry_set.has<MeshComponent>()) {
+    attribute_calc_proximity(
+        geometry_set.get_component_for_write<MeshComponent>(), geometry_set_target, params);
+  }
+  if (geometry_set.has<PointCloudComponent>()) {
+    attribute_calc_proximity(
+        geometry_set.get_component_for_write<PointCloudComponent>(), geometry_set_target, params);
+  }
+
+  params.set_output("Geometry", geometry_set);
+}
+
+}  // namespace blender::nodes
+
+void register_node_type_geo_attribute_proximity()
+{
+  static bNodeType ntype;
+
+  geo_node_type_base(
+      &ntype, GEO_NODE_ATTRIBUTE_PROXIMITY, "Attribute Proximity", NODE_CLASS_ATTRIBUTE, 0);
+  node_type_socket_templates(
+      &ntype, geo_node_attribute_proximity_in, geo_node_attribute_proximity_out);
+  node_type_init(&ntype, geo_attribute_proximity_init);
+  node_type_storage(&ntype,
+                    "NodeGeometryAttributeProximity",
+                    node_free_standard_storage,
+                    node_copy_standard_storage);
+  ntype.geometry_node_execute = blender::nodes::geo_node_attribute_proximity_exec;
+  nodeRegisterType(&ntype);
+}