From 553e72ec77485c1350b5282ce611d4cca4656ec2 Mon Sep 17 00:00:00 2001
From: Bogdan Opanchuk <bogdan@opanchuk.net>
Date: Fri, 29 Dec 2023 11:05:46 -0800
Subject: [PATCH] Additional methods for `Integer` and `Monty` (#533)

* Add `Integer::from_limb_like()`, `one_like()`, `zero_like()`.
* Add `Monty::params()` and `as_montgomery()`
---
 src/modular/boxed_monty_form.rs |  8 ++++++++
 src/modular/monty_form.rs       |  8 ++++++++
 src/traits.rs                   | 24 ++++++++++++++++++++++++
 src/uint.rs                     |  4 ++++
 src/uint/boxed.rs               |  6 ++++++
 5 files changed, 50 insertions(+)

diff --git a/src/modular/boxed_monty_form.rs b/src/modular/boxed_monty_form.rs
index 2c6fe07e..5098d5db 100644
--- a/src/modular/boxed_monty_form.rs
+++ b/src/modular/boxed_monty_form.rs
@@ -256,6 +256,14 @@ impl Monty for BoxedMontyForm {
         BoxedMontyForm::one(params)
     }
 
+    fn params(&self) -> &Self::Params {
+        &self.params
+    }
+
+    fn as_montgomery(&self) -> &Self::Integer {
+        &self.montgomery_form
+    }
+
     fn div_by_2(&self) -> Self {
         BoxedMontyForm::div_by_2(self)
     }
diff --git a/src/modular/monty_form.rs b/src/modular/monty_form.rs
index 905a7267..2e018790 100644
--- a/src/modular/monty_form.rs
+++ b/src/modular/monty_form.rs
@@ -250,6 +250,14 @@ impl<const LIMBS: usize> Monty for MontyForm<LIMBS> {
         MontyForm::one(params)
     }
 
+    fn params(&self) -> &Self::Params {
+        &self.params
+    }
+
+    fn as_montgomery(&self) -> &Self::Integer {
+        &self.montgomery_form
+    }
+
     fn div_by_2(&self) -> Self {
         MontyForm::div_by_2(self)
     }
diff --git a/src/traits.rs b/src/traits.rs
index cff9782c..410f3c68 100644
--- a/src/traits.rs
+++ b/src/traits.rs
@@ -158,6 +158,14 @@ pub trait Integer:
     /// The value `1`.
     fn one() -> Self;
 
+    /// The value `1` with the same precision as `other`.
+    fn one_like(other: &Self) -> Self {
+        Self::from_limb_like(Limb::ONE, other)
+    }
+
+    /// Returns an integer with the first limb set to `limb`, and the same precision as `other`.
+    fn from_limb_like(limb: Limb, other: &Self) -> Self;
+
     /// Number of limbs in this integer.
     fn nlimbs(&self) -> usize;
 
@@ -244,6 +252,16 @@ pub trait Zero: ConstantTimeEq + Sized {
     fn set_zero(&mut self) {
         *self = Zero::zero();
     }
+
+    /// Return the value `0` with the same precision as `other`.
+    fn zero_like(other: &Self) -> Self
+    where
+        Self: Clone,
+    {
+        let mut ret = other.clone();
+        ret.set_zero();
+        ret
+    }
 }
 
 /// Trait for associating a constant representing zero.
@@ -787,6 +805,12 @@ pub trait Monty:
     /// Returns one in this representation.
     fn one(params: Self::Params) -> Self;
 
+    /// Returns the parameter struct used to initialize this object.
+    fn params(&self) -> &Self::Params;
+
+    /// Access the value in Montgomery form.
+    fn as_montgomery(&self) -> &Self::Integer;
+
     /// Performs division by 2, that is returns `x` such that `x + x = self`.
     fn div_by_2(&self) -> Self;
 }
diff --git a/src/uint.rs b/src/uint.rs
index 7ab7ead5..e483e0b5 100644
--- a/src/uint.rs
+++ b/src/uint.rs
@@ -249,6 +249,10 @@ impl<const LIMBS: usize> Integer for Uint<LIMBS> {
         Self::ONE
     }
 
+    fn from_limb_like(limb: Limb, _other: &Self) -> Self {
+        Self::from(limb)
+    }
+
     fn nlimbs(&self) -> usize {
         Self::LIMBS
     }
diff --git a/src/uint/boxed.rs b/src/uint/boxed.rs
index d8b50647..9e0935e0 100644
--- a/src/uint/boxed.rs
+++ b/src/uint/boxed.rs
@@ -307,6 +307,12 @@ impl Integer for BoxedUint {
         Self::one()
     }
 
+    fn from_limb_like(limb: Limb, other: &Self) -> Self {
+        let mut ret = Self::zero_with_precision(other.bits_precision());
+        ret.limbs[0] = limb;
+        ret
+    }
+
     fn nlimbs(&self) -> usize {
         self.nlimbs()
     }