source/blender/blenlib/BLI_index_mask.hh

@@ -292,6 +292,22 @@ class IndexMask : private IndexMaskData {
    */
   IndexMaskSegment segment(int64_t segment_i) const;
 
+  /**
+   * Iterate over the indices in multiple masks which have the same size. The given function is
+   * called for groups of segments where each segment has the same size and comes from a different
+   * input mask.
+   * For example, if the input masks are (both have size 18):
+   *   A: [0, 15), {20, 24, 25}
+   *   B: [0, 5), [10, 15], {20, 30, 40, 50, 60, 70, 80, 90}
+   * Then the function will be called multiple times, each time with two segments:
+   *   1. [0, 5), [0, 5)
+   *   2. [5, 10), [10, 15)
+   *   3. [10, 15), {20, 30, 40, 50, 60}
+   *   4. {20, 24, 25}, {70, 80, 90}
+   */
+  static void foreach_segment_zipped(Span<IndexMask> masks,
+                                     FunctionRef<bool(Span<IndexMaskSegment> segments)> fn);
+
   /**
    * Calls the function once for every index.
    *
@@ -389,6 +405,9 @@ class IndexMask : private IndexMaskData {
    * Is used by some functions to get low level access to the mask in order to construct it.
    */
   IndexMaskData &data_for_inplace_construction();
+
+  friend bool operator==(const IndexMask &a, const IndexMask &b);
+  friend bool operator!=(const IndexMask &a, const IndexMask &b);
 };
 
 /**
@@ -926,6 +945,11 @@ inline Vector<std::variant<IndexRange, IndexMaskSegment>, N> IndexMask::to_spans
   return segments;
 }
 
+inline bool operator!=(const IndexMask &a, const IndexMask &b)
+{
+  return !(a == b);
+}
+
 }  // namespace blender::index_mask
 
 namespace blender {

source/blender/blenlib/intern/index_mask.cc

@@ -868,4 +868,111 @@ template IndexMask IndexMask::from_indices(Span<int64_t>, IndexMaskMemory &);
 template void IndexMask::to_indices(MutableSpan<int32_t>) const;
 template void IndexMask::to_indices(MutableSpan<int64_t>) const;
 
+void IndexMask::foreach_segment_zipped(const Span<IndexMask> masks,
+                                       const FunctionRef<bool(Span<IndexMaskSegment> segments)> fn)
+{
+  BLI_assert(!masks.is_empty());
+  BLI_assert(std::all_of(masks.begin() + 1, masks.end(), [&](const IndexMask &maks) {
+    return masks[0].size() == maks.size();
+  }));
+
+  Array<int64_t, 8> segment_iter(masks.size(), 0);
+  Array<int16_t, 8> start_iter(masks.size(), 0);
+
+  Array<IndexMaskSegment, 8> segments(masks.size());
+  Array<IndexMaskSegment, 8> sequences(masks.size());
+
+  /* This function only take positions of indices in to account.
+   * Masks with the same size is fragmented in positions space.
+   * So, all last segments (index in mask does not matter) of all masks will be ended in the same
+   * position. All segment iterators will be out of range at the same time. */
+  while (segment_iter[0] != masks[0].segments_num()) {
+    for (const int64_t mask_i : masks.index_range()) {
+      if (start_iter[mask_i] == 0) {
+        segments[mask_i] = masks[mask_i].segment(segment_iter[mask_i]);
+      }
+    }
+
+    int16_t next_common_sequence_size = std::numeric_limits<int16_t>::max();
+    for (const int64_t mask_i : masks.index_range()) {
+      next_common_sequence_size = math::min(next_common_sequence_size,
+                                            int16_t(segments[mask_i].size() - start_iter[mask_i]));
+    }
+
+    for (const int64_t mask_i : masks.index_range()) {
+      sequences[mask_i] = segments[mask_i].slice(start_iter[mask_i], next_common_sequence_size);
+    }
+
+    if (!fn(sequences)) {
+      break;
+    }
+
+    for (const int64_t mask_i : masks.index_range()) {
+      if (segments[mask_i].size() - start_iter[mask_i] == next_common_sequence_size) {
+        segment_iter[mask_i]++;
+        start_iter[mask_i] = 0;
+      }
+      else {
+        start_iter[mask_i] += next_common_sequence_size;
+      }
+    }
+  }
+}
+
+static bool segments_is_equal(const IndexMaskSegment &a, const IndexMaskSegment &b)
+{
+  if (a.size() != b.size()) {
+    return false;
+  }
+  if (a.is_empty()) {
+    /* Both segments are empty. */
+    return true;
+  }
+  if (a[0] != b[0]) {
+    return false;
+  }
+
+  const bool a_is_range = unique_sorted_indices::non_empty_is_range(a.base_span());
+  const bool b_is_range = unique_sorted_indices::non_empty_is_range(b.base_span());
+  if (a_is_range || b_is_range) {
+    return a_is_range && b_is_range;
+  }
+
+  const Span<int16_t> a_indices = a.base_span();
+  [[maybe_unused]] const Span<int16_t> b_indices = b.base_span();
+
+  const int64_t offset_difference = int16_t(b.offset() - a.offset());
+
+  BLI_assert(a_indices[0] >= 0 && b_indices[0] >= 0);
+  BLI_assert(b_indices[0] == a_indices[0] - offset_difference);
+
+  return std::equal(a_indices.begin(),
+                    a_indices.end(),
+                    b.base_span().begin(),
+                    [offset_difference](const int16_t a_index, const int16_t b_index) -> bool {
+                      return a_index - offset_difference == b_index;
+                    });
+}
+
+bool operator==(const IndexMask &a, const IndexMask &b)
+{
+  if (a.size() != b.size()) {
+    return false;
+  }
+
+  const std::optional<IndexRange> a_as_range = a.to_range();
+  const std::optional<IndexRange> b_as_range = b.to_range();
+  if (a_as_range.has_value() || b_as_range.has_value()) {
+    return a_as_range == b_as_range;
+  }
+
+  bool equals = true;
+  IndexMask::foreach_segment_zipped({a, b}, [&](const Span<IndexMaskSegment> segments) {
+    equals &= segments_is_equal(segments[0], segments[1]);
+    return equals;
+  });
+
+  return equals;
+}
+
 }  // namespace blender::index_mask

