Add modular square root calculation for prime modulus ConstMontyForm

benches/const_monty.rs

@@ -193,10 +193,94 @@ fn bench_montgomery_ops<M: Measurement>(group: &mut BenchmarkGroup<'_, M>) {
     }
 }
 
+fn bench_montgomery_sqrt<M: Measurement>(group: &mut BenchmarkGroup<'_, M>) {
+    use crypto_bigint::{U256, const_prime_monty_params, modular::ConstPrimeMontyParams};
+
+    {
+        // P-256 field modulus
+        // p = 3 mod 4, s = 1, uses Shanks algorithm
+        const_prime_monty_params!(
+            P256Field,
+            U256,
+            "ffffffff00000001000000000000000000000000ffffffffffffffffffffffff"
+        );
+        assert_eq!(P256Field::PRIME_PARAMS.s().get(), 1);
+        type ConstForm = crypto_bigint::modular::ConstMontyForm<P256Field, { U256::LIMBS }>;
+
+        let mut rng = ChaCha8Rng::from_seed([7u8; 32]);
+        group.bench_function("sqrt, U256, s=1", |b| {
+            b.iter_batched(
+                || {
+                    let x =
+                        U256::random_mod_vartime(&mut rng, P256Field::PARAMS.modulus().as_nz_ref());
+                    ConstForm::new(&x)
+                },
+                |x| x.sqrt(),
+                BatchSize::SmallInput,
+            );
+        });
+    }
+
+    {
+        // P-256 scalar modulus
+        // p = 17 mod 32, s = 4
+        const_prime_monty_params!(
+            P256Scalar,
+            U256,
+            "ffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551"
+        );
+        assert_eq!(P256Scalar::PRIME_PARAMS.s().get(), 4);
+        type ConstForm = crypto_bigint::modular::ConstMontyForm<P256Scalar, { U256::LIMBS }>;
+
+        let mut rng = ChaCha8Rng::from_seed([7u8; 32]);
+        group.bench_function("sqrt, U256, s=4", |b| {
+            b.iter_batched(
+                || {
+                    let x = U256::random_mod_vartime(
+                        &mut rng,
+                        P256Scalar::PARAMS.modulus().as_nz_ref(),
+                    );
+                    ConstForm::new(&x)
+                },
+                |x| x.sqrt(),
+                BatchSize::SmallInput,
+            );
+        });
+    }
+
+    {
+        // K-256 scalar modulus
+        // s = 6, uses Tonelli-Shanks
+        const_prime_monty_params!(
+            K256Scalar,
+            U256,
+            "FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141"
+        );
+        assert_eq!(K256Scalar::PRIME_PARAMS.s().get(), 6);
+        type ConstForm = crypto_bigint::modular::ConstMontyForm<K256Scalar, { U256::LIMBS }>;
+
+        let mut rng = ChaCha8Rng::from_seed([7u8; 32]);
+        group.bench_function("sqrt, U256, s=6", |b| {
+            b.iter_batched(
+                || {
+                    let x = U256::random_mod_vartime(
+                        &mut rng,
+                        K256Scalar::PARAMS.modulus().as_nz_ref(),
+                    );
+                    ConstForm::new(&x)
+                },
+                |x| x.sqrt(),
+                BatchSize::SmallInput,
+            );
+        });
+    }
+}
+
 fn bench_montgomery(c: &mut Criterion) {
     let mut group = c.benchmark_group("Const Montgomery arithmetic");
     bench_montgomery_conversion(&mut group);
     bench_montgomery_ops(&mut group);
+    bench_montgomery_sqrt(&mut group);
     group.finish();
 }
 

src/modular.rs

@@ -28,16 +28,19 @@ mod div_by_2;
 mod monty_params;
 mod mul;
 mod pow;
+mod prime_params;
 pub(crate) mod safegcd;
+mod sqrt;
 mod sub;
 
 #[cfg(feature = "alloc")]
 pub(crate) mod boxed_monty_form;
 
 pub use self::{
-    const_monty_form::{ConstMontyForm, ConstMontyParams},
+    const_monty_form::{ConstMontyForm, ConstMontyParams, ConstPrimeMontyParams},
     fixed_monty_form::FixedMontyForm,
     monty_params::{FixedMontyParams, MontyParams},
+    prime_params::PrimeParams,
 };
 
 pub(crate) use self::safegcd::SafeGcdInverter;

src/modular/const_monty_form.rs

@@ -8,10 +8,11 @@ mod mul;
 mod neg;
 mod pow;
 mod reduce;
+mod sqrt;
 mod sub;
 
 use super::{
-    FixedMontyParams, Retrieve, div_by_2::div_by_2, mul::mul_montgomery_form,
+    FixedMontyParams, PrimeParams, Retrieve, div_by_2::div_by_2, mul::mul_montgomery_form,
     reduction::montgomery_retrieve,
 };
 use crate::{ConstOne, ConstZero, CtEq, Odd, One, Uint, Zero};
@@ -47,6 +48,16 @@ pub trait ConstMontyParams<const LIMBS: usize>:
     const PARAMS: FixedMontyParams<LIMBS>;
 }
 
