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

11 files changed, 344 insertions, 0 deletions

diff --git a/source/blender/blenkernel/BKE_bvhutils.h b/source/blender/blenkernel/BKE_bvhutils.h
index 84a17136361..78908908343 100644
--- a/source/blender/blenkernel/BKE_bvhutils.h
+++ b/source/blender/blenkernel/BKE_bvhutils.h
@@ -38,6 +38,7 @@ struct BMEditMesh;
 struct MFace;
 struct MVert;
 struct Mesh;
+struct PointCloud;
 
 struct BVHCache;
 
@@ -249,6 +250,20 @@ float bvhtree_sphereray_tri_intersection(const BVHTreeRay *ray,
                                          const float v1[3],
                                          const float v2[3]);
 
+typedef struct BVHTreeFromPointCloud {
+  struct BVHTree *tree;
+
+  BVHTree_NearestPointCallback nearest_callback;
+
+  const float (*coords)[3];
+} BVHTreeFromPointCloud;
+
+BVHTree *BKE_bvhtree_from_pointcloud_get(struct BVHTreeFromPointCloud *data,
+                                         const struct PointCloud *pointcloud,
+                                         const int tree_type);
+
+void free_bvhtree_from_pointcloud(struct BVHTreeFromPointCloud *data);
+
 /**
  * BVHCache
  */
diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index 9ee21c6e825..0b8ef70de2a 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1368,6 +1368,7 @@ int ntreeTexExecTree(struct bNodeTree *ntree,
 #define GEO_NODE_POINTS_TO_VOLUME 1022
 #define GEO_NODE_COLLECTION_INFO 1023
 #define GEO_NODE_IS_VIEWPORT 1024
+#define GEO_NODE_ATTRIBUTE_PROXIMITY 1025
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/bvhutils.c b/source/blender/blenkernel/intern/bvhutils.c
index fd5cb33f02d..790fb128c7c 100644
--- a/source/blender/blenkernel/intern/bvhutils.c
+++ b/source/blender/blenkernel/intern/bvhutils.c
@@ -27,6 +27,7 @@
 
 #include "DNA_mesh_types.h"
 #include "DNA_meshdata_types.h"
+#include "DNA_pointcloud_types.h"
 
 #include "BLI_linklist.h"
 #include "BLI_math.h"
@@ -1717,3 +1718,41 @@ void free_bvhtree_from_mesh(struct BVHTreeFromMesh *data)
 
   memset(data, 0, sizeof(*data));
 }
+
+/** \} */
+
+/* -------------------------------------------------------------------- */
+/** \name Point Cloud BVH Building
+ * \{ */
+
+BVHTree *BKE_bvhtree_from_pointcloud_get(BVHTreeFromPointCloud *data,
+                                         const PointCloud *pointcloud,
+                                         const int tree_type)
+{
+  BVHTree *tree = BLI_bvhtree_new(pointcloud->totpoint, 0.0f, tree_type, 6);
+  if (!tree) {
+    return NULL;
+  }
+
+  for (int i = 0; i < pointcloud->totpoint; i++) {
+    BLI_bvhtree_insert(tree, i, pointcloud->co[i], 1);
+  }
+  BLI_assert(BLI_bvhtree_get_len(tree) == pointcloud->totpoint);
+  BLI_bvhtree_balance(tree);
+
+  data->coords = pointcloud->co;
+  data->tree = tree;
+  data->nearest_callback = NULL;
+
+  return tree;
+}
+
+void free_bvhtree_from_pointcloud(BVHTreeFromPointCloud *data)
+{
+  if (data->tree) {
+    BLI_bvhtree_free(data->tree);
+  }
+  memset(data, 0, sizeof(*data));
+}
+
+/** \} */
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index d78ba2a414e..44efbe64202 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -4753,6 +4753,7 @@ static void registerGeometryNodes()
 
   register_node_type_geo_attribute_compare();
   register_node_type_geo_attribute_fill();
+  register_node_type_geo_attribute_proximity();
   register_node_type_geo_attribute_vector_math();
   register_node_type_geo_triangulate();
   register_node_type_geo_edge_split();
