From 79193649f74d8497e36e3b1ba8aae798a7d48edd Mon Sep 17 00:00:00 2001
From: Bogdan Opanchuk <bogdan@opanchuk.net>
Date: Wed, 14 Dec 2022 19:47:34 -0800
Subject: [PATCH] Add SubResidue trait and impls for Residue and DynResidue

---
 src/uint/modular.rs                         |  7 ++
 src/uint/modular/constant_mod.rs            |  2 +
 src/uint/modular/constant_mod/const_sub.rs  | 72 +++++++++++++++++++
 src/uint/modular/runtime_mod.rs             |  2 +
 src/uint/modular/runtime_mod/runtime_sub.rs | 80 +++++++++++++++++++++
 src/uint/modular/sub.rs                     |  9 +++
 6 files changed, 172 insertions(+)
 create mode 100644 src/uint/modular/constant_mod/const_sub.rs
 create mode 100644 src/uint/modular/runtime_mod/runtime_sub.rs
 create mode 100644 src/uint/modular/sub.rs

diff --git a/src/uint/modular.rs b/src/uint/modular.rs
index a93d772f..59a99434 100644
--- a/src/uint/modular.rs
+++ b/src/uint/modular.rs
@@ -13,6 +13,7 @@ mod add;
 mod inv;
 mod mul;
 mod pow;
