From 25ac6fe121014a649bc6fe50d767d4de6e128ac4 Mon Sep 17 00:00:00 2001 From: Shramee Srivastav Date: Wed, 15 May 2024 17:48:31 +0530 Subject: [PATCH 1/3] u256 sqr mod --- src/math/src/mod_arithmetics.cairo | 49 +++++++++++++++++++++++++++++- 1 file changed, 48 insertions(+), 1 deletion(-) diff --git a/src/math/src/mod_arithmetics.cairo b/src/math/src/mod_arithmetics.cairo index d424c417..467c2c7c 100644 --- a/src/math/src/mod_arithmetics.cairo +++ b/src/math/src/mod_arithmetics.cairo @@ -1,4 +1,7 @@ -use core::integer::{u512, u512_safe_div_rem_by_u256, u256_wide_mul}; +use core::integer::{ + u512, u512_safe_div_rem_by_u256, u256_wide_mul, u128_wide_mul, u128_overflowing_add, + u128_wrapping_add +}; use core::option::OptionTrait; use core::traits::TryInto; @@ -72,6 +75,50 @@ pub fn mult_mod(a: u256, b: u256, mod_non_zero: NonZero) -> u256 { rem_u256 } +#[inline(always)] +// core::integer::u128_add_with_carry +fn u128_add_with_carry(a: u128, b: u128) -> (u128, u128) nopanic { + match u128_overflowing_add(a, b) { + Result::Ok(v) => (v, 0), + Result::Err(v) => (v, 1), + } +} + +pub fn u256_wide_sqr(a: u256) -> u512 nopanic { + let (limb1, limb0) = u128_wide_mul(a.low, a.low); + let (limb2, limb1_part) = u128_wide_mul(a.low, a.high); + let (limb1, limb1_overflow0) = u128_add_with_carry(limb1, limb1_part); + let (limb1, limb1_overflow1) = u128_add_with_carry(limb1, limb1_part); + let (limb2, limb2_overflow) = u128_add_with_carry(limb2, limb2); + let (limb3, limb2_part) = u128_wide_mul(a.high, a.high); + // No overflow since no limb4. + let limb3 = u128_wrapping_add(limb3, limb2_overflow); + let (limb2, limb2_overflow) = u128_add_with_carry(limb2, limb2_part); + // No overflow since no limb4. + let limb3 = u128_wrapping_add(limb3, limb2_overflow); + // No overflow possible in this addition since both operands are 0/1. + let limb1_overflow = u128_wrapping_add(limb1_overflow0, limb1_overflow1); + let (limb2, limb2_overflow) = u128_add_with_carry(limb2, limb1_overflow); + // No overflow since no limb4. + let limb3 = u128_wrapping_add(limb3, limb2_overflow); + u512 { limb0, limb1, limb2, limb3 } +} + +/// Function that performs modular multiplication. +/// # Arguments +/// * `a` - Left hand side of multiplication. +/// * `b` - Right hand side of multiplication. +/// * `modulo` - modulo. +/// # Returns +/// * `u256` - result of modular multiplication +#[inline(always)] +pub fn sqr_mod(a: u256, modulo: u256) -> u256 { + let mult: u512 = u256_wide_sqr(a); + let mod_non_zero: NonZero = modulo.try_into().unwrap(); + let (_, rem_u256) = u512_safe_div_rem_by_u256(mult, mod_non_zero); + rem_u256 +} + /// Function that performs modular division. /// # Arguments /// * `a` - Left hand side of division. From c76f71e040ac240b6d30ffcaaf531663629cd033 Mon Sep 17 00:00:00 2001 From: Shramee Srivastav Date: Wed, 15 May 2024 17:56:09 +0530 Subject: [PATCH 2/3] tests for mod sqr --- src/math/src/tests/mod_arithmetics_test.cairo | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/src/math/src/tests/mod_arithmetics_test.cairo b/src/math/src/tests/mod_arithmetics_test.cairo index df2a906f..9301d778 100644 --- a/src/math/src/tests/mod_arithmetics_test.cairo +++ b/src/math/src/tests/mod_arithmetics_test.cairo @@ -1,4 +1,4 @@ -use alexandria_math::mod_arithmetics::{add_mod, sub_mod, mult_mod, div_mod, pow_mod}; +use alexandria_math::mod_arithmetics::{add_mod, sub_mod, mult_mod, sqr_mod, div_mod, pow_mod}; use core::traits::TryInto; const p: u256 = @@ -115,6 +115,22 @@ fn mult_mod_2_test() { assert_eq!(mult_mod(pow_256_minus_1, 1, 2), 1, "Incorrect result"); } +#[test] +#[available_gas(500000000)] +fn sqr_mod_test() { + assert_eq!(sqr_mod(p, 2), 1, "Incorrect result"); + assert_eq!( + sqr_mod(p, pow_256_minus_1), + mult_mod(p, p, pow_256_minus_1.try_into().unwrap()), + "Incorrect result" + ); + assert_eq!( + sqr_mod(pow_256_minus_1, p), + mult_mod(pow_256_minus_1, pow_256_minus_1, p.try_into().unwrap()), + "Incorrect result" + ); +} + #[test] #[available_gas(500000000)] fn div_mod_test() { From 2bae25e96a3bf3af89a414a44b1add2615864516 Mon Sep 17 00:00:00 2001 From: Shramee Srivastav Date: Wed, 15 May 2024 23:17:11 +0530 Subject: [PATCH 3/3] use non zero modulus function sig adjusted to `mult_mod` --- src/math/src/mod_arithmetics.cairo | 3 +-- src/math/src/tests/mod_arithmetics_test.cairo | 4 ++-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/math/src/mod_arithmetics.cairo b/src/math/src/mod_arithmetics.cairo index 467c2c7c..106c1985 100644 --- a/src/math/src/mod_arithmetics.cairo +++ b/src/math/src/mod_arithmetics.cairo @@ -112,9 +112,8 @@ pub fn u256_wide_sqr(a: u256) -> u512 nopanic { /// # Returns /// * `u256` - result of modular multiplication #[inline(always)] -pub fn sqr_mod(a: u256, modulo: u256) -> u256 { +pub fn sqr_mod(a: u256, mod_non_zero: NonZero) -> u256 { let mult: u512 = u256_wide_sqr(a); - let mod_non_zero: NonZero = modulo.try_into().unwrap(); let (_, rem_u256) = u512_safe_div_rem_by_u256(mult, mod_non_zero); rem_u256 } diff --git a/src/math/src/tests/mod_arithmetics_test.cairo b/src/math/src/tests/mod_arithmetics_test.cairo index 9301d778..a8b52136 100644 --- a/src/math/src/tests/mod_arithmetics_test.cairo +++ b/src/math/src/tests/mod_arithmetics_test.cairo @@ -120,12 +120,12 @@ fn mult_mod_2_test() { fn sqr_mod_test() { assert_eq!(sqr_mod(p, 2), 1, "Incorrect result"); assert_eq!( - sqr_mod(p, pow_256_minus_1), + sqr_mod(p, pow_256_minus_1.try_into().unwrap()), mult_mod(p, p, pow_256_minus_1.try_into().unwrap()), "Incorrect result" ); assert_eq!( - sqr_mod(pow_256_minus_1, p), + sqr_mod(pow_256_minus_1, p.try_into().unwrap()), mult_mod(pow_256_minus_1, pow_256_minus_1, p.try_into().unwrap()), "Incorrect result" );