BLI: add IndexMask.shift method

source/blender/blenlib/BLI_index_mask.hh

@@ -299,6 +299,11 @@ class IndexMask : private IndexMaskData {
                             int64_t offset,
                             IndexMaskMemory &memory) const;
 
+  /**
+   * Adds an offset to every index in the mask.
+   */
+  IndexMask shift(const int64_t offset, IndexMaskMemory &memory) const;
+
   /**
    * \return A new index mask that contains all the indices from the universe that are not in the
    * current mask.

source/blender/blenlib/intern/index_mask.cc

@@ -190,20 +190,28 @@ IndexMask IndexMask::slice_and_shift(const int64_t start,
   if (std::optional<IndexRange> range = this->to_range()) {
     return range->slice(start, size).shift(offset);
   }
-  IndexMask sliced_mask = this->slice(start, size);
-  if (offset == 0) {
-    return sliced_mask;
+  return this->slice(start, size).shift(offset, memory);
+}
+
+IndexMask IndexMask::shift(const int64_t offset, IndexMaskMemory &memory) const
+{
+  if (indices_num_ == 0) {
+    return {};
   }
-  if (std::optional<IndexRange> range = sliced_mask.to_range()) {
+  BLI_assert(this->first() + offset >= 0);
+  if (offset == 0) {
+    return *this;
+  }
+  if (std::optional<IndexRange> range = this->to_range()) {
     return range->shift(offset);
   }
-  MutableSpan<int64_t> new_segment_offsets = memory.allocate_array<int64_t>(
-      sliced_mask.segments_num_);
-  for (const int64_t i : new_segment_offsets.index_range()) {
-    new_segment_offsets[i] = sliced_mask.segment_offsets_[i] + offset;
+  IndexMask shifted_mask = *this;
+  MutableSpan<int64_t> new_segment_offsets = memory.allocate_array<int64_t>(segments_num_);
+  for (const int64_t i : IndexRange(segments_num_)) {
+    new_segment_offsets[i] = segment_offsets_[i] + offset;
   }
-  sliced_mask.segment_offsets_ = new_segment_offsets.data();
-  return sliced_mask;
+  shifted_mask.segment_offsets_ = new_segment_offsets.data();
+  return shifted_mask;
 }
 
 /**

source/blender/blenlib/tests/BLI_index_mask_test.cc

@@ -797,4 +797,58 @@ TEST(index_mask, FromEveryNth)
   }
 }
 
+TEST(index_mask, Shift)
+{
+  IndexMaskMemory memory;
+  {
+    const IndexMask mask;
+    const IndexMask shifted_mask = mask.shift(10, memory);
+    EXPECT_TRUE(shifted_mask.is_empty());
+    EXPECT_EQ(mask, shifted_mask);
+  }
+  {
+    const IndexMask mask{IndexRange(100, 10)};
+    const IndexMask shifted_mask = mask.shift(1000, memory);
+    EXPECT_EQ(shifted_mask.size(), 10);
+    EXPECT_EQ(shifted_mask[0], 1100);
+    EXPECT_EQ(shifted_mask[9], 1109);
+  }
+  {
+    const IndexMask mask = IndexMask::from_initializers({4, 6, 7, IndexRange(100, 100)}, memory);
+    const IndexMask shifted_mask = mask.shift(1000, memory).shift(-1000, memory);
+    EXPECT_EQ(mask, shifted_mask);
+  }
+  {
+    const IndexMask mask{IndexRange(100, 10)};
+    const IndexMask shifted_mask = mask.shift(0, memory);
+    EXPECT_EQ(mask, shifted_mask);
+  }
+}
+
+TEST(index_mask, SliceAndShift)
+{
+  IndexMaskMemory memory;
+  {
+    const IndexMask mask{IndexRange(100, 10)};
+    const IndexMask new_mask = mask.slice_and_shift(5, 5, 1000, memory);
+    EXPECT_EQ(new_mask.size(), 5);
+    EXPECT_EQ(new_mask[0], 1105);
+    EXPECT_EQ(new_mask[1], 1106);
+  }
+  {
+    const IndexMask mask = IndexMask::from_indices<int>({10, 100, 1'000, 10'000, 100'000}, memory);
+    const IndexMask new_mask = mask.slice_and_shift(IndexRange(1, 4), -100, memory);
+    EXPECT_EQ(new_mask.size(), 4);
+    EXPECT_EQ(new_mask[0], 0);
+    EXPECT_EQ(new_mask[1], 900);
+    EXPECT_EQ(new_mask[2], 9'900);
+    EXPECT_EQ(new_mask[3], 99'900);
+  }
+  {
+    const IndexMask mask = IndexMask::from_indices<int>({10, 100}, memory);
+    const IndexMask new_mask = mask.slice_and_shift(1, 0, 100, memory);
+    EXPECT_TRUE(new_mask.is_empty());
+  }
+}
+
 }  // namespace blender::index_mask::tests