source/blender/blenlib/BLI_index_mask.hh

@@ -186,6 +186,23 @@ class IndexMask : private IndexMaskData {
   static IndexMask from_bools(const IndexMask &universe,
                               const VArray<bool> &bools,
                               IndexMaskMemory &memory);
+  /**
+   * Constructs a mask by repeating the indices inthe given mask with a stride.
+   * For example, with an input mask containing `{3, 5}` and a stride of 10 the resulting mask
+   * would contain `{3, 5, 13, 15, 23, 25, ...}`.
+   */
+  static IndexMask from_repeating(const IndexMask &mask_to_repeat,
+                                  int64_t repetitions,
+                                  int64_t stride,
+                                  int64_t initial_offset,
+                                  IndexMaskMemory &memory);
+  /**
+   * Constructs a mask that contains every nth index the given number of times.
+   */
+  static IndexMask from_every_nth(int64_t n,
+                                  int64_t indices_num,
+                                  const int64_t initial_offset,
+                                  IndexMaskMemory &memory);
   /**
    * Construct a mask from the given segments. The provided segments are expected to be
    * sorted and owned by #memory already.
@@ -221,6 +238,11 @@ class IndexMask : private IndexMaskData {
   int64_t first() const;
   int64_t last() const;
 
+  /**
+   * Returns the smallest range that contains all indices stored in this mask.
+   */
+  IndexRange bounds() const;
+
   /**
    * \return Minimum number of elements an array has to have so that it can be indexed by every
    * index stored in the mask.
@@ -568,6 +590,17 @@ inline IndexRange IndexMask::index_range() const
   return IndexRange(indices_num_);
 }
 
+inline IndexRange IndexMask::bounds() const
+{
+  if (this->is_empty()) {
+    return IndexRange();
+  }
+  const int64_t first = this->first();
+  const int64_t last = this->last();
+  const int64_t range = last - first + 1;
+  return IndexRange(first, range);
+}
+
 inline int64_t IndexMask::first() const
 {
   BLI_assert(indices_num_ > 0);

source/blender/blenlib/intern/index_mask.cc

@@ -863,10 +863,137 @@ bool IndexMask::contains(const int64_t query_index) const
   return this->find(query_index).has_value();
 }
 
-template IndexMask IndexMask::from_indices(Span<int32_t>, IndexMaskMemory &);
-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;
+static Array<int16_t> build_every_nth_index_array(const int64_t n)
+{
+  Array<int16_t> data(max_segment_size / n);
+  for (const int64_t i : data.index_range()) {
+    const int64_t index = i * n;
+    BLI_assert(index < max_segment_size);
+    data[i] = int16_t(index);
+  }
+  return data;
+}
+
+/**
+ * Returns a span containting every nth index. This is optimized for a few special values of n
+ * which are cached. The returned indices have either static life-time, or they are freed when the
+ * given memory is feed.
+ */
+static Span<int16_t> get_every_nth_index(const int64_t n,
+                                         const int64_t repetitions,
+                                         IndexMaskMemory &memory)
+{
+  BLI_assert(n >= 2);
+  BLI_assert(n * repetitions <= max_segment_size);
+
+  switch (n) {
+    case 2: {
+      static auto data = build_every_nth_index_array(2);
+      return data.as_span().take_front(repetitions);
+    }
+    case 3: {
+      static auto data = build_every_nth_index_array(3);
+      return data.as_span().take_front(repetitions);
+    }
+    case 4: {
+      static auto data = build_every_nth_index_array(4);
+      return data.as_span().take_front(repetitions);
+    }
+    default: {
+      MutableSpan<int16_t> data = memory.allocate_array<int16_t>(repetitions);
+      for (const int64_t i : IndexRange(repetitions)) {
+        const int64_t index = i * n;
+        BLI_assert(index < max_segment_size);
+        data[i] = int16_t(index);
+      }
+      return data;
+    }
+  }
+}
+
+IndexMask IndexMask::from_repeating(const IndexMask &mask_to_repeat,
+                                    const int64_t repetitions,
+                                    const int64_t stride,
+                                    const int64_t initial_offset,
+                                    IndexMaskMemory &memory)
+{
+  if (mask_to_repeat.is_empty()) {
+    return {};
+  }
+  BLI_assert(mask_to_repeat.last() < stride);
+  if (repetitions == 0) {
+    return {};
+  }
+  if (repetitions == 1 && initial_offset == 0) {
+    /* The output is the same as the input mask. */
+    return mask_to_repeat;
+  }
+  const std::optional<IndexRange> range_to_repeat = mask_to_repeat.to_range();
+  if (range_to_repeat && range_to_repeat->first() == 0 && range_to_repeat->size() == stride) {
+    /* The output is a range. */
+    return IndexRange(initial_offset, repetitions * stride);
+  }
+  const int64_t segments_num = mask_to_repeat.segments_num();
+  const IndexRange bounds = mask_to_repeat.bounds();
+
+  /* Avoid having many very small segments by creating a single segment that contains the input
+   * multiple times already. This way, a lower total number of segments is necessary. */
+  if (segments_num == 1 && stride <= max_segment_size / 2 && mask_to_repeat.size() <= 256) {
+    const IndexMaskSegment src_segment = mask_to_repeat.segment(0);
+    /* Number of repetitions that fit into a single segment. */
+    const int64_t inline_repetitions_num = std::min(repetitions, max_segment_size / stride);
+    Span<int16_t> repeated_indices;
+    if (src_segment.size() == 1) {
+      /* Optimize the case when a single index is repeated. */
+      repeated_indices = get_every_nth_index(stride, inline_repetitions_num, memory);
+    }
+    else {
+      /* More general case that repeats multiple indices. */
+      MutableSpan<int16_t> repeated_indices_mut = memory.allocate_array<int16_t>(
+          inline_repetitions_num * src_segment.size());
+      for (const int64_t repetition : IndexRange(inline_repetitions_num)) {
+        for (const int64_t i : src_segment.index_range()) {
+          const int64_t index = src_segment[i] - src_segment[0] + repetition * stride;
+          BLI_assert(index < max_segment_size);
+          repeated_indices_mut[repetition * src_segment.size() + i] = int16_t(index);
+        }
+      }
+      repeated_indices = repeated_indices_mut;
+    }
+    BLI_assert(repeated_indices[0] == 0);
+
+    Vector<IndexMaskSegment, 16> repeated_segments;
+    const int64_t result_segments_num = ceil_division(repetitions, inline_repetitions_num);
+    for (const int64_t i : IndexRange(result_segments_num)) {
+      const int64_t used_repetitions = std::min(inline_repetitions_num,
+                                                repetitions - i * inline_repetitions_num);
+      repeated_segments.append(
+          IndexMaskSegment(initial_offset + bounds.first() + i * stride * inline_repetitions_num,
+                           repeated_indices.take_front(used_repetitions * src_segment.size())));
+    }
+    return IndexMask::from_segments(repeated_segments, memory);
+  }
+
+  /* Simply repeat and offset the existing segments in the input mask. */
+  Vector<IndexMaskSegment, 16> repeated_segments;
+  for (const int64_t repetition : IndexRange(repetitions)) {
+    for (const int64_t segment_i : IndexRange(segments_num)) {
+      const IndexMaskSegment segment = mask_to_repeat.segment(segment_i);
+      repeated_segments.append(IndexMaskSegment(
+          segment.offset() + repetition * stride + initial_offset, segment.base_span()));
+    }
+  }
+  return IndexMask::from_segments(repeated_segments, memory);
+}
+
+IndexMask IndexMask::from_every_nth(const int64_t n,
+                                    const int64_t indices_num,
+                                    const int64_t initial_offset,
+                                    IndexMaskMemory &memory)
+{
+  BLI_assert(n >= 1);
+  return IndexMask::from_repeating(IndexRange(1), indices_num, n, initial_offset, memory);
+}
 
 void IndexMask::foreach_segment_zipped(const Span<IndexMask> masks,
                                        const FunctionRef<bool(Span<IndexMaskSegment> segments)> fn)
@@ -975,4 +1102,9 @@ bool operator==(const IndexMask &a, const IndexMask &b)
   return equals;
 }
 
