Implement Weekday::nth_next (#544)

Co-authored-by: Jacob Pratt <jacob@jhpratt.dev>
time-rs · Apr 20, 2023 · 6795714 · 6795714
1 parent 4507d07
commit 6795714
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 1 deletion.
diff --git a/benchmarks/weekday.rs b/benchmarks/weekday.rs
@@ -24,6 +24,21 @@ setup_benchmark! {
         ben.iter(|| Saturday.next());
     }
 
+    fn nth(ben: &mut Bencher<'_>) {
+        ben.iter(|| Sunday.nth_next(0));
+        ben.iter(|| Sunday.nth_next(1));
+        ben.iter(|| Sunday.nth_next(2));
+        ben.iter(|| Sunday.nth_next(3));
+        ben.iter(|| Sunday.nth_next(4));
+        ben.iter(|| Sunday.nth_next(5));
+        ben.iter(|| Sunday.nth_next(6));
+
+        ben.iter(|| Sunday.nth_next(7));
+        ben.iter(|| Sunday.nth_next(u8::MAX));
+        ben.iter(|| Monday.nth_next(7));
+        ben.iter(|| Monday.nth_next(u8::MAX));
+    }
+
     fn number_from_monday(ben: &mut Bencher<'_>) {
         ben.iter(|| Monday.number_from_monday());
         ben.iter(|| Tuesday.number_from_monday());

diff --git a/tests/quickcheck.rs b/tests/quickcheck.rs
@@ -1,5 +1,6 @@
 use quickcheck::{Arbitrary, TestResult};
 use quickcheck_macros::quickcheck;
+use time::Weekday::*;
 use time::{Date, Duration, Month, OffsetDateTime, PrimitiveDateTime, Time, UtcOffset, Weekday};
 
 macro_rules! test_shrink {
@@ -112,6 +113,26 @@ fn unix_timestamp_nanos_roundtrip(odt: OffsetDateTime) -> TestResult {
     }
 }
 
+#[quickcheck]
+fn number_from_monday_roundtrip(w: Weekday) -> bool {
+    Monday.nth_next(w.number_from_monday() + 7 - 1) == w
+}
+
+#[quickcheck]
+fn number_from_sunday_roundtrip(w: Weekday) -> bool {
+    Sunday.nth_next(w.number_from_sunday() + 7 - 1) == w
+}
+
+#[quickcheck]
+fn number_days_from_monday_roundtrip(w: Weekday) -> bool {
+    Monday.nth_next(w.number_days_from_monday()) == w
+}
+
+#[quickcheck]
+fn number_days_from_sunday_roundtrip(w: Weekday) -> bool {
+    Sunday.nth_next(w.number_days_from_sunday()) == w
+}
+
 #[quickcheck]
 fn weekday_supports_arbitrary(w: Weekday) -> bool {
     (1..=7).contains(&w.number_from_monday())
@@ -120,7 +141,7 @@ fn weekday_supports_arbitrary(w: Weekday) -> bool {
 #[quickcheck]
 fn weekday_can_shrink(w: Weekday) -> bool {
     match w {
-        Weekday::Monday => w.shrink().next().is_none(),
+        Monday => w.shrink().next().is_none(),
         _ => w.shrink().next() == Some(w.previous()),
     }
 }

diff --git a/tests/weekday.rs b/tests/weekday.rs
@@ -22,6 +22,30 @@ fn next() {
     assert_eq!(Saturday.next(), Sunday);
 }
 
+#[test]
+fn nth_next() {
+    assert_eq!(Sunday.nth_next(0), Sunday);
+    assert_eq!(Sunday.nth_next(1), Monday);
+    assert_eq!(Sunday.nth_next(2), Tuesday);
+    assert_eq!(Sunday.nth_next(3), Wednesday);
+    assert_eq!(Sunday.nth_next(4), Thursday);
+    assert_eq!(Sunday.nth_next(5), Friday);
+    assert_eq!(Sunday.nth_next(6), Saturday);
+
+    assert_eq!(Monday.nth_next(0), Monday);
+    assert_eq!(Monday.nth_next(1), Tuesday);
+    assert_eq!(Monday.nth_next(2), Wednesday);
+    assert_eq!(Monday.nth_next(3), Thursday);
+    assert_eq!(Monday.nth_next(4), Friday);
+    assert_eq!(Monday.nth_next(5), Saturday);
+    assert_eq!(Monday.nth_next(6), Sunday);
+
+    assert_eq!(Sunday.nth_next(7), Sunday);
+    assert_eq!(Sunday.nth_next(u8::MAX), Wednesday);
+    assert_eq!(Monday.nth_next(7), Monday);
+    assert_eq!(Monday.nth_next(u8::MAX), Thursday);
+}
+
 #[test]
 fn number_from_monday() {
     assert_eq!(Monday.number_from_monday(), 1);

diff --git a/time/src/weekday.rs b/time/src/weekday.rs
@@ -66,6 +66,28 @@ impl Weekday {
         }
     }
 
+    /// Get n-th next day.
+    ///
+    /// ```rust
+    /// # use time::Weekday;
+    /// assert_eq!(Weekday::Monday.nth_next(1), Weekday::Tuesday);
+    /// assert_eq!(Weekday::Sunday.nth_next(10), Weekday::Wednesday);
+    /// ```
+    pub const fn nth_next(self, n: u8) -> Self {
+        match (self.number_days_from_monday() + n % 7) % 7 {
+            0 => Monday,
+            1 => Tuesday,
+            2 => Wednesday,
+            3 => Thursday,
+            4 => Friday,
+            5 => Saturday,
+            val => {
+                debug_assert!(val == 6);
+                Sunday
+            }
+        }
+    }
+
     /// Get the one-indexed number of days from Monday.
     ///
     /// ```rust