+/// Trait representing a prime modulus and its associated constants for converting in
+/// and out of Montgomery form.
+///
+/// To define a type which impls this trait, use the
+/// [`const_prime_monty_params!`][`crate::const_prime_monty_params`] macro.
+pub trait ConstPrimeMontyParams<const LIMBS: usize>: ConstMontyParams<LIMBS> {
+    /// Prime parameters constant.
+    const PRIME_PARAMS: PrimeParams<LIMBS>;
+}
+
 /// An integer in Montgomery form modulo `MOD`, represented using `LIMBS` limbs.
 /// The modulus is constant, so it cannot be set at runtime.
 ///

src/modular/const_monty_form/macros.rs

@@ -1,7 +1,7 @@
 //! [`ConstMontyForm`]/[`ConstMontyParams`] support macros.
 
 #[cfg(doc)]
-use crate::modular::{ConstMontyForm, ConstMontyParams};
+use crate::modular::{ConstMontyForm, ConstMontyParams, ConstPrimeMontyParams};
 
 /// Create a type representing a modulus which impls the [`ConstMontyParams`] trait with the given
 /// name, type, value (in big endian hex), and optional documentation string.
@@ -45,6 +45,52 @@ macro_rules! const_monty_params {
     };
 }
 
+/// Create a type representing a prime modulus which impls the [`ConstPrimeMontyParams`]
+/// trait with the given name, type, value (in big endian hex), and optional documentation
+/// string.
+///
+/// # Usage
+///
+/// ```
+/// use crypto_bigint::{U256, const_prime_monty_params};
+///
+/// const_prime_monty_params!(
+///     MyModulus,
+///     U256,
+///     "73eda753299d7d483339d80809a1d80553bda402fffe5bfeffffffff00000001",
+///     "Docs for my modulus"
+/// );
+/// ```
+///
+/// The modulus _must_ be odd and prime, or this will panic.
+#[macro_export]
+macro_rules! const_prime_monty_params {
+    ($name:ident, $uint_type:ty, $value:expr) => {
+        $crate::const_prime_monty_params!(
+            $name,
+            $uint_type,
+            $value,
+            "Modulus which impls `ConstPrimeMontyParams`"
+        );
+    };
+    ($name:ident, $uint_type:ty, $value:expr, $doc:expr) => {
+        $crate::const_monty_params!(
+            $name,
+            $uint_type,
+            $value,
+            "Modulus which impls `ConstPrimeMontyParams`"
+        );
+
+        impl $crate::modular::ConstPrimeMontyParams<{ <$uint_type>::LIMBS }> for $name {
+            const PRIME_PARAMS: $crate::modular::PrimeParams<{ <$uint_type>::LIMBS }> =
+                $crate::modular::PrimeParams::new_vartime(
+                    &<$name as $crate::modular::ConstMontyParams<{ <$uint_type>::LIMBS }>>::PARAMS,
+                )
+                .expect("cannot derive prime parameters");
+        }
+    };
+}
+
 /// Creates a type alias to [`ConstMontyForm`] with the given [`ConstMontyParams`].
 ///
 /// # Usage

src/modular/const_monty_form/sqrt.rs

@@ -0,0 +1,59 @@
+use core::marker::PhantomData;
+
+use super::ConstPrimeMontyParams;
+use crate::{
+    CtOption,
+    modular::{ConstMontyForm, sqrt::sqrt_montgomery_form},
+};
+
+impl<const LIMBS: usize, MOD> ConstMontyForm<MOD, LIMBS>
+where
+    MOD: ConstPrimeMontyParams<LIMBS>,
+{
+    /// Compute the modular square root for `self`, if it exists.
+    #[must_use]
+    pub const fn sqrt(&self) -> CtOption<Self> {
+        let res = sqrt_montgomery_form(self.as_montgomery(), &MOD::PARAMS, &MOD::PRIME_PARAMS);
+        let is_some = res.is_some();
+        CtOption::new(
+            Self {
+                montgomery_form: *res.as_inner_unchecked(),
+                phantom: PhantomData,
+            },
+            is_some,
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        U256, const_prime_monty_params,
+        modular::{ConstMontyForm, ConstPrimeMontyParams},
+    };
+
+    #[test]
+    fn check_sqrt() {
+        // P-256 field modulus
+        // p = 3 mod 4, s = 1, uses Shanks algorithm
+        const_prime_monty_params!(
+            P256Field,
+            U256,
+            "ffffffff00000001000000000000000000000000ffffffffffffffffffffffff"
+        );
+        assert_eq!(P256Field::PRIME_PARAMS.s().get(), 1);
+        type ConstForm = ConstMontyForm<P256Field, { U256::LIMBS }>;
+
+        // four is square
+        let four_monty = ConstForm::new(&U256::from(4u32));
+        assert_eq!(
+            four_monty.sqrt().expect("ensured square"),
+            ConstForm::new(&U256::from(2u32))
+        );
+
+        // generator must be non-residue
+        let generator = U256::from_u32(P256Field::PRIME_PARAMS.generator().get());
+        let gen_monty = ConstForm::new(&generator);
+        assert!(gen_monty.sqrt().is_none().to_bool_vartime());
+    }
+}