source/blender/blenlib/BLI_index_mask.hh

@@ -73,8 +73,8 @@ class IndexMask {
   int64_t first() const;
   int64_t last() const;
   int64_t min_array_size() const;
-  std::optional<RawMaskIterator> find(const int64_t index) const;
-  bool contains(const int64_t index) const;
+  std::optional<RawMaskIterator> find(const int64_t query_index) const;
+  bool contains(const int64_t query_index) const;
 
   RawMaskIterator index_to_iterator(const int64_t index) const;
   int64_t iterator_to_index(const RawMaskIterator &it) const;

source/blender/blenlib/intern/index_mask.cc

@@ -374,32 +374,42 @@ IndexMask IndexMask::from_predicate_impl(
   return mask;
 }
 
-std::optional<RawMaskIterator> IndexMask::find(const int64_t index) const
+std::optional<RawMaskIterator> IndexMask::find(const int64_t query_index) const
 {
   if (this->is_empty()) {
     return std::nullopt;
   }
-  if (index < this->first()) {
+  if (query_index < this->first()) {
     return std::nullopt;
   }
-  if (index > this->last()) {
+  if (query_index > this->last()) {
     return std::nullopt;
   }
 
-  for (const int64_t segment_i : IndexRange(data_.segments_num)) {
-    const OffsetSpan<int64_t, int16_t> segment = this->segment(segment_i);
-    for (const int64_t index_in_segment : segment.index_range()) {
-      if (segment[index_in_segment] == index) {
-        return RawMaskIterator{segment_i, int16_t(index_in_segment)};
-      }
-    }
+  const int64_t segment_i = std::upper_bound(data_.segment_offsets,
+                                             data_.segment_offsets + data_.segments_num,
+                                             query_index) -
+                            data_.segment_offsets - 1;
+  const OffsetSpan<int64_t, int16_t> segment = this->segment(segment_i);
+  const Span<int16_t> local_segment = segment.base_span();
+  const int64_t local_query_index = query_index - segment.offset();
+  if (local_query_index > local_segment.last()) {
+    return std::nullopt;
   }
-  return std::nullopt;
+  const int64_t index_in_segment = std::lower_bound(local_segment.begin(),
+                                                    local_segment.end(),
+                                                    local_query_index) -
+                                   local_segment.begin();
+  BLI_assert(index_in_segment < local_segment.size());
+  if (local_segment[index_in_segment] != local_query_index) {
+    return std::nullopt;
+  }
+  return RawMaskIterator{segment_i, int16_t(index_in_segment)};
 }
 
-bool IndexMask::contains(const int64_t index) const
+bool IndexMask::contains(const int64_t query_index) const
 {
-  return this->find(index).has_value();
+  return this->find(query_index).has_value();
 }
 
 template IndexMask IndexMask::from_indices(Span<int32_t>, IndexMaskMemory &);

source/blender/blenlib/tests/BLI_index_mask_test.cc

@@ -255,6 +255,19 @@ TEST(index_mask, IndexIteratorConversionFuzzy)
     const int64_t new_index = sub_mask.iterator_to_index(it);
     EXPECT_EQ(index, new_index);
   }
+
+  for ([[maybe_unused]] const int64_t _ : IndexRange(100)) {
+    const int64_t index = rng.get_int32(int(indices.size() - 1000));
+    for (const int64_t offset : {0, 1, 2, 100, 500}) {
+      const int64_t index_to_search = indices[index] + offset;
+      const bool contained = std::binary_search(indices.begin(), indices.end(), index_to_search);
+      const std::optional<RawMaskIterator> it = mask.find(index_to_search);
+      EXPECT_EQ(contained, it.has_value());
+      if (contained) {
+        EXPECT_EQ(index_to_search, mask[*it]);
+      }
+    }
+  }
 }
 
 TEST(index_mask, FromPredicateFuzzy)