Month::overflowing_add/sub

remove wayward debug

fix formatting, add better docs, improve impl, check for valid years

src/month.rs

@@ -81,6 +81,52 @@ impl Month {
         }
     }
 
+    /// Add `Months` to the current `Month` returning the new `Month`
+    /// and any years that have been crossed.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use chrono::{Month, Months};
+    ///
+    /// assert_eq!(Month::January.overflowing_add(Months::new(3)), (0, Month::April));
+    /// assert_eq!(Month::January.overflowing_add(Months::new(15)), (1, Month::April));
+    /// ```
+    pub fn overflowing_add(&self, months: Months) -> (i32, Month) {
+        let month = self.number_from_month() + months.0 % 12;
+
+        let (years, month) = if month > 12 {
+            (months.as_i32() / 12 + 1, month % 12)
+        } else {
+            (months.as_i32() / 12, month)
+        };
+
+        (years, num_traits::FromPrimitive::from_u32(month).unwrap())
+    }
+
+    /// Subtract `Months` from the current `Month` returning the new `Month`
+    /// and any years that have been crossed
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use chrono::{Month, Months};
+    ///
+    /// assert_eq!(Month::April.overflowing_sub(Months::new(3)), (0, Month::January));
+    /// assert_eq!(Month::April.overflowing_sub(Months::new(15)), (-1, Month::January));
+    /// ```
+    pub fn overflowing_sub(&self, months: Months) -> (i32, Month) {
+        let rem_mths = months.0 % 12;
+
+        let (years, month) = if self.number_from_month() <= rem_mths {
+            (-months.as_i32() / 12 - 1, (12 + self.number_from_month()) - rem_mths)
+        } else {
+            (-months.as_i32() / 12, self.number_from_month() - rem_mths)
+        };
+
+        (years, num_traits::FromPrimitive::from_u32(month).unwrap())
+    }
+
     /// The previous month.
     ///
     /// `m`:        | `January`  | `February` | `...` | `December`
@@ -195,8 +241,23 @@ pub struct Months(pub(crate) u32);
 
 impl Months {
     /// Construct a new `Months` from a number of months
+    pub fn new_opt(num: u32) -> Option<Self> {
+        if num <= core::i32::MAX as u32 {
+            Some(Months(num))
+        } else {
+            None
+        }
+    }
+
+    /// Construct a new `Months` from a number of months
+    /// # panics
+    /// if the u32 is greater than i32::MAX
     pub fn new(num: u32) -> Self {
-        Self(num)
+        Months::new_opt(num).unwrap()
+    }
+
+    fn as_i32(&self) -> i32 {
+        self.0 as i32
     }
 }
 
@@ -318,7 +379,7 @@ mod month_serde {
 
 #[cfg(test)]
 mod tests {
-    use super::Month;
+    use super::{Month, Months};
     use crate::{Datelike, TimeZone, Utc};
 
     #[test]
@@ -342,6 +403,28 @@ mod tests {
         assert_eq!((dt.year(), dt.month(), dt.day()), (2019, 1, 28));
     }
 
+    #[test]
+    fn test_overflowing_add() {
+        assert_eq!(Month::December.overflowing_add(Months::new(5)), (1, Month::May));
+        assert_eq!(Month::December.overflowing_add(Months::new(15)), (2, Month::March));
+        assert_eq!(
+            Month::December.overflowing_add(Months::new(15 + 12 * 200)),
+            (202, Month::March)
+        );
+        assert_eq!(Month::January.overflowing_add(Months::new(3)), (0, Month::April));
+    }
+
+    #[test]
+    fn test_overflowing_sub() {
+        assert_eq!(Month::December.overflowing_sub(Months::new(5)), (0, Month::July));
+        assert_eq!(Month::January.overflowing_sub(Months::new(15)), (-2, Month::October));
+        assert_eq!(
+            Month::January.overflowing_sub(Months::new(15 + 12 * 200)),
+            (-2 - 200, Month::October)
+        );
+        assert_eq!(Month::January.overflowing_sub(Months::new(3)), (-1, Month::October));
+    }
+
     #[test]
     fn test_month_enum_succ_pred() {
         assert_eq!(Month::January.succ(), Month::February);

src/naive/date.rs

@@ -619,10 +619,7 @@ impl NaiveDate {
             return Some(self);
         }
 
-        match months.0 <= core::i32::MAX as u32 {
-            true => self.diff_months(months.0 as i32),
-            false => None,
-        }
+        self.diff_months(months, true)
     }
 
     /// Subtract a duration in [`Months`] from the date
@@ -643,44 +640,30 @@ impl NaiveDate {
             return Some(self);
         }
 
-        // Copy `i32::MIN` here so we don't have to do a complicated cast
-        match months.0 <= 2_147_483_648 {
-            true => self.diff_months(-(months.0 as i32)),
-            false => None,
-        }
+        self.diff_months(months, false)
     }
 
-    fn diff_months(self, months: i32) -> Option<Self> {
-        let (years, left) = ((months / 12), (months % 12));
-
-        // Determine new year (without taking months into account for now
+    fn diff_months(self, months: Months, forwards: bool) -> Option<Self> {
+        // safe as .month() returns [1..12]
+        // can be further improved when .month() returns a Month
+        let month: crate::Month = num_traits::FromPrimitive::from_u32(self.month()).unwrap();
 
-        let year = if (years > 0 && years > (MAX_YEAR - self.year()))
-            || (years < 0 && years < (MIN_YEAR - self.year()))
-        {
-            return None;
-        } else {
-            self.year() + years
+        let (years, month) = match forwards {
+            true => month.overflowing_add(months),
+            false => month.overflowing_sub(months),
         };
 
-        // Determine new month
+        // Add the delta years to the current year
 
-        let month = self.month() as i32 + left;
-        let (year, month) = if month <= 0 {
-            if year == MIN_YEAR {
-                return None;
-            }
+        let year = self.year() + years;
 
-            (year - 1, month + 12)
-        } else if month > 12 {
-            if year == MAX_YEAR {
-                return None;
-            }
+        // Check whether the target year is in the valid range
 
-            (year + 1, month - 12)
-        } else {
-            (year, month)
-        };
+        if year < MIN_YEAR || year > MAX_YEAR {
+            return None;
+        }
+
+        let month = month.number_from_month();
 
         // Clamp original day in case new month is shorter
 
@@ -689,7 +672,7 @@ impl NaiveDate {
         let days = [31, feb_days, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31];
         let day = Ord::min(self.day(), days[(month - 1) as usize]);
 
-        NaiveDate::from_mdf(year, Mdf::new(month as u32, day, flags))
+        NaiveDate::from_mdf(year, Mdf::new(month, day, flags))
     }
 
     /// Makes a new `NaiveDateTime` from the current date and given `NaiveTime`.
@@ -2160,19 +2143,6 @@ mod tests {
             Some(NaiveDate::from_ymd(2022, 8, 3))
         );
 
-        // add with months exceeding `i32::MAX`
-        assert_eq!(
-            NaiveDate::from_ymd(2022, 8, 3).checked_add_months(Months::new(i32::MAX as u32 + 1)),
-            None
-        );
-
-        // sub with months exceeindg `i32::MIN`
-        assert_eq!(
-            NaiveDate::from_ymd(2022, 8, 3)
-                .checked_sub_months(Months::new((i32::MIN as i64).abs() as u32 + 1)),
-            None
-        );
-
         // add overflowing year
         assert_eq!(NaiveDate::MAX.checked_add_months(Months::new(1)), None);