+template IndexMask IndexMask::from_indices(Span<int32_t>, IndexMaskMemory &);
+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;
+
 }  // namespace blender::index_mask

source/blender/blenlib/tests/BLI_index_mask_test.cc

@@ -23,6 +23,7 @@ TEST(index_mask, IndicesToMask)
   EXPECT_EQ(mask.first(), 5);
   EXPECT_EQ(mask.last(), 101000);
   EXPECT_EQ(mask.min_array_size(), 101001);
+  EXPECT_EQ(mask.bounds(), IndexRange(5, 101001 - 5));
 }
 
 TEST(index_mask, FromBits)
@@ -51,6 +52,7 @@ TEST(index_mask, FromSize)
     EXPECT_EQ(mask.first(), 0);
     EXPECT_EQ(mask.last(), 4);
     EXPECT_EQ(mask.min_array_size(), 5);
+    EXPECT_EQ(mask.bounds(), IndexRange(5));
   }
   {
     const IndexMask mask(max_segment_size);
@@ -61,6 +63,7 @@ TEST(index_mask, FromSize)
     EXPECT_EQ(mask.first(), 0);
     EXPECT_EQ(mask.last(), max_segment_size - 1);
     EXPECT_EQ(mask.min_array_size(), max_segment_size);
+    EXPECT_EQ(mask.bounds(), IndexRange(max_segment_size));
   }
 }
 
