Geometry Nodes: Initial merge by distance node

This implements a merge by distance operation for point clouds.
Besides the geometry input, there are two others-- a selection
input to limit the operation to certain points, and the merge
distance. While it would be a reasonable feature, the distance
does not support a field currently, since that would make
the algorithm significantly more complex.

All attributes are merged to the merged points, with the values
mixed together. This same generic method is used for all attributes,
including `position`. The `id` attribute uses the value from the
first merged index for each point.

For the implementation, most of the effort goes into creating a
merge map to speed up attribute mixing. Some parts are inherently
single-threaded, like finding the final indices accounting for the
merged points. By far most of the time is spend balancing the
KD tree.

Mesh support will be added in the next commit.

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

10 files changed, 309 insertions, 0 deletions

diff --git a/release/scripts/startup/nodeitems_builtins.py b/release/scripts/startup/nodeitems_builtins.py
index d7bf5d82c6a..b841cb5dd13 100644
--- a/release/scripts/startup/nodeitems_builtins.py
+++ b/release/scripts/startup/nodeitems_builtins.py
@@ -181,6 +181,7 @@ def geometry_node_items(context):
     yield NodeItem("GeometryNodeConvexHull")
     yield NodeItem("GeometryNodeDeleteGeometry")
     yield NodeItem("GeometryNodeGeometryToInstance")
+    yield NodeItem("GeometryNodeMergeByDistance")
     yield NodeItem("GeometryNodeProximity")
     yield NodeItem("GeometryNodeJoinGeometry")
     yield NodeItem("GeometryNodeRaycast")
diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h
index b2c91fafe97..16d8ba2e6dd 100644
--- a/source/blender/blenkernel/BKE_node.h
+++ b/source/blender/blenkernel/BKE_node.h
@@ -1523,6 +1523,7 @@ struct TexResult;
 #define GEO_NODE_FLIP_FACES 1150
 #define GEO_NODE_SCALE_ELEMENTS 1151
 #define GEO_NODE_EXTRUDE_MESH 1152
+#define GEO_NODE_MERGE_BY_DISTANCE 1153
 
 /** \} */
 
diff --git a/source/blender/blenkernel/intern/node.cc b/source/blender/blenkernel/intern/node.cc
index 074e38dea87..40d0c24c9af 100644
--- a/source/blender/blenkernel/intern/node.cc
+++ b/source/blender/blenkernel/intern/node.cc
@@ -4802,6 +4802,7 @@ static void registerGeometryNodes()
   register_node_type_geo_join_geometry();
   register_node_type_geo_material_replace();
   register_node_type_geo_material_selection();
+  register_node_type_geo_merge_by_distance();
   register_node_type_geo_mesh_primitive_circle();
   register_node_type_geo_mesh_primitive_cone();
   register_node_type_geo_mesh_primitive_cube();
