From 5063f7a725d03fcd8394f8c8fb5275abe75e62c0 Mon Sep 17 00:00:00 2001
From: Hans Goudey <h.goudey@me.com>
Date: Wed, 24 May 2023 20:05:29 -0400
Subject: [PATCH] Cleanup: Simplify IndexMask usage in two geometry nodes

Use the weird but at least consistent "from_groups" method in the sample
curves node, and avoid one logarithmic lookups when merging points.
---
 .../geometry/intern/point_merge_by_distance.cc   | 11 ++++-------
 .../geometry/nodes/node_geo_curve_sample.cc      | 16 +++++++++++-----
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/source/blender/geometry/intern/point_merge_by_distance.cc b/source/blender/geometry/intern/point_merge_by_distance.cc
index 6c1444ca4eb..3518643d7ef 100644
--- a/source/blender/geometry/intern/point_merge_by_distance.cc
+++ b/source/blender/geometry/intern/point_merge_by_distance.cc
@@ -47,18 +47,15 @@ PointCloud *point_merge_by_distance(const PointCloud &src_points,
    * 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;
-  }
+  std::iota(merge_indices.begin(), merge_indices.end(), 0);
 
-  for (const int i : selection_merge_indices.index_range()) {
-    const int merge_index = selection_merge_indices[i];
+  selection.foreach_index([&](const int src_index, const int pos) {
+    const int merge_index = selection_merge_indices[pos];
     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. */
diff --git a/source/blender/nodes/geometry/nodes/node_geo_curve_sample.cc b/source/blender/nodes/geometry/nodes/node_geo_curve_sample.cc
index 942e1e0025d..c031924833c 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_curve_sample.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_curve_sample.cc
@@ -395,12 +395,12 @@ class SampleCurveFunction : public mf::MultiFunction {
     }
     else {
       Vector<int> invalid_indices;
-      MultiValueMap<int, int> indices_per_curve;
+      VectorSet<int> used_curves;
       devirtualize_varray(curve_indices, [&](const auto curve_indices) {
         mask.foreach_index([&](const int i) {
           const int curve_i = curve_indices[i];
           if (curves.curves_range().contains(curve_i)) {
-            indices_per_curve.add(curve_i, i);
+            used_curves.add(curve_i);
           }
           else {
             invalid_indices.append(i);
@@ -409,9 +409,15 @@ class SampleCurveFunction : public mf::MultiFunction {
       });
 
       IndexMaskMemory memory;
-      for (const int curve_i : indices_per_curve.keys()) {
-        sample_curve(curve_i,
-                     IndexMask::from_indices<int>(indices_per_curve.lookup(curve_i), memory));
+      Array<IndexMask> mask_by_curve(used_curves.size());
+      IndexMask::from_groups<int>(
+          mask,
+          memory,
+          [&](const int i) { return used_curves.index_of(curve_indices[i]); },
+          mask_by_curve);
+
+      for (const int i : mask_by_curve.index_range()) {
+        sample_curve(used_curves[i], mask_by_curve[i]);
       }
       fill_invalid(IndexMask::from_indices<int>(invalid_indices, memory));
     }