BLI: support content based slicing in IndexMask

This allows slicing an `IndexMask` so that it only contains certain indices.

Pull Request: https://projects.blender.org/blender/blender/pulls/117857

source/blender/blenlib/BLI_index_mask.hh

@@ -260,6 +260,13 @@ class IndexMask : private IndexMaskData {
    */
   IndexMask slice(IndexRange range) const;
   IndexMask slice(int64_t start, int64_t size) const;
+  IndexMask slice(RawMaskIterator first_it, RawMaskIterator last_it, int64_t size) const;
+  /**
+   * Slices the mask based on the stored indices. The resulting mask only contains the indices that
+   * are within the given range.
+   */
+  IndexMask slice_content(IndexRange range) const;
+  IndexMask slice_content(int64_t start, int64_t size) const;
   /**
    * Same as above but can also add an offset to every index in the mask.
    * Takes O(log n + range.size()) time but with a very small constant factor.

source/blender/blenlib/intern/index_mask.cc

@@ -133,6 +133,46 @@ IndexMask IndexMask::slice(const int64_t start, const int64_t size) const
   return sliced;
 }
 
+IndexMask IndexMask::slice(const RawMaskIterator first_it,
+                           const RawMaskIterator last_it,
+                           const int64_t size) const
+{
+  BLI_assert(this->iterator_to_index(last_it) - this->iterator_to_index(first_it) + 1 == size);
+  IndexMask sliced = *this;
+  sliced.indices_num_ = size;
+  sliced.segments_num_ = last_it.segment_i - first_it.segment_i + 1;
+  sliced.indices_by_segment_ += first_it.segment_i;
+  sliced.segment_offsets_ += first_it.segment_i;
+  sliced.cumulative_segment_sizes_ += first_it.segment_i;
+  sliced.begin_index_in_segment_ = first_it.index_in_segment;
+  sliced.end_index_in_segment_ = last_it.index_in_segment + 1;
+  return sliced;
+}
+
+IndexMask IndexMask::slice_content(const IndexRange range) const
+{
+  return this->slice_content(range.start(), range.size());
+}
+
+IndexMask IndexMask::slice_content(const int64_t start, const int64_t size) const
+{
+  if (size <= 0) {
+    return {};
+  }
+  const std::optional<RawMaskIterator> first_it = this->find_larger_equal(start);
+  const std::optional<RawMaskIterator> last_it = this->find_smaller_equal(start + size - 1);
+  if (!first_it || !last_it) {
+    return {};
+  }
+  const int64_t first_index = this->iterator_to_index(*first_it);
+  const int64_t last_index = this->iterator_to_index(*last_it);
+  if (last_index < first_index) {
+    return {};
+  }
+  const int64_t sliced_mask_size = last_index - first_index + 1;
+  return this->slice(*first_it, *last_it, sliced_mask_size);
+}
+
 IndexMask IndexMask::slice_and_offset(const IndexRange range,
                                       const int64_t offset,
                                       IndexMaskMemory &memory) const

source/blender/blenlib/tests/BLI_index_mask_test.cc

@@ -418,4 +418,39 @@ TEST(index_mask, FindSmallerEqual)
   }
 }
 
+TEST(index_mask, SliceContent)
+{
+  IndexMaskMemory memory;
+  {
+    const IndexMask mask;
+    EXPECT_TRUE(mask.slice_content(IndexRange(50, 10)).is_empty());
+  }
+  {
+    const IndexMask mask{IndexRange(10, 90)};
+    const IndexMask a = mask.slice_content(IndexRange(30));
+    EXPECT_EQ(a.size(), 20);
+    const IndexMask b = mask.slice_content(IndexRange(10, 90));
+    EXPECT_EQ(b.size(), 90);
+    const IndexMask c = mask.slice_content(IndexRange(80, 100));
+    EXPECT_EQ(c.size(), 20);
+    const IndexMask d = mask.slice_content(IndexRange(1000, 100));
+    EXPECT_EQ(d.size(), 0);
+  }
+  {
+    const IndexMask mask = IndexMask::from_initializers(
+        {4, 5, 100, 1'000, 10'000, 20'000, 25'000, 100'000}, memory);
+    EXPECT_EQ(mask.slice_content(IndexRange(10)).size(), 2);
+    EXPECT_EQ(mask.slice_content(IndexRange(200)).size(), 3);
+    EXPECT_EQ(mask.slice_content(IndexRange(2'000)).size(), 4);
+    EXPECT_EQ(mask.slice_content(IndexRange(10'000)).size(), 4);
+    EXPECT_EQ(mask.slice_content(IndexRange(10'001)).size(), 5);
+    EXPECT_EQ(mask.slice_content(IndexRange(1'000'000)).size(), 8);
+    EXPECT_EQ(mask.slice_content(IndexRange(10'000, 100'000)).size(), 4);
+    EXPECT_EQ(mask.slice_content(IndexRange(1'001, 100'000)).size(), 4);
+    EXPECT_EQ(mask.slice_content(IndexRange(1'000, 100'000)).size(), 5);
+    EXPECT_EQ(mask.slice_content(IndexRange(1'000, 99'000)).size(), 4);
+    EXPECT_EQ(mask.slice_content(IndexRange(1'000, 10'000)).size(), 2);
+  }
+}
+
 }  // namespace blender::index_mask::tests