scripts/startup/bl_ui/node_add_menu_geometry.py

@@ -220,6 +220,7 @@ class NODE_MT_geometry_node_GEO_GEOMETRY_SAMPLE(Menu):
     def draw(self, _context):
         layout = self.layout
         node_add_menu.add_node_type(layout, "GeometryNodeProximity")
+        node_add_menu.add_node_type(layout, "GeometryNodeIndexOfNearest")
         node_add_menu.add_node_type(layout, "GeometryNodeRaycast")
         node_add_menu.add_node_type(layout, "GeometryNodeSampleIndex")
         node_add_menu.add_node_type(layout, "GeometryNodeSampleNearest")

source/blender/blenkernel/BKE_node.h

@@ -1581,6 +1581,7 @@ void BKE_nodetree_remove_layer_n(struct bNodeTree *ntree, struct Scene *scene, i
 #define GEO_NODE_SDF_VOLUME_SPHERE 1196
 #define GEO_NODE_MEAN_FILTER_SDF_VOLUME 1197
 #define GEO_NODE_OFFSET_SDF_VOLUME 1198
+#define GEO_NODE_INDEX_OF_NEAREST 1199
 
 /** \} */
 

source/blender/blenlib/BLI_kdtree_impl.h

@@ -105,6 +105,23 @@ inline void BLI_kdtree_nd_(range_search_cb_cpp)(const KDTree *tree,
       },
       const_cast<Fn *>(&fn));
 }
+
+template<typename Fn>
+inline int BLI_kdtree_nd_(find_nearest_cb_cpp)(const KDTree *tree,
+                                               const float co[KD_DIMS],
+                                               KDTreeNearest *r_nearest,
+                                               Fn &&fn)
+{
+  return BLI_kdtree_nd_(find_nearest_cb)(
+      tree,
+      co,
+      [](void *user_data, const int index, const float *co, const float dist_sq) {
+        Fn &fn = *static_cast<Fn *>(user_data);
+        return fn(index, co, dist_sq);
+      },
+      &fn,
+      r_nearest);
+}
 #endif
 
 #undef _BLI_CONCAT_AUX

source/blender/nodes/NOD_static_types.h

@@ -324,6 +324,7 @@ DefNode(GeometryNode, GEO_NODE_FLIP_FACES, 0, "FLIP_FACES", FlipFaces, "Flip Fac
 DefNode(GeometryNode, GEO_NODE_GEOMETRY_TO_INSTANCE, 0, "GEOMETRY_TO_INSTANCE", GeometryToInstance, "Geometry to Instance", "Convert each input geometry into an instance, which can be much faster than the Join Geometry node when the inputs are large")
 DefNode(GeometryNode, GEO_NODE_IMAGE_INFO, 0, "IMAGE_INFO", ImageInfo, "Image Info", "Retrieve information about an image")
 DefNode(GeometryNode, GEO_NODE_IMAGE_TEXTURE, def_geo_image_texture, "IMAGE_TEXTURE", ImageTexture, "Image Texture", "Sample values from an image texture")
+DefNode(GeometryNode, GEO_NODE_INDEX_OF_NEAREST, 0, "INDEX_OF_NEAREST", IndexOfNearest, "Index of Nearest", "Find the nearest element in the a group. Similar to the \"Sample Nearest\" node")
 DefNode(GeometryNode, GEO_NODE_IMAGE, def_geo_image, "IMAGE", InputImage, "Image", "Input image")
 DefNode(GeometryNode, GEO_NODE_INPUT_CURVE_HANDLES, 0, "INPUT_CURVE_HANDLES", InputCurveHandlePositions,"Curve Handle Positions", "Retrieve the position of each Bézier control point's handles")
 DefNode(GeometryNode, GEO_NODE_INPUT_CURVE_TILT, 0, "INPUT_CURVE_TILT", InputCurveTilt, "Curve Tilt", "Retrieve the angle at each control point used to twist the curve's normal around its tangent")
@@ -398,7 +399,7 @@ DefNode(GeometryNode, GEO_NODE_ROTATE_INSTANCES, 0, "ROTATE_INSTANCES", RotateIn
 DefNode(GeometryNode, GEO_NODE_SAMPLE_CURVE, def_geo_curve_sample, "SAMPLE_CURVE", SampleCurve, "Sample Curve", "Retrieve data from a point on a curve at a certain distance from its start")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_INDEX, def_geo_sample_index, "SAMPLE_INDEX", SampleIndex, "Sample Index", "Retrieve values from specific geometry elements")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_NEAREST_SURFACE, def_geo_sample_nearest_surface, "SAMPLE_NEAREST_SURFACE", SampleNearestSurface, "Sample Nearest Surface", "Calculate the interpolated value of a mesh attribute on the closest point of its surface")