diff --git a/source/blender/geometry/CMakeLists.txt b/source/blender/geometry/CMakeLists.txt
index de508ddc540..9e0b53f4980 100644
--- a/source/blender/geometry/CMakeLists.txt
+++ b/source/blender/geometry/CMakeLists.txt
@@ -31,9 +31,11 @@ set(INC
 
 set(SRC
   intern/mesh_to_curve_convert.cc
+  intern/point_merge_by_distance.cc
   intern/realize_instances.cc
 
   GEO_mesh_to_curve.hh
+  GEO_point_merge_by_distance.hh
   GEO_realize_instances.hh
 )
 
diff --git a/source/blender/geometry/GEO_point_merge_by_distance.hh b/source/blender/geometry/GEO_point_merge_by_distance.hh
new file mode 100644
index 00000000000..6766f3c559d
--- /dev/null
+++ b/source/blender/geometry/GEO_point_merge_by_distance.hh
@@ -0,0 +1,38 @@
+/*
+ * 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_index_mask.hh"
+
+#pragma once
+
+struct PointCloud;
+class PointCloudComponent;
+
+/** \file
+ * \ingroup geo
+ */
+
+namespace blender::geometry {
+
+/**
+ * Merge selected points into other selected points within the \a merge_distance. The merged
+ * indices favor speed over accuracy, since the results will depend on the order of the points.
+ */
+PointCloud *point_merge_by_distance(const PointCloudComponent &src_points,
+                                    const float merge_distance,
+                                    const IndexMask selection);
+
+}  // namespace blender::geometry
diff --git a/source/blender/geometry/intern/point_merge_by_distance.cc b/source/blender/geometry/intern/point_merge_by_distance.cc
new file mode 100644
index 00000000000..daa08a3bce6
--- /dev/null
+++ b/source/blender/geometry/intern/point_merge_by_distance.cc
@@ -0,0 +1,183 @@
+/*
+ * 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 "DNA_pointcloud_types.h"
+
+#include "BKE_attribute_math.hh"
+#include "BKE_geometry_set.hh"
+#include "BKE_pointcloud.h"
+
+#include "GEO_point_merge_by_distance.hh"
+
+namespace blender::geometry {
+
+PointCloud *point_merge_by_distance(const PointCloudComponent &src_points,
+                                    const float merge_distance,
+                                    const IndexMask selection)
+{
+  const PointCloud &src_pointcloud = *src_points.get_for_read();
+  const int src_size = src_pointcloud.totpoint;
+  Span<float3> positions{reinterpret_cast<float3 *>(src_pointcloud.co), src_size};
+
+  /* Create the KD tree based on only the selected points, to speed up merge detection and
+   * balancing. */
+  KDTree_3d *tree = BLI_kdtree_3d_new(selection.size());
+  for (const int i : selection.index_range()) {
+    BLI_kdtree_3d_insert(tree, i, positions[selection[i]]);
+  }
+  BLI_kdtree_3d_balance(tree);
+
+  /* Find the duplicates in the KD tree. Because the tree only contains the selected points, the
+   * resulting indices are indices into the selection, rather than indices of the source point
+   * cloud. */
+  Array<int> selection_merge_indices(selection.size(), -1);
+  const int duplicate_count = BLI_kdtree_3d_calc_duplicates_fast(
+      tree, merge_distance, false, selection_merge_indices.data());
+  BLI_kdtree_3d_free(tree);
+
+  /* Create the new point cloud and add it to a temporary component for the attribute API. */
+  const int dst_size = src_size - duplicate_count;
+  PointCloud *dst_pointcloud = BKE_pointcloud_new_nomain(dst_size);
+  PointCloudComponent dst_points;
+  dst_points.replace(dst_pointcloud, GeometryOwnershipType::Editable);
+
+  /* By default, every point is just "merged" with itself. Then fill in the results of the merge
+   * finding, converting from indices into the selection to indices into the full input point
+   * cloud. */
+  Array<int> merge_indices(src_size);
+  for (const int i : merge_indices.index_range()) {
+    merge_indices[i] = i;
+  }
+  for (const int i : selection_merge_indices.index_range()) {
+    const int merge_index = selection_merge_indices[i];
+    if (merge_index != -1) {
+      const int src_merge_index = selection[merge_index];
+      const int src_index = selection[i];
+      merge_indices[src_index] = src_merge_index;
+    }
+  }
+
+  /* For every source index, find the corresponding index in the result by iterating through the
+   * source indices and counting how many merges happened before that point. */
+  int merged_points = 0;
+  Array<int> src_to_dst_indices(src_size);
+  for (const int i : IndexRange(src_size)) {
+    src_to_dst_indices[i] = i - merged_points;
+    if (merge_indices[i] != i) {
+      merged_points++;
+    }
+  }
+
+  /* In order to use a contiguous array as the storage for every destination point's source
+   * indices, first the number of source points must be counted for every result point. */
+  Array<int> point_merge_counts(dst_size, 0);
+  for (const int i : IndexRange(src_size)) {
+    const int merge_index = merge_indices[i];
+    const int dst_index = src_to_dst_indices[merge_index];
+    point_merge_counts[dst_index]++;
+  }
+
+  /* This array stores an offset into `merge_map` for every result point. */
+  Array<int> map_offsets(dst_size + 1);
+  int offset = 0;
+  for (const int i : IndexRange(dst_size)) {
+    map_offsets[i] = offset;
+    offset += point_merge_counts[i];
+  }
+  map_offsets.last() = offset;
+
+  point_merge_counts.fill(0);
+
+  /* This array stores all of the source indices for every result point. The size is the source
+   * size because every input point is either merged with another or copied directly. */
+  Array<int> merge_map(src_size);
+  for (const int i : IndexRange(src_size)) {
+    const int merge_index = merge_indices[i];
+    const int dst_index = src_to_dst_indices[merge_index];
+
+    const IndexRange point_range(map_offsets[dst_index],
+                                 map_offsets[dst_index + 1] - map_offsets[dst_index]);
+    MutableSpan<int> point_merge_indices = merge_map.as_mutable_span().slice(point_range);
+    point_merge_indices[point_merge_counts[dst_index]] = i;
+    point_merge_counts[dst_index]++;
+  }
+
+  Set<bke::AttributeIDRef> attributes = src_points.attribute_ids();
+
+  /* Transfer the ID attribute if it exists, using the ID of the first merged point. */
+  if (attributes.contains("id")) {
+    VArray<int> src = src_points.attribute_get_for_read<int>("id", ATTR_DOMAIN_POINT, 0);
+    bke::OutputAttribute_Typed<int> dst = dst_points.attribute_try_get_for_output_only<int>(
+        "id", ATTR_DOMAIN_POINT);
+    Span<int> src_ids = src.get_internal_span();
+    MutableSpan<int> dst_ids = dst.as_span();
+
+    threading::parallel_for(IndexRange(dst_size), 1024, [&](IndexRange range) {
+      for (const int i_dst : range) {
+        const IndexRange point_range(map_offsets[i_dst],
+                                     map_offsets[i_dst + 1] - map_offsets[i_dst]);
+        dst_ids[i_dst] = src_ids[point_range.first()];
+      }
+    });
+
+    dst.save();
+    attributes.remove_contained("id");
+  }
+
+  /* Transfer all other attributes. */
+  for (const bke::AttributeIDRef &id : attributes) {
+    if (!id.should_be_kept()) {
+      continue;
+    }
+
+    bke::ReadAttributeLookup src_attribute = src_points.attribute_try_get_for_read(id);
+    attribute_math::convert_to_static_type(src_attribute.varray.type(), [&](auto dummy) {
+      using T = decltype(dummy);
+      if constexpr (!std::is_void_v<attribute_math::DefaultMixer<T>>) {
+        bke::OutputAttribute_Typed<T> dst_attribute =
+            dst_points.attribute_try_get_for_output_only<T>(id, ATTR_DOMAIN_POINT);
+        Span<T> src = src_attribute.varray.get_internal_span().typed<T>();
+        MutableSpan<T> dst = dst_attribute.as_span();
+
+        threading::parallel_for(IndexRange(dst_size), 1024, [&](IndexRange range) {
+          for (const int i_dst : range) {
+            /* Create a separate mixer for every point to avoid allocating temporary buffers
+             * in the mixer the size of the result point cloud and to improve memory locality. */
+            attribute_math::DefaultMixer<T> mixer{dst.slice(i_dst, 1)};
+
+            const IndexRange point_range(map_offsets[i_dst],
+                                         map_offsets[i_dst + 1] - map_offsets[i_dst]);
+            Span<int> src_merge_indices = merge_map.as_span().slice(point_range);
+            for (const int i_src : src_merge_indices) {
+              mixer.mix_in(0, src[i_src]);
+            }
+
+            mixer.finalize();
+          }
+        });
+
+        dst_attribute.save();
+      }
+    });
+  }
+
+  return dst_pointcloud;
+}
+
+}  // namespace blender::geometry
diff --git a/source/blender/nodes/NOD_geometry.h b/source/blender/nodes/NOD_geometry.h
index e5c005f8c95..609d92c09df 100644
--- a/source/blender/nodes/NOD_geometry.h
+++ b/source/blender/nodes/NOD_geometry.h
@@ -132,6 +132,7 @@ void register_node_type_geo_is_viewport(void);
 void register_node_type_geo_join_geometry(void);
 void register_node_type_geo_material_replace(void);
 void register_node_type_geo_material_selection(void);