@@ -105,6 +108,7 @@ TEST(index_mask, DefaultConstructor)
   IndexMask mask;
   EXPECT_EQ(mask.size(), 0);
   EXPECT_EQ(mask.min_array_size(), 0);
+  EXPECT_EQ(mask.bounds(), IndexRange());
 }
 
 TEST(index_mask, ForeachRange)
@@ -667,4 +671,130 @@ TEST(index_mask, ZippedForeachEqual)
   EXPECT_EQ(index, 4);
 }
 
+TEST(index_mask, FromRepeatingEmpty)
+{
+  IndexMaskMemory memory;
+  const IndexMask mask = IndexMask::from_repeating(IndexMask(), 100, 0, 10, memory);
+  EXPECT_TRUE(mask.is_empty());
+}
+
+TEST(index_mask, FromRepeatingSingle)
+{
+  IndexMaskMemory memory;
+  const IndexMask mask = IndexMask::from_repeating(IndexMask(1), 5, 10, 2, memory);
+  EXPECT_EQ(mask, IndexMask::from_initializers({2, 12, 22, 32, 42}, memory));
+}
+
+TEST(index_mask, FromRepeatingSame)
+{
+  IndexMaskMemory memory;
+  const IndexMask mask = IndexMask::from_indices<int>({4, 6, 7}, memory);
+  const IndexMask repeated_mask = IndexMask::from_repeating(mask, 1, 100, 0, memory);
+  EXPECT_EQ(mask, repeated_mask);
+}
+
+TEST(index_mask, FromRepeatingMultiple)
+{
+  IndexMaskMemory memory;
+  const IndexMask mask = IndexMask::from_repeating(
+      IndexMask::from_indices<int>({5, 6, 7, 50}, memory), 3, 100, 1000, memory);
+  EXPECT_EQ(mask[0], 1005);
+  EXPECT_EQ(mask[1], 1006);
+  EXPECT_EQ(mask[2], 1007);
+  EXPECT_EQ(mask[3], 1050);
+  EXPECT_EQ(mask[4], 1105);
+  EXPECT_EQ(mask[5], 1106);
+  EXPECT_EQ(mask[6], 1107);
+  EXPECT_EQ(mask[7], 1150);
+  EXPECT_EQ(mask[8], 1205);
+  EXPECT_EQ(mask[9], 1206);
+  EXPECT_EQ(mask[10], 1207);
+  EXPECT_EQ(mask[11], 1250);
+}
+
+TEST(index_mask, FromRepeatingRangeFromSingle)
+{
+  IndexMaskMemory memory;
+  const IndexMask mask = IndexMask::from_repeating(IndexMask(IndexRange(1)), 50'000, 1, 0, memory);
+  EXPECT_EQ(*mask.to_range(), IndexRange(50'000));
+}
+
+TEST(index_mask, FromRepeatingRangeFromRange)
+{
+  IndexMaskMemory memory;
+  const IndexMask mask = IndexMask::from_repeating(
+      IndexMask(IndexRange(100)), 50'000, 100, 100, memory);
+  EXPECT_EQ(*mask.to_range(), IndexRange(100, 5'000'000));
+}
+
+TEST(index_mask, FromRepeatingEverySecond)
+{
+  IndexMaskMemory memory;
+  const IndexMask mask = IndexMask::from_repeating(IndexMask(1), 500'000, 2, 0, memory);
+  EXPECT_EQ(mask[0], 0);
+  EXPECT_EQ(mask[1], 2);
+  EXPECT_EQ(mask[2], 4);
+  EXPECT_EQ(mask[3], 6);
+  EXPECT_EQ(mask[20'000], 40'000);
+}
+
+TEST(index_mask, FromRepeatingMultipleRanges)
+{
+  IndexMaskMemory memory;
+  const IndexMask mask = IndexMask::from_repeating(
+      IndexMask::from_initializers({IndexRange(0, 100), IndexRange(10'000, 100)}, memory),
+      5,
+      100'000,
+      0,
+      memory);
+  EXPECT_EQ(mask[0], 0);
+  EXPECT_EQ(mask[1], 1);
+  EXPECT_EQ(mask[2], 2);
+  EXPECT_EQ(mask[100], 10'000);
+  EXPECT_EQ(mask[101], 10'001);
+  EXPECT_EQ(mask[102], 10'002);
+  EXPECT_EQ(mask[200], 100'000);
+  EXPECT_EQ(mask[201], 100'001);
+  EXPECT_EQ(mask[202], 100'002);
+  EXPECT_EQ(mask[300], 110'000);
+  EXPECT_EQ(mask[301], 110'001);
+  EXPECT_EQ(mask[302], 110'002);
+}
+
+TEST(index_mask, FromRepeatingNoRepetitions)
+{
+  IndexMaskMemory memory;
+  const IndexMask mask = IndexMask::from_repeating(IndexMask(IndexRange(5)), 0, 100, 0, memory);
+  EXPECT_TRUE(mask.is_empty());
+}
+
+TEST(index_mask, FromEveryNth)
+{
+  IndexMaskMemory memory;
+  {
+    const IndexMask mask = IndexMask::from_every_nth(2, 5, 0, memory);
+    EXPECT_EQ(mask, IndexMask::from_initializers({0, 2, 4, 6, 8}, memory));
+  }
+  {
+    const IndexMask mask = IndexMask::from_every_nth(3, 5, 100, memory);
+    EXPECT_EQ(mask, IndexMask::from_initializers({100, 103, 106, 109, 112}, memory));
+  }
+  {
+    const IndexMask mask = IndexMask::from_every_nth(4, 5, 0, memory);
+    EXPECT_EQ(mask, IndexMask::from_initializers({0, 4, 8, 12, 16}, memory));
+  }
+  {
+    const IndexMask mask = IndexMask::from_every_nth(10, 5, 100, memory);
+    EXPECT_EQ(mask, IndexMask::from_initializers({100, 110, 120, 130, 140}, memory));
+  }
+  {
+    const IndexMask mask = IndexMask::from_every_nth(1, 5, 100, memory);
+    EXPECT_EQ(mask, IndexMask::from_initializers({100, 101, 102, 103, 104}, memory));
+  }
+  {
+    const IndexMask mask = IndexMask::from_every_nth(100'000, 5, 0, memory);
+    EXPECT_EQ(mask, IndexMask::from_initializers({0, 100'000, 200'000, 300'000, 400'000}, memory));
+  }
+}
+
 }  // namespace blender::index_mask::tests