u256 sqr mod

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>) -> 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<u256> = 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.