-DefNode(GeometryNode, GEO_NODE_SAMPLE_NEAREST, def_geo_sample_nearest, "SAMPLE_NEAREST", SampleNearest, "Sample Nearest", "Find the element of a geometry closest to a position")
+DefNode(GeometryNode, GEO_NODE_SAMPLE_NEAREST, def_geo_sample_nearest, "SAMPLE_NEAREST", SampleNearest, "Sample Nearest", "Find the element of a geometry closest to a position. Similar to the \"Index of Nearest\" node")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_UV_SURFACE, def_geo_sample_uv_surface, "SAMPLE_UV_SURFACE", SampleUVSurface, "Sample UV Surface", "Calculate the interpolated values of a mesh attribute at a UV coordinate")
 DefNode(GeometryNode, GEO_NODE_SCALE_ELEMENTS, def_geo_scale_elements, "SCALE_ELEMENTS", ScaleElements, "Scale Elements", "Scale groups of connected edges and faces")
 DefNode(GeometryNode, GEO_NODE_SCALE_INSTANCES, 0, "SCALE_INSTANCES", ScaleInstances, "Scale Instances", "Scale geometry instances in local or global space")

source/blender/nodes/geometry/CMakeLists.txt

@@ -78,6 +78,7 @@ set(SRC
   nodes/node_geo_image.cc
   nodes/node_geo_image_info.cc
   nodes/node_geo_image_texture.cc
+  nodes/node_geo_index_of_nearest.cc
   nodes/node_geo_input_curve_handles.cc
   nodes/node_geo_input_curve_tilt.cc
   nodes/node_geo_input_id.cc

source/blender/nodes/geometry/node_geometry_register.cc

@@ -62,6 +62,7 @@ void register_geometry_nodes()
   register_node_type_geo_image_info();
   register_node_type_geo_image_texture();
   register_node_type_geo_image();
+  register_node_type_geo_index_of_nearest();
   register_node_type_geo_input_curve_handles();
   register_node_type_geo_input_curve_tilt();
   register_node_type_geo_input_id();

source/blender/nodes/geometry/node_geometry_register.hh

@@ -59,6 +59,7 @@ void register_node_type_geo_geometry_to_instance();
 void register_node_type_geo_image_info();
 void register_node_type_geo_image_texture();
 void register_node_type_geo_image();
+void register_node_type_geo_index_of_nearest();
 void register_node_type_geo_input_curve_handles();
 void register_node_type_geo_input_curve_tilt();
 void register_node_type_geo_input_id();

source/blender/nodes/geometry/nodes/node_geo_index_of_nearest.cc

@@ -0,0 +1,270 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+
+#include "BLI_array.hh"
+#include "BLI_kdtree.h"
+#include "BLI_map.hh"
+#include "BLI_task.hh"
+
+#include "node_geometry_util.hh"
+
+namespace blender::nodes::node_geo_index_of_nearest_cc {
+
+static void node_declare(NodeDeclarationBuilder &b)
+{
+  b.add_input<decl::Vector>(N_("Position")).implicit_field(implicit_field_inputs::position);
+  b.add_input<decl::Int>(N_("Group ID")).supports_field().hide_value();
+
+  b.add_output<decl::Int>(N_("Index")).field_source().description(N_("Index of nearest element"));
+  b.add_output<decl::Bool>(N_("Has Neighbor")).field_source();
+}
+
+static KDTree_3d *build_kdtree(const Span<float3> positions, const IndexMask mask)
+{
+  KDTree_3d *tree = BLI_kdtree_3d_new(mask.size());
+  for (const int index : mask) {
+    BLI_kdtree_3d_insert(tree, index, positions[index]);
+  }
+  BLI_kdtree_3d_balance(tree);
+  return tree;
+}
+
+static int find_nearest_non_self(const KDTree_3d &tree, const float3 &position, const int index)
+{
+  return BLI_kdtree_3d_find_nearest_cb_cpp(
+      &tree, position, 0, [index](const int other, const float * /*co*/, const float /*dist_sq*/) {
+        return index == other ? 0 : 1;
+      });
+}
+
+static void find_neighbors(const KDTree_3d &tree,
+                           const Span<float3> positions,
+                           const IndexMask mask,
+                           MutableSpan<int> r_indices)
+{
+  threading::parallel_for(mask.index_range(), 512, [&](const IndexRange range) {
+    for (const int index : mask.slice(range)) {
+      r_indices[index] = find_nearest_non_self(tree, positions[index], index);
+    }
+  });
+}
+
+class IndexOfNearestFieldInput final : public bke::GeometryFieldInput {
+ private:
+  const Field<float3> positions_field_;
+  const Field<int> group_field_;
+
+ public:
+  IndexOfNearestFieldInput(Field<float3> positions_field, Field<int> group_field)
+      : bke::GeometryFieldInput(CPPType::get<int>(), "Index of Nearest"),
+        positions_field_(std::move(positions_field)),
+        group_field_(std::move(group_field))
+  {
+  }
+
+  GVArray get_varray_for_context(const bke::GeometryFieldContext &context,
+                                 const IndexMask mask) const final
+  {
+    if (!context.attributes()) {
+      return {};
+    }
+    const int domain_size = context.attributes()->domain_size(context.domain());
+    fn::FieldEvaluator evaluator{context, domain_size};
+    evaluator.add(positions_field_);
+    evaluator.add(group_field_);
+    evaluator.evaluate();
+    const VArraySpan<float3> positions = evaluator.get_evaluated<float3>(0);
+    const VArray<int> group = evaluator.get_evaluated<int>(1);
+
+    Array<int> result;
+
+    if (group.is_single()) {
+      result.reinitialize(mask.min_array_size());
+      KDTree_3d *tree = build_kdtree(positions, IndexRange(domain_size));
+      find_neighbors(*tree, positions, mask, result);
+      BLI_kdtree_3d_free(tree);
+      return VArray<int>::ForContainer(std::move(result));
+    }
+
+    VectorSet<int> group_indexing;
+    for (const int index : mask) {
+      const int group_id = group[index];
+      group_indexing.add(group_id);
+    }
+
+    /* Each group id has two corresponding index masks. One that contains all the points in the
+     * group, one that contains all the points in the group that should be looked up (this is the
+     * intersection of the points in the group and `mask`). In many cases, both of these masks are
+     * the same or very similar, so there is no benefit two separate masks. */
+    const bool use_separate_lookup_indices = mask.size() < domain_size / 2;
+
+    Array<Vector<int64_t>> all_indices_by_group_id(group_indexing.size());
+    Array<Vector<int64_t>> lookup_indices_by_group_id;
+
+    if (use_separate_lookup_indices) {
+      result.reinitialize(mask.min_array_size());
+      lookup_indices_by_group_id.reinitialize(group_indexing.size());
+    }
+    else {
+      result.reinitialize(domain_size);
+    }
+
+    const auto build_group_masks = [&](const IndexMask mask,
+                                       MutableSpan<Vector<int64_t>> r_groups) {
+      for (const int index : mask) {
+        const int group_id = group[index];
+        const int index_of_group = group_indexing.index_of_try(group_id);
+        if (index_of_group != -1) {
+          r_groups[index_of_group].append(index);
+        }
+      }
+    };
+
+    threading::parallel_invoke(
+        domain_size > 1024 && use_separate_lookup_indices,
+        [&]() {
+          if (use_separate_lookup_indices) {
+            build_group_masks(mask, lookup_indices_by_group_id);
+          }
+        },
+        [&]() { build_group_masks(IndexMask(domain_size), all_indices_by_group_id); });
+
+    threading::parallel_for(group_indexing.index_range(), 256, [&](const IndexRange range) {
+      for (const int index : range) {
+        const IndexMask tree_mask = all_indices_by_group_id[index].as_span();
+        const IndexMask lookup_mask = use_separate_lookup_indices ?
+                                          IndexMask(lookup_indices_by_group_id[index]) :
+                                          tree_mask;
+        KDTree_3d *tree = build_kdtree(positions, tree_mask);
+        find_neighbors(*tree, positions, lookup_mask, result);
+        BLI_kdtree_3d_free(tree);
+      }
+    });
+
+    return VArray<int>::ForContainer(std::move(result));
+  }
+
+ public:
+  void for_each_field_input_recursive(FunctionRef<void(const FieldInput &)> fn) const
+  {
+    positions_field_.node().for_each_field_input_recursive(fn);
+    group_field_.node().for_each_field_input_recursive(fn);
+  }
+
+  uint64_t hash() const final
+  {
+    return get_default_hash_2(positions_field_, group_field_);
+  }
+
+  bool is_equal_to(const fn::FieldNode &other) const final
+  {
+    if (const auto *other_field = dynamic_cast<const IndexOfNearestFieldInput *>(&other)) {
+      return positions_field_ == other_field->positions_field_ &&
+             group_field_ == other_field->group_field_;
+    }
+    return false;
+  }
+
+  std::optional<eAttrDomain> preferred_domain(const GeometryComponent &component) const final
+  {
+    return bke::try_detect_field_domain(component, positions_field_);
+  }
+};
+
+class HasNeighborFieldInput final : public bke::GeometryFieldInput {
+ private:
+  const Field<int> group_field_;
+
+ public:
+  HasNeighborFieldInput(Field<int> group_field)
+      : bke::GeometryFieldInput(CPPType::get<bool>(), "Has Neighbor"),
+        group_field_(std::move(group_field))
+  {
+  }
+
+  GVArray get_varray_for_context(const bke::GeometryFieldContext &context,
+                                 const IndexMask mask) const final
+  {
+    if (!context.attributes()) {
+      return {};
+    }
+    const int domain_size = context.attributes()->domain_size(context.domain());
+    if (domain_size == 1) {
+      return VArray<bool>::ForSingle(false, mask.min_array_size());
+    }
+
+    fn::FieldEvaluator evaluator{context, domain_size};
+    evaluator.add(group_field_);
+    evaluator.evaluate();
+    const VArray<int> group = evaluator.get_evaluated<int>(0);
+
+    if (group.is_single()) {
+      return VArray<bool>::ForSingle(true, mask.min_array_size());
+    }
+
+    Map<int, int> counts;
+    const VArraySpan<int> group_span(group);
+    mask.foreach_index([&](const int i) {
+      counts.add_or_modify(
+          group_span[i], [](int *count) { *count = 0; }, [](int *count) { (*count)++; });
+    });
+    Array<bool> result(mask.min_array_size());
+    mask.foreach_index([&](const int i) { result[i] = counts.lookup(group_span[i]) > 1; });
+    return VArray<bool>::ForContainer(std::move(result));
+  }
+
+ public:
+  void for_each_field_input_recursive(FunctionRef<void(const FieldInput &)> fn) const
+  {
+    group_field_.node().for_each_field_input_recursive(fn);
+  }
+
+  uint64_t hash() const final
+  {
+    return get_default_hash_2(3984756934876, group_field_);
+  }
+
+  bool is_equal_to(const fn::FieldNode &other) const final
+  {
+    if (const auto *other_field = dynamic_cast<const HasNeighborFieldInput *>(&other)) {
+      return group_field_ == other_field->group_field_;
+    }
+    return false;
+  }
+
+  std::optional<eAttrDomain> preferred_domain(const GeometryComponent &component) const final
+  {
+    return bke::try_detect_field_domain(component, group_field_);
+  }
+};
+
+static void node_geo_exec(GeoNodeExecParams params)
+{
+  Field<float3> position_field = params.extract_input<Field<float3>>("Position");
+  Field<int> group_field = params.extract_input<Field<int>>("Group ID");
+
+  if (params.output_is_required("Index")) {
+    params.set_output("Index",
+                      Field<int>(std::make_shared<IndexOfNearestFieldInput>(
+                          std::move(position_field), group_field)));
+  }
+
+  if (params.output_is_required("Has Neighbor")) {
+    params.set_output(
+        "Has Neighbor",
+        Field<bool>(std::make_shared<HasNeighborFieldInput>(std::move(group_field))));
+  }
+}
+
+}  // namespace blender::nodes::node_geo_index_of_nearest_cc
+
+void register_node_type_geo_index_of_nearest()
+{
+  namespace file_ns = blender::nodes::node_geo_index_of_nearest_cc;
+
+  static bNodeType ntype;
+
+  geo_node_type_base(&ntype, GEO_NODE_INDEX_OF_NEAREST, "Index of Nearest", NODE_CLASS_CONVERTER);
+  ntype.geometry_node_execute = file_ns::node_geo_exec;
+  ntype.declare = file_ns::node_declare;
+  nodeRegisterType(&ntype);
+}