source/blender/blenlib/tests/BLI_index_mask_test.cc

@@ -453,4 +453,218 @@ TEST(index_mask, SliceContent)
   }
 }
 
+TEST(index_mask, EqualsRangeSelf)
+{
+  IndexMask mask = IndexRange(16384);
+  EXPECT_EQ(mask, mask);
+}
+
+TEST(index_mask, EqualsRange)
+{
+  IndexMask mask_a = IndexRange(16384);
+  IndexMask mask_b = IndexRange(16384);
+  EXPECT_EQ(mask_a, mask_b);
+}
+
+TEST(index_mask, EqualsRangeLarge)
+{
+  IndexMask mask_a = IndexRange(96384);
+  IndexMask mask_b = IndexRange(96384);
+  EXPECT_EQ(mask_a, mask_b);
+}
+
+TEST(index_mask, EqualsRangeBegin)
+{
+  IndexMask mask_a = IndexRange(102, 16384 - 102);
+  IndexMask mask_b = IndexRange(102, 16384 - 102);
+  EXPECT_EQ(mask_a, mask_b);
+}
+
+TEST(index_mask, EqualsRangeEnd)
+{
+  IndexMask mask_a = IndexRange(16384 + 1);
+  IndexMask mask_b = IndexRange(16384 + 1);
+  EXPECT_EQ(mask_a, mask_b);
+}
+
+TEST(index_mask, NonEqualsRange)
+{
+  IndexMask mask_a = IndexRange(16384);
+  IndexMask mask_b = IndexRange(1, 16384);
+  EXPECT_NE(mask_a, mask_b);
+}
+
+TEST(index_mask, EqualsSelf)
+{
+  IndexMaskMemory memory;
+  IndexMask mask = IndexMask::from_union(IndexRange(16384), IndexRange(16384 * 3, 533), memory);
+  EXPECT_EQ(mask, mask);
+}
+
+TEST(index_mask, Equals)
+{
+  IndexMaskMemory memory;
+  IndexMask mask_a = IndexMask::from_union(IndexRange(16384), IndexRange(16384 * 3, 533), memory);
+  IndexMask mask_b = IndexMask::from_union(IndexRange(16384), IndexRange(16384 * 3, 533), memory);
+  EXPECT_EQ(mask_a, mask_b);
+}
+
+TEST(index_mask, NonEquals)
+{
+  IndexMaskMemory memory;
+  IndexMask mask_a = IndexMask::from_union(IndexRange(16384), IndexRange(16384 * 3, 533), memory);
+  IndexMask mask_b = IndexMask::from_union(
+      IndexRange(55, 16384), IndexRange(16384 * 5, 533), memory);
+  EXPECT_NE(mask_a, mask_b);
+}
+
+TEST(index_mask, NotEqualsRangeAndIndices)
+{
+  IndexMaskMemory memory;
+  IndexMask mask_a = IndexMask::from_union(
+      IndexRange(2040), IndexMask::from_indices<int>({2072, 2073, 2075}, memory), memory);
+  IndexMask mask_b = IndexMask::from_union(
+      IndexRange(2040), IndexMask::from_indices<int>({2072, 2073 + 1, 2075}, memory), memory);
+
+  EXPECT_NE(mask_a, mask_b);
+}
+
+static bool mask_segments_equals(const IndexMaskSegment &a, const IndexMaskSegment &b)
+{
+  if (a.size() != b.size()) {
+    return false;
+  }
+  for (const int64_t i : a.index_range()) {
+    if (a[i] != b[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+TEST(index_mask, ZippedForeachSelf)
+{
+  IndexMaskMemory memory;
+  IndexMask mask = IndexMask::from_initializers({IndexRange(500), 555, 699, 222, 900, 100},
+                                                memory);
+  {
+    int calls_num = 0;
+    IndexMask::foreach_segment_zipped({mask}, [&](Span<IndexMaskSegment> segments) {
+      EXPECT_FALSE(segments.is_empty());
+      calls_num++;
+      return true;
+    });
+    EXPECT_EQ(calls_num, 2);
+  }
+
+  {
+    int calls_num = 0;
+    IndexMask::foreach_segment_zipped({mask, mask}, [&](Span<IndexMaskSegment> segments) {
+      EXPECT_FALSE(segments.is_empty());
+      EXPECT_TRUE(mask_segments_equals(segments[0], segments[1]));
+      calls_num++;
+      return true;
+    });
+    EXPECT_EQ(calls_num, 2);
+  }
+
+  {
+    int calls_num = 0;
+    IndexMask::foreach_segment_zipped({mask, mask, mask}, [&](Span<IndexMaskSegment> segments) {
+      EXPECT_FALSE(segments.is_empty());
+      EXPECT_TRUE(mask_segments_equals(segments[0], segments[1]));
+      EXPECT_TRUE(mask_segments_equals(segments[0], segments[2]));
+      calls_num++;
+      return true;
+    });
+    EXPECT_EQ(calls_num, 2);
+  }
+
+  {
+    int calls_num = 0;
+    IndexMask::foreach_segment_zipped(
+        {mask, mask, mask, mask}, [&](Span<IndexMaskSegment> segments) {
+          EXPECT_FALSE(segments.is_empty());
+          EXPECT_TRUE(mask_segments_equals(segments[0], segments[1]));
+          EXPECT_TRUE(mask_segments_equals(segments[0], segments[2]));
+          EXPECT_TRUE(mask_segments_equals(segments[0], segments[3]));
+          calls_num++;
+          return true;
+        });
+    EXPECT_EQ(calls_num, 2);
+  }
+}
+
+TEST(index_mask, ZippedForeachSameSegments)
+{
+  IndexMaskMemory memory;
+  IndexMask mask_a = IndexMask::from_initializers({0, 1, 2}, memory);
+  IndexMask mask_b = IndexMask::from_initializers({3, 4, 5}, memory);
+  IndexMask mask_c = IndexMask::from_initializers({6, 7, 8}, memory);
+  {
+    int calls_num = 0;
+    IndexMask::foreach_segment_zipped({mask_a}, [&](Span<IndexMaskSegment> segments) {
+      EXPECT_FALSE(segments.is_empty());
+      calls_num++;
+      return true;
+    });
+    EXPECT_EQ(calls_num, 1);
+  }
+  {
+    int calls_num = 0;
+    IndexMask::foreach_segment_zipped({mask_a, mask_b}, [&](Span<IndexMaskSegment> segments) {
+      EXPECT_FALSE(segments.is_empty());
+      EXPECT_EQ(segments[0].size(), segments[1].size());
+      EXPECT_FALSE(mask_segments_equals(segments[0], segments[1]));
+      calls_num++;
+      return true;
+    });
+    EXPECT_EQ(calls_num, 1);
+  }
+  {
+    int calls_num = 0;
+    IndexMask::foreach_segment_zipped(
+        {mask_a, mask_b, mask_c}, [&](Span<IndexMaskSegment> segments) {
+          EXPECT_FALSE(segments.is_empty());
+          EXPECT_EQ(segments[0].size(), segments[1].size());
+          EXPECT_EQ(segments[0].size(), segments[2].size());
+          EXPECT_FALSE(mask_segments_equals(segments[0], segments[1]));
+          EXPECT_FALSE(mask_segments_equals(segments[0], segments[2]));
+          EXPECT_FALSE(mask_segments_equals(segments[1], segments[2]));
+          calls_num++;
+          return true;
+        });
+    EXPECT_EQ(calls_num, 1);
+  }
+}
+
+TEST(index_mask, ZippedForeachEqual)
+{
+  Span<int16_t> indices(get_static_indices_array());
+
+  IndexMaskMemory memory;
+  IndexMask mask_a = IndexMask::from_segments(
+      {{0, indices.take_front(5)}, {5, indices.take_front(5)}}, memory);
+  IndexMask mask_b = IndexMask::from_segments(
+      {{0, indices.take_front(3)}, {3, indices.take_front(4)}, {7, indices.take_front(3)}},
+      memory);
+  IndexMask mask_c = IndexMask::from_segments({{0, indices.take_front(10)}}, memory);
+
+  int index = 0;
+  Array<IndexMaskSegment> reference_segments{{0, indices.take_front(3)},
+                                             {3, indices.take_front(2)},
+                                             {5, indices.take_front(2)},
+                                             {7, indices.take_front(3)}};
+
+  IndexMask::foreach_segment_zipped(
+      {mask_a, mask_b, mask_c}, [&](Span<IndexMaskSegment> segments) {
+        EXPECT_TRUE(mask_segments_equals(reference_segments[index], segments[0]));
+        EXPECT_TRUE(mask_segments_equals(reference_segments[index], segments[1]));
+        EXPECT_TRUE(mask_segments_equals(reference_segments[index], segments[2]));
+        index++;
+        return true;
+      });
+  EXPECT_EQ(index, 4);
+}
+
 }  // namespace blender::index_mask::tests