diff --git a/source/blender/editors/space_node/drawnode.c b/source/blender/editors/space_node/drawnode.c
index d7edb1ecbed..b16ccb9bce4 100644
--- a/source/blender/editors/space_node/drawnode.c
+++ b/source/blender/editors/space_node/drawnode.c
@@ -3381,6 +3381,13 @@ static void node_geometry_buts_collection_info(uiLayout *layout,
   uiItemR(layout, ptr, "transform_space", UI_ITEM_R_EXPAND, NULL, ICON_NONE);
 }
 
+static void node_geometry_buts_attribute_proximity(uiLayout *layout,
+                                                   bContext *UNUSED(C),
+                                                   PointerRNA *ptr)
+{
+  uiItemR(layout, ptr, "target_geometry_element", DEFAULT_FLAGS, "", ICON_NONE);
+}
+
 static void node_geometry_set_butfunc(bNodeType *ntype)
 {
   switch (ntype->type) {
@@ -3444,6 +3451,9 @@ static void node_geometry_set_butfunc(bNodeType *ntype)
     case GEO_NODE_COLLECTION_INFO:
       ntype->draw_buttons = node_geometry_buts_collection_info;
       break;
+    case GEO_NODE_ATTRIBUTE_PROXIMITY:
+      ntype->draw_buttons = node_geometry_buts_attribute_proximity;
+      break;
   }
 }
 
diff --git a/source/blender/makesdna/DNA_node_types.h b/source/blender/makesdna/DNA_node_types.h
index a69af18ded2..edab43c752e 100644
--- a/source/blender/makesdna/DNA_node_types.h
+++ b/source/blender/makesdna/DNA_node_types.h
@@ -1202,6 +1202,13 @@ typedef struct NodeGeometryCollectionInfo {
   char _pad[7];
 } NodeGeometryCollectionInfo;
 
+typedef struct NodeGeometryAttributeProximity {
+  /* GeometryNodeAttributeProximityTargetGeometryElement. */
+  uint8_t target_geometry_element;
+
+  char _pad[7];
+} NodeGeometryAttributeProximity;
+
 /* script node mode */
 #define NODE_SCRIPT_INTERNAL 0
 #define NODE_SCRIPT_EXTERNAL 1
@@ -1589,6 +1596,12 @@ typedef enum NodeShaderOutputTarget {
 
 /* Geometry Nodes */
 
+typedef enum GeometryNodeAttributeProximityTargetGeometryElement {
+  GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_POINTS = 0,
+  GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_EDGES = 1,
+  GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_FACES = 2,
+} GeometryNodeAttributeProximityTargetGeometryElement;
+
 /* Boolean Node */
 typedef enum GeometryNodeBooleanOperation {
   GEO_NODE_BOOLEAN_INTERSECT = 0,
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index ab6d4fa985c..7f66d5ff70f 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -8938,6 +8938,39 @@ 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 const EnumPropertyItem target_geometry_element[] = {
+      {GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_POINTS,
+       "POINTS",
+       ICON_NONE,
+       "Points",
+       "Calculate proximity to the target's points (usually faster than the other two modes)"},
+      {GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_EDGES,
+       "EDGES",
+       ICON_NONE,
+       "Edges",
+       "Calculate proximity to the target's edges"},
+      {GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_FACES,
+       "FACES",
+       ICON_NONE,
+       "Faces",
+       "Calculate proximity to the target's faces"},
+      {0, NULL, 0, NULL, NULL},
+  };
+
+  PropertyRNA *prop;
+
+  RNA_def_struct_sdna_from(srna, "NodeGeometryAttributeProximity", "storage");
+
+  prop = RNA_def_property(srna, "target_geometry_element", PROP_ENUM, PROP_NONE);
+  RNA_def_property_enum_items(prop, target_geometry_element);
+  RNA_def_property_enum_default(prop, GEO_NODE_ATTRIBUTE_PROXIMITY_TARGET_GEOMETRY_ELEMENT_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 rna_def_shader_node(BlenderRNA *brna)
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);
+}