From 5d66127457fd1cac73bb5f0fdfa2cddeb0d8c8ad Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Sat, 23 Dec 2023 09:27:38 -0700
Subject: [PATCH] Add `InvMod` trait

Adds a trait for computing modular inverses with the modulus provided
as an argument, impl'd for both `Uint` and `BoxedUint`.
---
 src/traits.rs             |  6 ++++++
 src/uint/boxed/inv_mod.rs | 17 +++++++++++++++--
 src/uint/inv_mod.rs       | 19 +++++++++++++++----
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/src/traits.rs b/src/traits.rs
index 2888ee672..881b16466 100644
--- a/src/traits.rs
+++ b/src/traits.rs
@@ -346,6 +346,12 @@ pub trait MulMod<Rhs = Self> {
     fn mul_mod(&self, rhs: &Rhs, p: &Self) -> Self::Output;
 }
 
+/// Compute `1 / self mod p`.
+pub trait InvMod: Sized {
+    /// Compute `1 / self mod p`.
+    fn inv_mod(&self, p: &Self) -> CtOption<Self>;
+}
+
 /// Checked addition.
 pub trait CheckedAdd<Rhs = Self>: Sized {
     /// Perform checked addition, returning a [`CtOption`] which `is_some` only if the operation
diff --git a/src/uint/boxed/inv_mod.rs b/src/uint/boxed/inv_mod.rs
index 32513dcdb..e642daed0 100644
--- a/src/uint/boxed/inv_mod.rs
+++ b/src/uint/boxed/inv_mod.rs
@@ -1,8 +1,8 @@
 //! [`BoxedUint`] modular inverse (i.e. reciprocal) operations.
 
 use crate::{
-    modular::BoxedBernsteinYangInverter, BoxedUint, ConstantTimeSelect, Integer, Inverter, Odd,
-    PrecomputeInverter, PrecomputeInverterWithAdjuster,
+    modular::BoxedBernsteinYangInverter, BoxedUint, ConstantTimeSelect, Integer, InvMod, Inverter,
+    Odd, PrecomputeInverter, PrecomputeInverterWithAdjuster,
 };
 use subtle::{Choice, ConstantTimeEq, ConstantTimeLess, CtOption};
 
@@ -45,6 +45,19 @@ impl BoxedUint {
     }
 }
 
+impl InvMod for BoxedUint {
+    /// Note: currently only supports odd modulus
+    fn inv_mod(&self, modulus: &Self) -> CtOption<Self> {
+        let modulus = Odd(modulus.clone());
+        let is_odd = modulus.is_odd();
+        let maybe_ret = self.inv_odd_mod(&modulus);
+        let is_some = maybe_ret.is_some() & is_odd;
+
+        // use cloned modulus as dummy value for is_none case
+        CtOption::new(Option::from(maybe_ret).unwrap_or(modulus.get()), is_some)
+    }
+}
+
 /// Precompute a Bernstein-Yang inverter using `self` as the modulus.
 impl PrecomputeInverter for Odd<BoxedUint> {
     type Inverter = BoxedBernsteinYangInverter;
diff --git a/src/uint/inv_mod.rs b/src/uint/inv_mod.rs
index 203da385c..1ce884d05 100644
--- a/src/uint/inv_mod.rs
+++ b/src/uint/inv_mod.rs
@@ -1,6 +1,8 @@
 use super::Uint;
-use crate::modular::BernsteinYangInverter;
-use crate::{ConstChoice, ConstCtOption, Odd, PrecomputeInverter};
+use crate::{
+    modular::BernsteinYangInverter, ConstChoice, ConstCtOption, InvMod, Odd, PrecomputeInverter,
+};
+use subtle::CtOption;
 
 impl<const LIMBS: usize> Uint<LIMBS> {
     /// Computes 1/`self` mod `2^k`.
@@ -92,8 +94,8 @@ impl<const LIMBS: usize> Uint<LIMBS> {
     }
 
     /// Computes the multiplicative inverse of `self` mod `modulus`.
-    /// Returns `(inverse, ConstChoice::TRUE)` if an inverse exists,
-    /// otherwise `(undefined, ConstChoice::FALSE)`.
+    ///
+    /// Returns some if an inverse exists, otherwise none.
     pub const fn inv_mod<const UNSAT_LIMBS: usize>(&self, modulus: &Self) -> ConstCtOption<Self>
     where
         Odd<Self>: PrecomputeInverter<Inverter = BernsteinYangInverter<LIMBS, UNSAT_LIMBS>>,
@@ -133,6 +135,15 @@ impl<const LIMBS: usize> Uint<LIMBS> {
     }
 }
 
+impl<const LIMBS: usize, const UNSAT_LIMBS: usize> InvMod for Uint<LIMBS>
+where
+    Odd<Self>: PrecomputeInverter<Inverter = BernsteinYangInverter<LIMBS, UNSAT_LIMBS>>,
+{
+    fn inv_mod(&self, modulus: &Self) -> CtOption<Self> {
+        self.inv_mod(modulus).into()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{U1024, U256, U64};