Add various arithmetic operations

CHANGELOG.md

@@ -7,6 +7,12 @@
 - Support `Step` so that arbitrary-int can be used in a range expression, e.g. `for n in u3::MIN..=u3::MAX { println!("{n}") }`. Note this trait is currently unstable, and so is only usable in nightly. Enable this feature with `step_trait`.
 - Support formatting via [defmt](https://crates.io/crates/defmt). Enable the option `defmt` feature
 - Support serializing and deserializing via [serde](https://crates.io/crates/serde). Enable the option `serde` feature
+- Implement `Mul`, `MulAssign`, `Div`, `DivAssign`
+- Implement `wrapping_add`, `wrapping_sub`, `wrapping_mul`, `wrapping_div`, `wrapping_shl`, `wrapping_shr`
+- Implement `saturating_add`, `saturating_sub`, `saturating_mul`, `saturating_div`, `saturating_pow`
+
+### Changed
+- In debug builds, `<<` (`Shl`, `ShlAssign`) and `>>` (`Shr`, `ShrAssign`) now bounds-check the shift amount using the same semantics as built-in shifts. For example, shifting a u5 by 5 or more bits will now panic as expected.
 
 ## arbitrary-int 1.2.6
 

src/lib.rs

@@ -12,8 +12,8 @@ use core::iter::Step;
 #[cfg(feature = "num-traits")]
 use core::num::Wrapping;
 use core::ops::{
-    Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Not, Shl,
-    ShlAssign, Shr, ShrAssign, Sub, SubAssign,
+    Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Div, DivAssign,
+    Mul, MulAssign, Not, Shl, ShlAssign, Shr, ShrAssign, Sub, SubAssign,
 };
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
@@ -340,6 +340,121 @@ macro_rules! uint_impl {
                     UInt::<$type, BITS_RESULT> { value: self.value }
                 }
 
+                pub const fn wrapping_add(self, rhs: Self) -> Self {
+                    let sum = self.value.wrapping_add(rhs.value);
+                    Self {
+                        value: sum & Self::MASK,
+                    }
+                }
+
+                pub const fn wrapping_sub(self, rhs: Self) -> Self {
+                    let sum = self.value.wrapping_sub(rhs.value);
+                    Self {
+                        value: sum & Self::MASK,
+                    }
+                }
+
+                pub const fn wrapping_mul(self, rhs: Self) -> Self {
+                    let sum = self.value.wrapping_mul(rhs.value);
+                    Self {
+                        value: sum & Self::MASK,
+                    }
+                }
+
+                pub const fn wrapping_div(self, rhs: Self) -> Self {
+                    let sum = self.value.wrapping_div(rhs.value);
+                    Self {
+                        // No need to mask here - divisions always produce a result that is <= self
+                        value: sum,
+                    }
+                }
+
+                pub const fn wrapping_shl(self, rhs: u32) -> Self {
+                    // modulo is expensive on some platforms, so only do it when necessary
+                    let shift_amount = if rhs >= (BITS as u32) {
+                        rhs % (BITS as u32)
+                    } else {
+                        rhs
+                    };
+
+                    Self {
+                        value: (self.value << shift_amount) & Self::MASK,
+                    }
+                }
+
+                pub const fn wrapping_shr(self, rhs: u32) -> Self {
+                    // modulo is expensive on some platforms, so only do it when necessary
+                    let shift_amount = if rhs >= (BITS as u32) {
+                        rhs % (BITS as u32)
+                    } else {
+                        rhs
+                    };
+
+                    Self {
+                        value: (self.value >> shift_amount) & Self::MASK,
+                    }
+                }
+
+                pub const fn saturating_add(self, rhs: Self) -> Self {
+                    let saturated = if core::mem::size_of::<$type>() << 3 == BITS {
+                        // We are something like a UInt::<u8; 8>. We can fallback to the base implementation
+                        self.value.saturating_add(rhs.value)
+                    } else {
+                        // We're dealing with fewer bits than the underlying type (e.g. u7).
+                        // That means the addition can never overflow the underlying type
+                        let sum = self.value.wrapping_add(rhs.value);
+                        let max = Self::MAX.value();
+                        if sum > max { max } else { sum }
+                    };
+                    Self {
+                        value: saturated,
+                    }
+                }
+
+                pub const fn saturating_sub(self, rhs: Self) -> Self {
+                    // For unsigned numbers, the only difference is when we reach 0 - which is the same
+                    // no matter the data size
+                    Self {
+                        value: self.value.saturating_sub(rhs.value),
+                    }
+                }
+
+                pub const fn saturating_mul(self, rhs: Self) -> Self {
+                    let product = if BITS << 1 <= (core::mem::size_of::<$type>() << 3) {
+                        // We have half the bits (e.g. u4 * u4) of the base type, so we can't overflow the base type
+                        self.value * rhs.value
+                    } else {
+                        // We have more than half the bits (e.g. u6 * u6)
+                        self.value.saturating_mul(rhs.value)
+                    };
+
+                    let max = Self::MAX.value();
+                    let saturated = if product > max { max } else { product };
+                    Self {
+                        value: saturated,
+                    }
+                }
+
+                pub const fn saturating_div(self, rhs: Self) -> Self {
+                    // When dividing unsigned numbers, we never need to saturate.
+                    // Divison by zero in saturating_div throws an exception (in debug and release mode),
+                    // so no need to do anything special there either
+                    Self {
+                        value: self.value.saturating_div(rhs.value),
+                    }
+                }
+
+                pub const fn saturating_pow(self, exp: u32) -> Self {
+                    // It might be possible to handwrite this to be slightly faster as both
+                    // saturating_pow has to do a bounds-check and then we do second one
+                    let powed = self.value.saturating_pow(exp);
+                    let max = Self::MAX.value();
+                    let saturated = if powed > max { max } else { powed };
+                    Self {
+                        value: saturated,
+                    }
+                }
+
                 /// Reverses the order of bits in the integer. The least significant bit becomes the most significant bit, second least-significant bit becomes second most-significant bit, etc.
                 pub const fn reverse_bits(self) -> Self {
                     let shift_right = (core::mem::size_of::<$type>() << 3) - BITS;
@@ -501,6 +616,102 @@ where
     }
 }
 
+impl<T, const BITS: usize> Mul for UInt<T, BITS>
+where
+    Self: Number,
+    T: PartialEq
+        + Copy
+        + BitAnd<T, Output = T>
+        + Not<Output = T>
+        + Mul<T, Output = T>
+        + Sub<T, Output = T>
+        + Shr<usize, Output = T>
+        + Shl<usize, Output = T>
+        + From<u8>,
+{
+    type Output = UInt<T, BITS>;
+
+    fn mul(self, rhs: Self) -> Self::Output {
+        let product = self.value * rhs.value;
+        #[cfg(debug_assertions)]
+        if (product & !Self::MASK) != T::from(0) {
+            panic!("attempt to multiply with overflow");
+        }
+        Self {
+            value: product & Self::MASK,
+        }
+    }
+}
+
+impl<T, const BITS: usize> MulAssign for UInt<T, BITS>
+where
+    Self: Number,
+    T: PartialEq
+        + Eq
+        + Not<Output = T>
+        + Copy
+        + MulAssign<T>
+        + BitAnd<T, Output = T>
+        + BitAndAssign<T>
+        + Sub<T, Output = T>
+        + Shr<usize, Output = T>
+        + Shl<usize, Output = T>
+        + From<u8>,
+{
+    fn mul_assign(&mut self, rhs: Self) {
+        self.value *= rhs.value;
+        #[cfg(debug_assertions)]
+        if (self.value & !Self::MASK) != T::from(0) {
+            panic!("attempt to multiply with overflow");
+        }
+        self.value &= Self::MASK;
+    }
+}
+
+impl<T, const BITS: usize> Div for UInt<T, BITS>
+where
+    Self: Number,
+    T: PartialEq
+        + Copy
+        + BitAnd<T, Output = T>
+        + Not<Output = T>
+        + Div<T, Output = T>
+        + Sub<T, Output = T>
+        + Shr<usize, Output = T>
+        + Shl<usize, Output = T>
+        + From<u8>,
+{
+    type Output = UInt<T, BITS>;
+
+    fn div(self, rhs: Self) -> Self::Output {
+        // Integer division can only make the value smaller. And as the result is same type as
+        // Self, there's no need to range-check or mask
+        Self {
+            value: self.value / rhs.value,
+        }
+    }
+}
+
+impl<T, const BITS: usize> DivAssign for UInt<T, BITS>
+where
+    Self: Number,
+    T: PartialEq
+        + Eq
+        + Not<Output = T>
+        + Copy
+        + DivAssign<T>
+        + BitAnd<T, Output = T>
+        + BitAndAssign<T>
+        + Sub<T, Output = T>
+        + Shr<usize, Output = T>
+        + Shl<usize, Output = T>
+        + From<u8>,
+{
+    fn div_assign(&mut self, rhs: Self) {
+        self.value /= rhs.value;
+    }
+}
+
 impl<T, const BITS: usize> BitAnd for UInt<T, BITS>
 where
     Self: Number,
@@ -603,10 +814,18 @@ where
         + Shl<usize, Output = T>
         + Shr<usize, Output = T>
         + From<u8>,
+    TSHIFTBITS: TryInto<usize> + Copy,
 {
     type Output = UInt<T, BITS>;
 
     fn shl(self, rhs: TSHIFTBITS) -> Self::Output {
+        // With debug assertions, the << and >> operators throw an exception if the shift amount
+        // is larger than the number of bits (in which case the result would always be 0)
+        #[cfg(debug_assertions)]
+        if rhs.try_into().unwrap_or(usize::MAX) >= BITS {
+            panic!("attempt to shift left with overflow")
+        }
+
         Self {
             value: (self.value << rhs) & Self::MASK,
         }
@@ -624,8 +843,15 @@ where
         + Shr<usize, Output = T>
         + Shl<usize, Output = T>
         + From<u8>,
+    TSHIFTBITS: TryInto<usize> + Copy,
 {
     fn shl_assign(&mut self, rhs: TSHIFTBITS) {
+        // With debug assertions, the << and >> operators throw an exception if the shift amount
+        // is larger than the number of bits (in which case the result would always be 0)
+        #[cfg(debug_assertions)]
+        if rhs.try_into().unwrap_or(usize::MAX) >= BITS {
+            panic!("attempt to shift left with overflow")
+        }
         self.value <<= rhs;
         self.value &= Self::MASK;
     }
@@ -634,10 +860,17 @@ where
 impl<T, TSHIFTBITS, const BITS: usize> Shr<TSHIFTBITS> for UInt<T, BITS>
 where
     T: Copy + Shr<TSHIFTBITS, Output = T> + Sub<T, Output = T> + Shl<usize, Output = T> + From<u8>,
+    TSHIFTBITS: TryInto<usize> + Copy,
 {
     type Output = UInt<T, BITS>;
 
     fn shr(self, rhs: TSHIFTBITS) -> Self::Output {
+        // With debug assertions, the << and >> operators throw an exception if the shift amount
+        // is larger than the number of bits (in which case the result would always be 0)
+        #[cfg(debug_assertions)]
+        if rhs.try_into().unwrap_or(usize::MAX) >= BITS {
+            panic!("attempt to shift left with overflow")
+        }
         Self {
             value: self.value >> rhs,
         }
@@ -647,8 +880,15 @@ where
 impl<T, TSHIFTBITS, const BITS: usize> ShrAssign<TSHIFTBITS> for UInt<T, BITS>
 where
     T: Copy + ShrAssign<TSHIFTBITS> + Sub<T, Output = T> + Shl<usize, Output = T> + From<u8>,
+    TSHIFTBITS: TryInto<usize> + Copy,
 {
     fn shr_assign(&mut self, rhs: TSHIFTBITS) {
+        // With debug assertions, the << and >> operators throw an exception if the shift amount
+        // is larger than the number of bits (in which case the result would always be 0)
+        #[cfg(debug_assertions)]
+        if rhs.try_into().unwrap_or(usize::MAX) >= BITS {
+            panic!("attempt to shift left with overflow")
+        }
         self.value >>= rhs;
     }
 }