introduced TiledRange1::make_shifted

ValeevGroup · Sep 8, 2024 · 918d4f6 · 918d4f6
1 parent e78e231
commit 918d4f6
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 0 deletions.
diff --git a/src/TiledArray/tiled_range1.h b/src/TiledArray/tiled_range1.h
@@ -50,6 +50,7 @@ class TiledRange1 {
  public:
   using range_type = Range1;
   using index1_type = range_type::index1_type;
+  using signed_index1_type = std::make_signed_t<index1_type>;
   using const_iterator = std::vector<range_type>::const_iterator;
 
   /// Default constructor creates an empty range (tile and element ranges are
@@ -305,6 +306,27 @@ class TiledRange1 {
     return make_uniform(Range1(0, range_extent), target_tile_size);
   }
 
+  /// same as make_uniform(const Range1&, std::size_t) for a 0-based range
+  /// specified by its extent
+  TiledRange1 make_shifted(signed_index1_type shift) {
+    // ensure that it's safe to shift
+    TA_ASSERT(shift <= 0 || shift <= (std::numeric_limits<index1_type>::max() -
+                                      elements_range().upbound()));
+    TA_ASSERT(shift >= 0 ||
+              std::abs(shift) <= (elements_range().lobound() -
+                                  std::numeric_limits<index1_type>::min()));
+    std::vector<index1_type> hashmarks;
+    hashmarks.reserve(tile_extent() + 1);
+    if (tiles_ranges_.empty())
+      hashmarks.emplace_back(elements_range_.lobound() + shift);
+    else {
+      for (auto& t : tiles_ranges_) {
+        hashmarks.push_back(t.first + shift);
+      }
+    }
+    return TiledRange1(hashmarks.begin(), hashmarks.end());
+  }
+
   /// swapper
 
   /// \param other the range with which the contents of this range will be

diff --git a/tests/tiled_range1.cpp b/tests/tiled_range1.cpp
@@ -360,4 +360,23 @@ BOOST_AUTO_TEST_CASE(make_uniform) {
               (TiledRange1{0, 10, 20, 30, 40, 50, 59}));
 }
 
+BOOST_AUTO_TEST_CASE(make_shifted) {
+  TiledRange1 r0;
+  BOOST_REQUIRE_NO_THROW(r0.make_shifted(1));
+  auto r0_plus_1 = r0.make_shifted(1);
+  BOOST_CHECK_EQUAL(r0_plus_1, TiledRange1(1));
+
+  BOOST_CHECK_TA_ASSERT(
+      TiledRange1{std::numeric_limits<index1_type>::max()}.make_shifted(1),
+      Exception);
+  BOOST_CHECK_TA_ASSERT(
+      TiledRange1{std::numeric_limits<index1_type>::min()}.make_shifted(-1),
+      Exception);
+
+  TiledRange1 r1{1, 3, 7, 9};
+  BOOST_REQUIRE_NO_THROW(r0.make_shifted(-1));
+  auto r1_minus_1 = r1.make_shifted(-1);
+  BOOST_CHECK_EQUAL(r1_minus_1, TiledRange1(0, 2, 6, 8));
+}
+
 BOOST_AUTO_TEST_SUITE_END()