From 6fdf08baeeb7e4f4ec3e3188b92043caffbf8714 Mon Sep 17 00:00:00 2001
From: Erik Takke <erik.takke@3milabs.tech>
Date: Fri, 24 Jan 2025 16:34:52 +0100
Subject: [PATCH] Introduce `Uint::new_overflowing_shr`

---
 benches/uint.rs | 8 ++++++++
 src/uint/shr.rs | 6 ++++++
 2 files changed, 14 insertions(+)

diff --git a/benches/uint.rs b/benches/uint.rs
index 3bb2a961..c6c9d858 100644
--- a/benches/uint.rs
+++ b/benches/uint.rs
@@ -405,6 +405,14 @@ fn bench_shr(c: &mut Criterion) {
         )
     });
 
+    group.bench_function("new shr, U2048", |b| {
+        b.iter_batched(
+            || U2048::ONE,
+            |x| x.new_overflowing_shr(1024 + 10),
+            BatchSize::SmallInput,
+        )
+    });
+
     group.finish();
 }
 
diff --git a/src/uint/shr.rs b/src/uint/shr.rs
index d60eb9a4..717be274 100644
--- a/src/uint/shr.rs
+++ b/src/uint/shr.rs
@@ -47,6 +47,12 @@ impl<const LIMBS: usize> Uint<LIMBS> {
         ConstCtOption::new(Uint::select(&result, &Self::ZERO, overflow), overflow.not())
     }
 
+    pub const fn new_overflowing_shr(&self, shift: u32) -> ConstCtOption<Self> {
+        let (intra_limb_shift, limb_shift) = (shift % Limb::BITS, shift / Limb::BITS);
+        self.intra_limb_carrying_shr_internal(intra_limb_shift)
+            .full_limb_shr(limb_shift)
+    }
+
     /// Computes `self >> shift`, for `shift < Limb::BITS`.
     ///
     /// Returns `None` if `shift >= Limb::BITS`.