+void register_node_type_geo_merge_by_distance(void);
 void register_node_type_geo_mesh_primitive_circle(void);
 void register_node_type_geo_mesh_primitive_cone(void);
 void register_node_type_geo_mesh_primitive_cube(void);
diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h
index df820bf6dcf..8cbde6adcad 100644
--- a/source/blender/nodes/NOD_static_types.h
+++ b/source/blender/nodes/NOD_static_types.h
@@ -387,6 +387,7 @@ DefNode(GeometryNode, GEO_NODE_INSTANCES_TO_POINTS, 0, "INSTANCES_TO_POINTS", In
 DefNode(GeometryNode, GEO_NODE_IS_VIEWPORT, 0, "IS_VIEWPORT", IsViewport, "Is Viewport", "")
 DefNode(GeometryNode, GEO_NODE_JOIN_GEOMETRY, 0, "JOIN_GEOMETRY", JoinGeometry, "Join Geometry", "")
 DefNode(GeometryNode, GEO_NODE_MATERIAL_SELECTION, 0, "MATERIAL_SELECTION", MaterialSelection, "Material Selection", "")
+DefNode(GeometryNode, GEO_NODE_MERGE_BY_DISTANCE, 0, "MERGE_BY_DISTANCE", MergeByDistance, "Merge by Distance", "")
 DefNode(GeometryNode, GEO_NODE_MESH_BOOLEAN, def_geo_boolean, "MESH_BOOLEAN", MeshBoolean, "Mesh Boolean", "")
 DefNode(GeometryNode, GEO_NODE_MESH_PRIMITIVE_CIRCLE, def_geo_mesh_circle, "MESH_PRIMITIVE_CIRCLE", MeshCircle, "Mesh Circle", "")
 DefNode(GeometryNode, GEO_NODE_MESH_PRIMITIVE_CONE, def_geo_mesh_cone, "MESH_PRIMITIVE_CONE", MeshCone, "Cone", "")
diff --git a/source/blender/nodes/geometry/CMakeLists.txt b/source/blender/nodes/geometry/CMakeLists.txt
index 0e5f90b58bf..b4add633b0c 100644
--- a/source/blender/nodes/geometry/CMakeLists.txt
+++ b/source/blender/nodes/geometry/CMakeLists.txt
@@ -150,6 +150,7 @@ set(SRC
   nodes/node_geo_join_geometry.cc
   nodes/node_geo_material_replace.cc
   nodes/node_geo_material_selection.cc
+  nodes/node_geo_merge_by_distance.cc
   nodes/node_geo_mesh_primitive_circle.cc
   nodes/node_geo_mesh_primitive_cone.cc
   nodes/node_geo_mesh_primitive_cube.cc
diff --git a/source/blender/nodes/geometry/nodes/node_geo_merge_by_distance.cc b/source/blender/nodes/geometry/nodes/node_geo_merge_by_distance.cc
new file mode 100644
index 00000000000..54006062374
--- /dev/null
+++ b/source/blender/nodes/geometry/nodes/node_geo_merge_by_distance.cc
@@ -0,0 +1,80 @@
+/*
+ * 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 "GEO_point_merge_by_distance.hh"
+
+#include "node_geometry_util.hh"
+
+namespace blender::nodes::node_geo_merge_by_distance_cc {
+
+static void node_declare(NodeDeclarationBuilder &b)
+{
+  b.add_input<decl::Geometry>(N_("Geometry")).supported_type({GEO_COMPONENT_TYPE_POINT_CLOUD});
+  b.add_input<decl::Bool>(N_("Selection")).default_value(true).hide_value().supports_field();
+  b.add_input<decl::Float>(N_("Distance")).default_value(0.1f).min(0.0f).subtype(PROP_DISTANCE);
+  b.add_output<decl::Geometry>(N_("Geometry"));
+}
+
+static PointCloud *pointcloud_merge_by_distance(const PointCloudComponent &src_points,
+                                                const float merge_distance,
+                                                const Field<bool> &selection_field)
+{
+  const int src_size = src_points.attribute_domain_size(ATTR_DOMAIN_POINT);
+  GeometryComponentFieldContext context{src_points, ATTR_DOMAIN_POINT};
+  FieldEvaluator evaluator{context, src_size};
+  evaluator.add(selection_field);
+  evaluator.evaluate();
+
+  const IndexMask selection = evaluator.get_evaluated_as_mask(0);
+  if (selection.is_empty()) {
+    return nullptr;
+  }
+
+  return geometry::point_merge_by_distance(src_points, merge_distance, selection);
+}
+
+static void node_geo_exec(GeoNodeExecParams params)
+{
+  GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
+
+  const Field<bool> selection = params.extract_input<Field<bool>>("Selection");
+  const float merge_distance = params.extract_input<float>("Distance");
+
+  geometry_set.modify_geometry_sets([&](GeometrySet &geometry_set) {
+    if (geometry_set.has_pointcloud()) {
+      PointCloud *result = pointcloud_merge_by_distance(
+          *geometry_set.get_component_for_read<PointCloudComponent>(), merge_distance, selection);
+      geometry_set.replace_pointcloud(result);
+    }
+  });
+
+  params.set_output("Geometry", std::move(geometry_set));
+}
+
+}  // namespace blender::nodes::node_geo_merge_by_distance_cc
+
+void register_node_type_geo_merge_by_distance()
+{
+  namespace file_ns = blender::nodes::node_geo_merge_by_distance_cc;
+
+  static bNodeType ntype;
+
+  geo_node_type_base(&ntype, GEO_NODE_MERGE_BY_DISTANCE, "Merge by Distance", NODE_CLASS_GEOMETRY);
+
+  ntype.declare = file_ns::node_declare;
+  ntype.geometry_node_execute = file_ns::node_geo_exec;
+  nodeRegisterType(&ntype);
+}