recmo · prestwich · May 29, 2025 · May 29, 2025 · May 29, 2025
@@ -9,6 +9,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Check limb overflow in shift ops ([#476])
+
+[#476]: https://github.com/recmo/uint/pull/476
+
 ## [1.15.0] - 2025-05-22
 
 ### Added

@@ -93,7 +93,7 @@ pub fn add_nx1(lhs: &mut [u64], mut a: u64) -> u64 {
 ///
 /// # Panics
 ///
-/// Panics if the lengts are not the same.
+/// Panics if the lengths are not the same.
 #[inline(always)]
 pub fn addmul_n(lhs: &mut [u64], a: &[u64], b: &[u64]) {
     assert_eq!(lhs.len(), a.len());

@@ -335,6 +335,24 @@ impl<const BITS: usize, const LIMBS: usize> Uint<BITS, LIMBS> {
         (r.masked(), carry != 0)
     }
 
+    /// Left shift by `rhs` bits with overflow detection, but with `Self` rhs.
+    ///
+    /// See [`overflowing_shl`](Self::overflowing_shl) for details.
+    #[inline]
+    pub(crate) fn overflowing_shl_big(self, rhs: Self) -> (Self, bool) {
+        if BITS == 0 {
+            return (Self::ZERO, false);
+        }
+        let Ok(rhs) = u64::try_from(rhs) else {
+            return (Self::ZERO, true);
+        };
+        // Rationale: if BITS is larger than 2**64 - 1, it means we're running
+        // on a 128-bit platform with 2.3 exabytes of memory. In this case,
+        // the code produces incorrect output.
+        #[allow(clippy::cast_possible_truncation)]
+        self.overflowing_shl(rhs as usize)
+    }
+
     /// Left shift by `rhs` bits.
     ///
     /// Returns $\mod{\mathtt{value} ⋅ 2^{\mathtt{rhs}}}_{2^{\mathtt{BITS}}}$.
@@ -400,6 +418,24 @@ impl<const BITS: usize, const LIMBS: usize> Uint<BITS, LIMBS> {
         (r, carry != 0)
     }
 
+    /// Right shift by `rhs` bits with underflow detection, but with `Self` rhs.
+    ///
+    /// See [`overflowing_shr`](Self::overflowing_shr) for details.
+    #[inline]
+    pub(crate) fn overflowing_shr_big(self, rhs: Self) -> (Self, bool) {
+        if BITS == 0 {
+            return (Self::ZERO, false);
+        }
+        let Ok(rhs) = u64::try_from(rhs) else {
+            return (Self::ZERO, true);
+        };
+        // Rationale: if BITS is larger than 2**64 - 1, it means we're running
+        // on a 128-bit platform with 2.3 exabytes of memory. In this case,
+        // the code produces incorrect output.
+        #[allow(clippy::cast_possible_truncation)]
+        self.overflowing_shr(rhs as usize)
+    }
+
     /// Right shift by `rhs` bits.
     ///
     /// $$
@@ -569,15 +605,7 @@ impl<const BITS: usize, const LIMBS: usize> Shl<Self> for Uint<BITS, LIMBS> {
 
     #[inline(always)]
     fn shl(self, rhs: Self) -> Self::Output {
-        // This check shortcuts, and prevents panics on the `[0]` later
-        if BITS == 0 {
-            return self;
-        }
-        // Rationale: if BITS is larger than 2**64 - 1, it means we're running
-        // on a 128-bit platform with 2.3 exabytes of memory. In this case,
-        // the code produces incorrect output.
-        #[allow(clippy::cast_possible_truncation)]
-        self.wrapping_shl(rhs.as_limbs()[0] as usize)
+        self.overflowing_shl_big(rhs).0
     }
 }
 
@@ -595,15 +623,7 @@ impl<const BITS: usize, const LIMBS: usize> Shr<Self> for Uint<BITS, LIMBS> {
 
     #[inline(always)]
     fn shr(self, rhs: Self) -> Self::Output {
-        // This check shortcuts, and prevents panics on the `[0]` later
-        if BITS == 0 {
-            return self;
-        }
-        // Rationale: if BITS is larger than 2**64 - 1, it means we're running
-        // on a 128-bit platform with 2.3 exabytes of memory. In this case,
-        // the code produces incorrect output.
-        #[allow(clippy::cast_possible_truncation)]
-        self.wrapping_shr(rhs.as_limbs()[0] as usize)
+        self.overflowing_shr_big(rhs).0
     }
 }
 
@@ -850,6 +870,22 @@ mod tests {
         });
     }
 
+    #[test]
+    fn test_shift_very_big_rhs() {
+        type U = Uint<128, 2>;
+
+        for rhs in [
+            U::from(u64::MAX),
+            U::from(u128::MAX),
+            U::from_limbs([0, 1]),
+            U::from_limbs([1, 1]),
+            U::from_limbs([1, u64::MAX]),
+        ] {
+            assert_eq!(U::ONE << rhs, U::ZERO, "{rhs}");
+            assert_eq!(U::ONE >> rhs, U::ZERO, "{rhs}");
+        }
+    }
+
     #[test]
     fn test_rotate() {
         const_for!(BITS in SIZES {

@@ -63,10 +63,7 @@ impl<const BITS: usize, const LIMBS: usize> Uint<BITS, LIMBS> {
     pub fn wrapping_mul(self, rhs: Self) -> Self {
         let mut result = Self::ZERO;
         algorithms::addmul_n(&mut result.limbs, self.as_limbs(), rhs.as_limbs());
-        if BITS > 0 {
-            result.apply_mask();
-        }
-        result
+        result.masked()
     }
 
     /// Computes the inverse modulo $2^{\mathtt{BITS}}$ of `self`, returning