+mod sub;
 
 /// Provides a consistent interface to add two residues of the same type together.
 pub trait AddResidue {
@@ -20,6 +21,12 @@ pub trait AddResidue {
     fn add(&self, rhs: &Self) -> Self;
 }
 
+/// Provides a consistent interface to subtract two residues of the same type.
+pub trait SubResidue {
+    /// Computes the (reduced) difference of two residues.
+    fn sub(&self, rhs: &Self) -> Self;
+}
+
 /// Provides a consistent interface to multiply two residues of the same type together.
 pub trait MulResidue
 where
diff --git a/src/uint/modular/constant_mod.rs b/src/uint/modular/constant_mod.rs
index d815f8c5..33b78968 100644
--- a/src/uint/modular/constant_mod.rs
+++ b/src/uint/modular/constant_mod.rs
@@ -14,6 +14,8 @@ mod const_inv;
 mod const_mul;
 /// Exponentiation of residues with a constant modulus
 mod const_pow;
+/// Subtractions between residues with a constant modulus
+mod const_sub;
 
 #[macro_use]
 /// Macros to remove the boilerplate code when dealing with constant moduli.
diff --git a/src/uint/modular/constant_mod/const_sub.rs b/src/uint/modular/constant_mod/const_sub.rs
new file mode 100644
index 00000000..9670c36d
--- /dev/null
+++ b/src/uint/modular/constant_mod/const_sub.rs
@@ -0,0 +1,72 @@
+use core::ops::SubAssign;
+
+use crate::{
+    modular::{sub::sub_montgomery_form, SubResidue},
+    Uint,
+};
+
+use super::{Residue, ResidueParams};
+
+impl<MOD: ResidueParams<LIMBS>, const LIMBS: usize> SubResidue for Residue<MOD, LIMBS> {
+    fn sub(&self, rhs: &Self) -> Self {
+        self.sub(rhs)
+    }
+}
+
+impl<MOD: ResidueParams<LIMBS>, const LIMBS: usize> Residue<MOD, LIMBS> {
+    /// Adds two residues together.
+    pub const fn sub(&self, rhs: &Self) -> Self {
+        Residue {
+            montgomery_form: sub_montgomery_form(
+                &self.montgomery_form,
+                &rhs.montgomery_form,
+                &MOD::MODULUS,
+            ),
+            phantom: core::marker::PhantomData,
+        }
+    }
+}
+
+impl<MOD: ResidueParams<LIMBS>, const LIMBS: usize> SubAssign<&Uint<LIMBS>>
+    for Residue<MOD, LIMBS>
+{
+    fn sub_assign(&mut self, rhs: &Uint<LIMBS>) {
+        *self -= &Residue::new(*rhs);
+    }
+}
+
+impl<MOD: ResidueParams<LIMBS>, const LIMBS: usize> SubAssign<&Self> for Residue<MOD, LIMBS> {
+    fn sub_assign(&mut self, rhs: &Self) {
+        *self = self.sub(rhs);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        const_residue, impl_modulus, modular::constant_mod::ResidueParams, traits::Encoding, U256,
+    };
+
+    impl_modulus!(
+        Modulus,
+        U256,
+        "ffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551"
+    );
+
+    #[test]
+    fn sub_overflow() {
+        let x =
+            U256::from_be_hex("44acf6b7e36c1342c2c5897204fe09504e1e2efb1a900377dbc4e7a6a133ec56");
+        let mut x_mod = const_residue!(x, Modulus);
+
+        let y =
+            U256::from_be_hex("d5777c45019673125ad240f83094d4252d829516fac8601ed01979ec1ec1a251");
+
+        x_mod -= &y;
+
+        let expected =
+            U256::from_be_hex("6f357a71e1d5a03167f34879d469352add829491c6df41ddff65387d7ed56f56");
+
+        assert_eq!(expected, x_mod.retrieve());
+    }
+}
diff --git a/src/uint/modular/runtime_mod.rs b/src/uint/modular/runtime_mod.rs
index 5ce7afd6..ea6836ca 100644
--- a/src/uint/modular/runtime_mod.rs
+++ b/src/uint/modular/runtime_mod.rs
@@ -10,6 +10,8 @@ mod runtime_inv;
 mod runtime_mul;
 /// Exponentiation of residues with a modulus set at runtime
 mod runtime_pow;
+/// Subtractions between residues with a modulus set at runtime
+mod runtime_sub;
 
 /// The parameters to efficiently go to and from the Montgomery form for a modulus provided at runtime.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
diff --git a/src/uint/modular/runtime_mod/runtime_sub.rs b/src/uint/modular/runtime_mod/runtime_sub.rs
new file mode 100644
index 00000000..81ee7d24
--- /dev/null
+++ b/src/uint/modular/runtime_mod/runtime_sub.rs
@@ -0,0 +1,80 @@
+use core::ops::{Sub, SubAssign};
+
+use crate::{
+    modular::{sub::sub_montgomery_form, SubResidue},
+    Uint,
+};
+
+use super::DynResidue;
+
+impl<const LIMBS: usize> SubResidue for DynResidue<LIMBS> {
+    fn sub(&self, rhs: &Self) -> Self {
+        debug_assert_eq!(self.residue_params, rhs.residue_params);
+        Self {
+            montgomery_form: sub_montgomery_form(
+                &self.montgomery_form,
+                &rhs.montgomery_form,
+                &self.residue_params.modulus,
+            ),
+            residue_params: self.residue_params,
+        }
+    }
+}
+
+impl<const LIMBS: usize> SubAssign for DynResidue<LIMBS> {
+    fn sub_assign(&mut self, rhs: Self) {
+        self.montgomery_form = sub_montgomery_form(
+            &self.montgomery_form,
+            &rhs.montgomery_form,
+            &self.residue_params.modulus,
+        );
+    }
+}
+
+impl<const LIMBS: usize> SubAssign<Uint<LIMBS>> for DynResidue<LIMBS> {
+    fn sub_assign(&mut self, rhs: Uint<LIMBS>) {
+        self.montgomery_form = sub_montgomery_form(
+            &self.montgomery_form,
+            &DynResidue::new(rhs, self.residue_params).montgomery_form,
+            &self.residue_params.modulus,
+        );
+    }
+}
+
+impl<const LIMBS: usize> Sub for DynResidue<LIMBS> {
+    type Output = DynResidue<LIMBS>;
+
+    fn sub(mut self, rhs: Self) -> Self::Output {
+        self -= rhs;
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        modular::runtime_mod::{DynResidue, DynResidueParams},
+        U256,
+    };
+
+    #[test]
+    fn sub_overflow() {
+        let params = DynResidueParams::new(U256::from_be_hex(
+            "ffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551",
+        ));
+
+        let x =
+            U256::from_be_hex("44acf6b7e36c1342c2c5897204fe09504e1e2efb1a900377dbc4e7a6a133ec56");
+        let mut x_mod = DynResidue::new(x, params);
+
+        let y =
+            U256::from_be_hex("d5777c45019673125ad240f83094d4252d829516fac8601ed01979ec1ec1a251");
+
+        x_mod -= y;
+
+        let expected =
+            U256::from_be_hex("6f357a71e1d5a03167f34879d469352add829491c6df41ddff65387d7ed56f56");
+
+        assert_eq!(expected, x_mod.retrieve());
+    }
+}
diff --git a/src/uint/modular/sub.rs b/src/uint/modular/sub.rs
new file mode 100644
index 00000000..9c471703
--- /dev/null
+++ b/src/uint/modular/sub.rs
@@ -0,0 +1,9 @@
+use crate::Uint;
+
+pub(crate) const fn sub_montgomery_form<const LIMBS: usize>(
+    a: &Uint<LIMBS>,
+    b: &Uint<LIMBS>,
+    modulus: &Uint<LIMBS>,
+) -> Uint<LIMBS> {
+    a.sub_mod(b, modulus)
+}