feat: support specifying chunk multiples (PROOF-831) (#125)

support specifying chunk multiples
spaceandtimelabs · Apr 29, 2024 · 2fe180b · 2fe180b
1 parent add495d
commit 2fe180b
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 12 deletions.
diff --git a/sxt/base/iterator/index_range.cc b/sxt/base/iterator/index_range.cc
@@ -23,10 +23,12 @@ namespace sxt::basit {
 // constructor
 //--------------------------------------------------------------------------------------------------
 index_range::index_range(size_t a, size_t b) noexcept
-    : index_range{a, b, 1, std::numeric_limits<size_t>::max()} {}
+    : index_range{a, b, 1, std::numeric_limits<size_t>::max(), 1} {}
 
-index_range::index_range(size_t a, size_t b, size_t min_chunk_size, size_t max_chunk_size) noexcept
-    : a_{a}, b_{b}, min_chunk_size_{min_chunk_size}, max_chunk_size_{max_chunk_size} {
+index_range::index_range(size_t a, size_t b, size_t min_chunk_size, size_t max_chunk_size,
+                         size_t chunk_multiple) noexcept
+    : a_{a}, b_{b}, min_chunk_size_{min_chunk_size}, max_chunk_size_{max_chunk_size},
+      chunk_multiple_{chunk_multiple} {
   SXT_DEBUG_ASSERT(
       // clang-format off
       0 <= a && a <= b &&
@@ -40,10 +42,7 @@ index_range::index_range(size_t a, size_t b, size_t min_chunk_size, size_t max_c
 //--------------------------------------------------------------------------------------------------
 index_range index_range::min_chunk_size(size_t val) const noexcept {
   return {
-      a_,
-      b_,
-      val,
-      max_chunk_size_,
+      a_, b_, val, max_chunk_size_, chunk_multiple_,
   };
 }
 
@@ -52,10 +51,16 @@ index_range index_range::min_chunk_size(size_t val) const noexcept {
 //--------------------------------------------------------------------------------------------------
 index_range index_range::max_chunk_size(size_t val) const noexcept {
   return {
-      a_,
-      b_,
-      min_chunk_size_,
-      val,
+      a_, b_, min_chunk_size_, val, chunk_multiple_,
+  };
+}
+
+//--------------------------------------------------------------------------------------------------
+// chunk_multiple
+//--------------------------------------------------------------------------------------------------
+index_range index_range::chunk_multiple(size_t val) const noexcept {
+  return {
+      a_, b_, min_chunk_size_, max_chunk_size_, val,
   };
 }
 } // namespace sxt::basit
diff --git a/sxt/base/iterator/index_range.h b/sxt/base/iterator/index_range.h
@@ -29,7 +29,8 @@ class index_range {
 
   index_range(size_t a, size_t b) noexcept;
 
-  index_range(size_t a, size_t b, size_t min_chunk_size, size_t max_chunk_size) noexcept;
+  index_range(size_t a, size_t b, size_t min_chunk_size, size_t max_chunk_size,
+              size_t chunk_multiple) noexcept;
 
   size_t a() const noexcept { return a_; }
   size_t b() const noexcept { return b_; }
@@ -40,15 +41,19 @@ class index_range {
 
   size_t min_chunk_size() const noexcept { return min_chunk_size_; }
   size_t max_chunk_size() const noexcept { return max_chunk_size_; }
+  size_t chunk_multiple() const noexcept { return chunk_multiple_; }
 
   [[nodiscard]] index_range min_chunk_size(size_t val) const noexcept;
 
   [[nodiscard]] index_range max_chunk_size(size_t val) const noexcept;
 
+  [[nodiscard]] index_range chunk_multiple(size_t val) const noexcept;
+
 private:
   size_t a_{0};
   size_t b_{0};
   size_t min_chunk_size_{1};
   size_t max_chunk_size_{std::numeric_limits<size_t>::max()};
+  size_t chunk_multiple_{1};
 };
 } // namespace sxt::basit
diff --git a/sxt/base/iterator/index_range_utility.cc b/sxt/base/iterator/index_range_utility.cc
@@ -34,6 +34,7 @@ std::pair<index_range_iterator, index_range_iterator> split(const index_range& r
   auto step = std::max(basn::divide_up(delta, n), size_t{1});
   step = std::max(step, rng.min_chunk_size());
   step = std::min(step, rng.max_chunk_size());
+  step = basn::divide_up(step, rng.chunk_multiple()) * rng.chunk_multiple();
   index_range_iterator first{index_range{rng.a(), rng.b()}, step};
   index_range_iterator last{index_range{rng.b(), rng.b()}, step};
   return {first, last};

diff --git a/sxt/base/iterator/index_range_utility.t.cc b/sxt/base/iterator/index_range_utility.t.cc
@@ -76,4 +76,12 @@ TEST_CASE("we can split an index_range") {
     REQUIRE(*iter++ == index_range{2, 4});
     REQUIRE(iter == last);
   }
+
+  SECTION("we respect the chunk multiple") {
+    auto [iter, last] = split(index_range{0, 4}.chunk_multiple(3), 4);
+    REQUIRE(std::distance(iter, last) == 2);
+    REQUIRE(*iter++ == index_range{0, 3});
+    REQUIRE(*iter++ == index_range{3, 4});
+    REQUIRE(iter == last);
+  }
 }