DO NOT MERGE chore: Refactor

src/bignum.nr

@@ -2,12 +2,12 @@ use crate::utils::map::map;
 
 use crate::params::BigNumParamsGetter;
 use crate::constrained_ops::{
-    derive_from_seed, conditional_select, assert_is_not_equal, eq, validate_in_field, validate_in_range, neg, add, sub,
-    mul, div, udiv_mod, udiv, umod
+    derive_from_seed, conditional_select, assert_is_not_equal, eq, validate_in_field, validate_in_range,
+    neg, add, sub, mul, div, udiv_mod, udiv, umod
 };
 use crate::unconstrained_ops::{
     __derive_from_seed, __eq, __is_zero, __neg, __add, __sub, __mul, __div, __udiv_mod, __invmod, __pow,
-    __batch_invert, __batch_invert_slice
+    __batch_invert, __batch_invert_slice, __tonelli_shanks_sqrt
 };
 use crate::expressions::{__compute_quadratic_expression, evaluate_quadratic_expression};
 use crate::serialization::{from_be_bytes, to_le_bytes};
@@ -51,6 +51,8 @@ pub(crate) trait BigNumTrait<let N: u32> {
     unconstrained fn __batch_invert<let M: u32>(to_invert: [Self; M]) -> [Self; M];
     unconstrained fn __batch_invert_slice<let M: u32>(to_invert: [Self]) -> [Self];
 
+    unconstrained fn __tonelli_shanks_sqrt(self) -> std::option::Option<Self>;
+
     unconstrained fn __compute_quadratic_expression<let LHS_N: u32, let RHS_N: u32, let NUM_PRODUCTS: u32, let ADD_N: u32>(
         lhs: [[Self; LHS_N]; NUM_PRODUCTS],
         lhs_flags: [[bool; LHS_N]; NUM_PRODUCTS],
@@ -216,6 +218,19 @@ impl<let N: u32, Params> BigNumTrait<N> for BigNum<N, Params> where Params: BigN
         __batch_invert_slice(params, x.map(|bn| Self::get_limbs(bn))).map(|limbs| Self { limbs })
     }
 
+    unconstrained fn __tonelli_shanks_sqrt(self) -> std::option::Option<Self> {
+        let params = Params::get_params();
+        let maybe_limbs = unsafe {
+            __tonelli_shanks_sqrt(params, self.limbs)
+        };
+        let result: std::option::Option<Self> = if maybe_limbs.is_some() {
+            std::option::Option::some(Self { limbs: maybe_limbs.unwrap_unchecked() })
+        } else {
+            std::option::Option::none()
+        };
+        result
+    }
+
     unconstrained fn __compute_quadratic_expression<let LHS_N: u32, let RHS_N: u32, let NUM_PRODUCTS: u32, let ADD_N: u32>(
         lhs_terms: [[Self; LHS_N]; NUM_PRODUCTS],
         lhs_flags: [[bool; LHS_N]; NUM_PRODUCTS],

src/runtime_bignum.nr

@@ -3,12 +3,12 @@ use crate::utils::map::map;
 
 use crate::params::BigNumParams;
 use crate::constrained_ops::{
-    derive_from_seed, conditional_select, assert_is_not_equal, eq, validate_in_field, validate_in_range, neg, add, sub,
-    mul, div, udiv_mod, udiv, umod
+    derive_from_seed, conditional_select, assert_is_not_equal, eq, validate_in_field, validate_in_range,
+    neg, add, sub, mul, div, udiv_mod, udiv, umod
 };
 use crate::unconstrained_ops::{
     __derive_from_seed, __eq, __is_zero, __neg, __add, __sub, __mul, __div, __udiv_mod, __invmod, __pow,
-    __batch_invert, __batch_invert_slice
+    __batch_invert, __batch_invert_slice, __tonelli_shanks_sqrt
 };
 use crate::expressions::{__compute_quadratic_expression, evaluate_quadratic_expression};
 use crate::serialization::{from_be_bytes, to_le_bytes};
@@ -54,6 +54,8 @@ pub(crate) trait RuntimeBigNumTrait<let N: u32> {
     fn __batch_invert<let M: u32>(x: [Self; M]) -> [Self; M];
     unconstrained fn __batch_invert_slice<let M: u32>(to_invert: [Self]) -> [Self];
 
+    unconstrained fn __tonelli_shanks_sqrt(self) -> std::option::Option<Self>;
+
     fn __compute_quadratic_expression<let LHS_N: u32, let RHS_N: u32, let NUM_PRODUCTS: u32, let ADD_N: u32>(
         params: BigNumParams<N>,
         lhs_terms: [[Self; LHS_N]; NUM_PRODUCTS],
@@ -268,6 +270,19 @@ impl<let N: u32> RuntimeBigNumTrait<N> for RuntimeBigNum<N> {
         all_limbs.map(|limbs| Self { limbs, params })
     }
 
+    unconstrained fn __tonelli_shanks_sqrt(self) -> std::option::Option<Self> {
+        let params = self.params;
+        let maybe_limbs = unsafe {
+            __tonelli_shanks_sqrt(params, self.limbs)
+        };
+        let result: std::option::Option<Self> = if maybe_limbs.is_some() {
+            std::option::Option::some(Self { limbs: maybe_limbs.unwrap_unchecked(), params })
+        } else {
+            std::option::Option::none()
+        };
+        result
+    }
+
     // UNCONSTRAINED! (Hence `__` prefix).
     fn __compute_quadratic_expression<let LHS_N: u32, let RHS_N: u32, let NUM_PRODUCTS: u32, let ADD_N: u32>(
         params: BigNumParams<N>,

src/runtime_bignum_test.nr

@@ -1,6 +1,6 @@
 use crate::utils::u60_representation::U60Repr;
 use crate::runtime_bignum::RuntimeBigNum;
-use crate::params::{ BigNumParams, BigNumParamsGetter };
+use crate::params::{BigNumParams, BigNumParamsGetter};
 
 use crate::fields::bn254Fq::BN254_Fq_Params;
 use crate::fields::secp256k1Fq::Secp256k1_Fq_Params;
@@ -410,7 +410,9 @@ fn test_div_BN() {
 #[test]
 fn test_invmod_BN() {
     let params = BN254_Fq_Params::get_params();
-    unsafe { test_invmod(params) };
+    unsafe {
+        test_invmod(params)
+    };
 }
 
 // N.B. witness generation times make these tests take ~15 minutes each! Uncomment at your peril
@@ -500,3 +502,25 @@ fn test_2048_bit_quadratic_expression() {
 
     RuntimeBigNum::evaluate_quadratic_expression(params, [[a_bn]], [[false]], [[b_bn]], [[false]], [c], [true]);
 }
+
+#[test]
+fn test_sqrt_BN() {
+    let params = BN254_Fq_Params::get_params();
+
+    let x = RuntimeBigNum { limbs: [144, 0, 0], params };
+    // let y = RuntimeBigNum { limbs: [143, 0, 0], params };
+
+    let maybe_sqrt_x = unsafe { x.__tonelli_shanks_sqrt() };
+    // let maybe_sqrt_y = unsafe { y.__tonelli_shanks_sqrt() };
+
+    // println(f"maybe_sqrt_x: {maybe_sqrt_x}");
+    // println(f"maybe_sqrt_y: {maybe_sqrt_y}");
+
+    // let sqrt_x = maybe_sqrt_x.unwrap();
+    // println(f"sqrt_x: {sqrt_x}");
+
+    assert(x == x);
+
+    // assert(maybe_sqrt_x.unwrap() == RuntimeBigNum { limbs: [12, 0, 0], params });
+    // assert(maybe_sqrt_y.is_none());
+}

src/unconstrained_helpers.nr

@@ -2,6 +2,8 @@ use crate::utils::u60_representation::U60Repr;
 use crate::utils::split_bits;
 
 use crate::params::BigNumParams as P;
+use crate::unconstrained_ops::{ __one, __eq, __neg, __add, __mul, __pow };
+
 
 /**
  * In this file:
@@ -255,153 +257,72 @@ unconstrained pub(crate) fn __barrett_reduction<let N: u32>(
     (q, r)
 }
 
+/**
+* @brief compute the log of the size of the primitive root
+* @details find the maximum value k where x^k = 1, where x = primitive root
+*          This is needed for our Tonelli-Shanks sqrt algorithm
+**/
+unconstrained pub(crate) fn __primitive_root_log_size<let N: u32>(params: P<N>) -> u32 {
+    let mut target: U60Repr<N, 2> = params.modulus_u60 - U60Repr::one();
+    let mut result: u32 = 0;
+    let modulus_bits_getter = params.modulus_bits_getter;
+    let modulus_bits = modulus_bits_getter();
+    for _ in 0..modulus_bits {
+        let lsb_is_one = (target.limbs[0] & 1) == 1;
+        if (!lsb_is_one) {
+            result += 1;
+            target.shr1();
+        } else {
+            break;
+        }
+    }
+    result
+}
+
+/**
+* @brief inner loop fn for `find_multiplive_generator`
+* @details recursive function to get around the lack of a `while` keyword
+**/
+unconstrained fn __recursively_find_multiplicative_generator<let N: u32>(
+    params: P<N>,
+    target: [Field; N],
+    p_minus_one_over_two: [Field; N]
+) -> (bool, [Field; N]) {
+    let exped = __pow(params, target, p_minus_one_over_two);
+    let one: [Field; N] = __one();
+    let neg_one = __neg(params, one);
+    let found = __eq(exped, neg_one);
+    let mut result: (bool, [Field; N]) = (found, target);
+    if (!found) {
+        let _target = unsafe {
+            __add(params, target, one)
+        };
+        result = __recursively_find_multiplicative_generator(params, _target, p_minus_one_over_two);
+    }
+    result
+}
+
+/**
+* @brief find multiplicative generator `g` where `g` is the smallest value that is not a quadratic residue
+*        i.e. smallest g where g^2 = -1
+* @note WARNING if multiplicative generator does not exist, this function will enter an infinite loop!
+**/
+unconstrained pub(crate) fn __multiplicative_generator<let N: u32>(params: P<N>) -> [Field; N] {
+    let mut target: [Field; N] = __one();
+    let p_minus_one_over_two: U60Repr<N, 2> = (params.modulus_u60 - U60Repr::one()).shr(1);
+    let p_minus_one_over_two: [Field; N] = U60Repr::into(p_minus_one_over_two);
+    let (_, target) = __recursively_find_multiplicative_generator(params, target, p_minus_one_over_two);
+    target
+}
 
-// /**
-//     * @brief compute the log of the size of the primitive root
-//     * @details find the maximum value k where x^k = 1, where x = primitive root
-//     *          This is needed for our Tonelli-Shanks sqrt algorithm
-//     **/
-// unconstrained fn primitive_root_log_size(self) -> u32 {
-//     let mut target: U60Repr<N, 2> = self.modulus_u60 - U60Repr::one();
-//     let mut result: u32 = 0;
-//     for _ in 0..Params::modulus_bits() {
-//         let lsb_is_one = (target.limbs[0] & 1) == 1;
-//         if (!lsb_is_one) {
-//             result += 1;
-//             target.shr1();
-//         } else {
-//             break;
-//         }
-//     }
-//     result
-// }
-// /**
-//     * @brief inner loop fn for `find_multiplive_generator`
-//     * @details recursive function to get around the lack of a `while` keyword
-//     **/
-// unconstrained fn recursively_find_multiplicative_generator(
-//     self,
-//     target: BigNum<N, Params>,
-//     p_minus_one_over_two: BigNum<N, Params>
-// ) -> (bool, BigNum<N, Params>) {
-//     let exped = (self.__pow(target, p_minus_one_over_two));
-//     let found = exped.__eq(self.__neg(BigNum::one()));
-//     let mut result: (bool, BigNum<N, Params>) = (found, target);
-//     if (!found) {
-//         let _target = unsafe {
-//             self.__add(target, BigNum::one())
-//         };
-//         result = self.recursively_find_multiplicative_generator(_target, p_minus_one_over_two);
-//     }
-//     result
-// }
-// /**
-//     * @brief find multiplicative generator `g` where `g` is the smallest value that is not a quadratic residue
-//     *        i.e. smallest g where g^2 = -1
-//     * @note WARNING if multiplicative generator does not exist, this function will enter an infinite loop!
-//     **/
-// unconstrained fn multiplicative_generator(self) -> BigNum<N, Params> {
-//     let mut target = BigNum::one();
-//     let p_minus_one_over_two: U60Repr<N, 2> = (self.modulus_u60 - U60Repr::one()).shr(1);
-//     let p_minus_one_over_two: BigNum<N, Params> = BigNum::from_array(U60Repr::into(p_minus_one_over_two));
-//     let (_, target) = self.recursively_find_multiplicative_generator(target, p_minus_one_over_two);
-//     target
-// }
-// unconstrained fn __tonelli_shanks_sqrt_inner_loop_check(self, t2m: BigNum<N, Params>, i: u32) -> u32 {
-//     let is_one = t2m.__eq(BigNum::one());
-//     let mut result = i;
-//     if (!is_one) {
-//         let t2m = self.__mul(t2m, t2m);
-//         let i = i + 1;
-//         result = self.__tonelli_shanks_sqrt_inner_loop_check(t2m, i);
-//     }
-//     result
-// }
-
-// /**
-// * @brief compute a modular square root using the Tonelli-Shanks algorithm
-// * @details only use for prime fields! Function may infinite loop if used for non-prime fields
-// * @note this is unconstrained fn. To constrain a square root, validate that output^2 = self
-// * TODO: create fn that constrains nonexistence of square root (i.e. find x where x^2 = -self)
-// **/
-// unconstrained fn __tonelli_shanks_sqrt(
-//     self,
-//     input: BigNum<N, Params>
-// ) -> std::option::Option<BigNum<N, Params>> {
-//     // Tonelli-shanks algorithm begins by finding a field element Q and integer S,
-//     // such that (p - 1) = Q.2^{s}
-//     // We can compute the square root of a, by considering a^{(Q + 1) / 2} = R
-//     // Once we have found such an R, we have
-//     // R^{2} = a^{Q + 1} = a^{Q}a
-//     // If a^{Q} = 1, we have found our square root.
-//     // Otherwise, we have a^{Q} = t, where t is a 2^{s-1}'th root of unity.
-//     // This is because t^{2^{s-1}} = a^{Q.2^{s-1}}.
-//     // We know that (p - 1) = Q.w^{s}, therefore t^{2^{s-1}} = a^{(p - 1) / 2}
-//     // From Euler's criterion, if a is a quadratic residue, a^{(p - 1) / 2} = 1
-//     // i.e. t^{2^{s-1}} = 1
-//     // To proceed with computing our square root, we want to transform t into a smaller subgroup,
-//     // specifically, the (s-2)'th roots of unity.
-//     // We do this by finding some value b,such that
-//     // (t.b^2)^{2^{s-2}} = 1 and R' = R.b
-//     // Finding such a b is trivial, because from Euler's criterion, we know that,
-//     // for any quadratic non-residue z, z^{(p - 1) / 2} = -1
-//     // i.e. z^{Q.2^{s-1}} = -1
-//     // => z^Q is a 2^{s-1}'th root of -1
-//     // => z^{Q^2} is a 2^{s-2}'th root of -1
-//     // Since t^{2^{s-1}} = 1, we know that t^{2^{s - 2}} = -1
-//     // => t.z^{Q^2} is a 2^{s - 2}'th root of unity.
-//     // We can iteratively transform t into ever smaller subgroups, until t = 1.
-//     // At each iteration, we need to find a new value for b, which we can obtain
-//     // by repeatedly squaring z^{Q}
-//     let one: U60Repr<N, 2> = unsafe {
-//         U60Repr::one()
-//     };
-//     let primitive_root_log_size = self.primitive_root_log_size();
-//     let mut Q = (self.modulus_u60 - one).shr(primitive_root_log_size - 1);
-//     let mut Q_minus_one_over_two = (Q - one).shr(2);
-//     let mut Q_minus_one_over_two = BigNum::from_array(U60Repr::into(Q_minus_one_over_two));
-//     let mut z = self.multiplicative_generator(); // the generator is a non-residue
-//     let mut b = self.__pow(input, Q_minus_one_over_two);
-//     let mut r = self.__mul(input, b);
-//     let mut t = self.__mul(r, b);
-//     let mut check: BigNum<N, Params> = t;
-//     for _ in 0..primitive_root_log_size - 1 {
-//         check = self.__mul(check, check);
-//     }
-//     let mut found_root = false;
-//     if (check.__eq(BigNum::one()) == false) {} else {
-//         let mut t1 = self.__pow(z, Q_minus_one_over_two);
-//         let mut t2 = self.__mul(t1, z);
-//         let mut c = self.__mul(t2, t1);
-//         let mut m: u32 = primitive_root_log_size;
-//         // tonelli shanks inner 1
-//         // (if t2m == 1) then skip
-//         // else increase i and square t2m and go again
-//         // algorithm runtime should only be max the number of bits in modulus
-//         let num_bits: u32 = Params::modulus_bits();
-//         for _ in 0..num_bits {
-//             if (t.__eq(BigNum::one())) {
-//                 found_root = true;
-//                 break;
-//             }
-//             let mut t2m = t;
-//             // while loop time
-//             let i = self.__tonelli_shanks_sqrt_inner_loop_check(t2m, 0);
-//             let mut j = m - i - 1;
-//             b = c;
-//             for _ in 0..j { // how big
-//                 if (j == 0) {
-//                     break;
-//                 }
-//                 b = self.__mul(b, b);
-//                 //j -= 1;
-//             }
-//             c = self.__mul(b, b);
-//             t = self.__mul(t, c);
-//             r = self.__mul(r, b);
-//             m = i;
-//         }
-//     }
-//     let mut result = std::option::Option { _value: r, _is_some: found_root };
-//     result
-// }
+unconstrained pub(crate) fn __tonelli_shanks_sqrt_inner_loop_check<let N: u32>(params: P<N>, t2m: [Field; N], i: u32) -> u32 {
+    let one: [Field; N] = __one();
+    let is_one = __eq(t2m, one);
+    let mut result = i;
+    if (!is_one) {
+        let t2m = __mul(params, t2m, t2m);
+        let i = i + 1;
+        result = __tonelli_shanks_sqrt_inner_loop_check(params, t2m, i);
+    }
+    result
+}