feat: Implement hints on field_arithmetic lib

CHANGELOG.md

@@ -2,6 +2,57 @@
 
 #### Upcoming Changes
 
+* Implement hints on field_arithmetic lib[#985](https://github.com/lambdaclass/cairo-rs/pull/983)
+
+    `BuiltinHintProcessor` now supports the following hint:
+
+    ```python
+        %{
+            from starkware.python.math_utils import is_quad_residue, sqrt
+
+            def split(num: int, num_bits_shift: int = 128, length: int = 3):
+                a = []
+                for _ in range(length):
+                    a.append( num & ((1 << num_bits_shift) - 1) )
+                    num = num >> num_bits_shift
+                return tuple(a)
+
+            def pack(z, num_bits_shift: int = 128) -> int:
+                limbs = (z.d0, z.d1, z.d2)
+                return sum(limb << (num_bits_shift * i) for i, limb in enumerate(limbs))
+
+
+            generator = pack(ids.generator)
+            x = pack(ids.x)
+            p = pack(ids.p)
+
+            success_x = is_quad_residue(x, p)
+            root_x = sqrt(x, p) if success_x else None
+
+            success_gx = is_quad_residue(generator*x, p)
+            root_gx = sqrt(generator*x, p) if success_gx else None
+
+            # Check that one is 0 and the other is 1
+            if x != 0:
+                assert success_x + success_gx ==1
+
+            # `None` means that no root was found, but we need to transform these into a felt no matter what
+            if root_x == None:
+                root_x = 0
+            if root_gx == None:
+                root_gx = 0
+            ids.success_x = int(success_x)
+            split_root_x = split(root_x)
+            split_root_gx = split(root_gx)
+            ids.sqrt_x.d0 = split_root_x[0]
+            ids.sqrt_x.d1 = split_root_x[1]
+            ids.sqrt_x.d2 = split_root_x[2]
+            ids.sqrt_gx.d0 = split_root_gx[0]
+            ids.sqrt_gx.d1 = split_root_gx[1]
+            ids.sqrt_gx.d2 = split_root_gx[2]
+        %}
+    ```
+
 * Add missing hint on uint256_improvements lib [#1016](https://github.com/lambdaclass/cairo-rs/pull/1016):
 
     `BuiltinHintProcessor` now supports the following hint:
@@ -85,6 +136,7 @@
   * The new version carries an 85% reduction in execution time for ECDSA signature verification
 
 * BREAKING CHANGE: refactor `Program` to optimize `Program::clone` [#999](https://github.com/lambdaclass/cairo-rs/pull/999)
+
     * Breaking change: many fields that were (unnecessarily) public become hidden by the refactor.
 
 * BREAKING CHANGE: Add _builtin suffix to builtin names e.g.: output -> output_builtin [#1005](https://github.com/lambdaclass/cairo-rs/pull/1005)

Cargo.toml

@@ -32,9 +32,11 @@ hooks = []
 
 [dependencies]
 mimalloc = { version = "0.1.29", default-features = false, optional = true }
-num-bigint = { version = "0.4", features = ["serde"], default-features = false }
+num-bigint = { version = "0.4", features = ["serde", "rand"], default-features = false }
+rand = { version = "0.8.3", features = ["small_rng"], default-features = false }
 num-traits = { version = "0.2", default-features = false }
 num-integer = { version = "0.1.45", default-features = false }
+num-prime = {version = "0.4.3", features = ["big-int"], default-features = false }
 serde = { version = "1.0", features = ["derive"], default-features = false }
 serde_bytes = { version = "0.11.9", default-features = false, features = [
     "alloc",

cairo_programs/benchmarks/field_arithmetic_get_square_benchmark.cairo

@@ -0,0 +1,36 @@
+from starkware.cairo.common.cairo_builtins import BitwiseBuiltin
+from starkware.cairo.common.bool import TRUE
+from cairo_programs.uint384 import uint384_lib, Uint384, Uint384_expand
+from cairo_programs.uint384_extension import uint384_extension_lib
+from cairo_programs.field_arithmetic import field_arithmetic
+
+
+func run_get_square{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}(prime: Uint384, generator: Uint384, num: Uint384, iterations: felt) {
+    alloc_locals;
+    if (iterations == 0) {
+        return ();
+    }
+
+    let (square) = field_arithmetic.mul(num, num, prime); 
+
+    let (success, root_1) = field_arithmetic.get_square_root(square, prime, generator);
+    assert success = 1;
+
+    // We calculate this before in order to prevent revoked range_check_ptr reference due to branching
+    let (root_2) = uint384_lib.sub(prime, root_1);
+    let (is_first_root) = uint384_lib.eq(root_1, num);
+
+    if ( is_first_root != TRUE) {
+        assert root_2 = num;
+    }
+
+    return run_get_square(prime, generator, square, iterations -1);
+}
+
+func main{range_check_ptr: felt, bitwise_ptr: BitwiseBuiltin*}() {
+    let p = Uint384(18446744069414584321, 0, 0); // Goldilocks Prime
+    let x = Uint384(5, 0, 0);
+    let g = Uint384(7, 0, 0);
+    run_get_square(p, g, x, 100);
+    return ();
+}

cairo_programs/field_arithmetic.cairo

@@ -0,0 +1,171 @@
+// Code taken from https://github.com/NethermindEth/research-basic-Cairo-operations-big-integers/blob/fbf532651959f27037d70cd70ec6dbaf987f535c/lib/field_arithmetic.cairo
+from starkware.cairo.common.bitwise import bitwise_and, bitwise_or, bitwise_xor
+from starkware.cairo.common.cairo_builtins import BitwiseBuiltin
+from starkware.cairo.common.math import assert_in_range, assert_le, assert_nn_le, assert_not_zero
+from starkware.cairo.common.math_cmp import is_le
+from starkware.cairo.common.pow import pow
+from starkware.cairo.common.registers import get_ap, get_fp_and_pc
+from cairo_programs.uint384 import uint384_lib, Uint384, Uint384_expand, SHIFT, HALF_SHIFT
+from cairo_programs.uint384_extension import uint384_extension_lib, Uint768
+
+// Functions for operating elements in a finite field F_p (i.e. modulo a prime p), with p of at most 384 bits
+namespace field_arithmetic {
+    // Computes a * b modulo p
+    func mul{range_check_ptr}(a: Uint384, b: Uint384, p: Uint384) -> (res: Uint384) {
+        let (low: Uint384, high: Uint384) = uint384_lib.mul_d(a, b);
+        let full_mul_result: Uint768 = Uint768(low.d0, low.d1, low.d2, high.d0, high.d1, high.d2);
+        let (
+            quotient: Uint768, remainder: Uint384
+        ) = uint384_extension_lib.unsigned_div_rem_uint768_by_uint384(full_mul_result, p);
+        return (remainder,);
+    }
+
+    // Computes a**2 modulo p
+    func square{range_check_ptr}(a: Uint384, p: Uint384) -> (res: Uint384) {
+        let (low: Uint384, high: Uint384) = uint384_lib.square_e(a);
+        let full_mul_result: Uint768 = Uint768(low.d0, low.d1, low.d2, high.d0, high.d1, high.d2);
+        let (
+            quotient: Uint768, remainder: Uint384
+        ) = uint384_extension_lib.unsigned_div_rem_uint768_by_uint384(full_mul_result, p);
+        return (remainder,);
+    }
+
+    // Finds a square of x in F_p, i.e. x ≅ y**2 (mod p) for some y
+    // To do so, the following is done in a hint:
+    // 0. Assume x is not  0 mod p
+    // 1. Check if x is a square, if yes, find a square root r of it
+    // 2. If (and only if not), then gx *is* a square (for g a generator of F_p^*), so find a square root r of it
+    // 3. Check in Cairo that r**2 = x (mod p) or r**2 = gx (mod p), respectively
+    // NOTE: The function assumes that 0 <= x < p
+    func get_square_root{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}(
+        x: Uint384, p: Uint384, generator: Uint384
+    ) -> (success: felt, res: Uint384) {
+        alloc_locals;
+
+        // TODO: Create an equality function within field_arithmetic to avoid overflow bugs
+        let (is_zero) = uint384_lib.eq(x, Uint384(0, 0, 0));
+        if (is_zero == 1) {
+            return (1, Uint384(0, 0, 0));
+        }
+
+        local success_x: felt;
+        local sqrt_x: Uint384;
+        local sqrt_gx: Uint384;
+
+        // Compute square roots in a hint
+        %{
+            from starkware.python.math_utils import is_quad_residue, sqrt
+
+            def split(num: int, num_bits_shift: int = 128, length: int = 3):
+                a = []
+                for _ in range(length):
+                    a.append( num & ((1 << num_bits_shift) - 1) )
+                    num = num >> num_bits_shift
+                return tuple(a)
+
+            def pack(z, num_bits_shift: int = 128) -> int:
+                limbs = (z.d0, z.d1, z.d2)
+                return sum(limb << (num_bits_shift * i) for i, limb in enumerate(limbs))
+
+
+            generator = pack(ids.generator)
+            x = pack(ids.x)
+            p = pack(ids.p)
+
+            success_x = is_quad_residue(x, p)
+            root_x = sqrt(x, p) if success_x else None
+
+            success_gx = is_quad_residue(generator*x, p)
+            root_gx = sqrt(generator*x, p) if success_gx else None
+
+            # Check that one is 0 and the other is 1
+            if x != 0:
+                assert success_x + success_gx ==1
+
+            # `None` means that no root was found, but we need to transform these into a felt no matter what
+            if root_x == None:
+                root_x = 0
+            if root_gx == None:
+                root_gx = 0
+            ids.success_x = int(success_x)
+            split_root_x = split(root_x)
+            split_root_gx = split(root_gx)
+            ids.sqrt_x.d0 = split_root_x[0]
+            ids.sqrt_x.d1 = split_root_x[1]
+            ids.sqrt_x.d2 = split_root_x[2]
+            ids.sqrt_gx.d0 = split_root_gx[0]
+            ids.sqrt_gx.d1 = split_root_gx[1]
+            ids.sqrt_gx.d2 = split_root_gx[2]
+        %}
+
+        // Verify that the values computed in the hint are what they are supposed to be
+        let (gx: Uint384) = mul(generator, x, p);
+        if (success_x == 1) {
+	    uint384_lib.check(sqrt_x);
+	    let (is_valid) = uint384_lib.lt(sqrt_x, p);
+	    assert is_valid = 1;
+            let (sqrt_x_squared: Uint384) = mul(sqrt_x, sqrt_x, p);
+            // Note these checks may fail if the input x does not satisfy 0<= x < p
+            // TODO: Create a equality function within field_arithmetic to avoid overflow bugs
+            let (check_x) = uint384_lib.eq(x, sqrt_x_squared);
+            assert check_x = 1;
+            return (1, sqrt_x);
+        } else {
+            // In this case success_gx = 1
+	    uint384_lib.check(sqrt_gx);
+	    let (is_valid) = uint384_lib.lt(sqrt_gx, p);
+	    assert is_valid = 1;
+            let (sqrt_gx_squared: Uint384) = mul(sqrt_gx, sqrt_gx, p);
+            let (check_gx) = uint384_lib.eq(gx, sqrt_gx_squared);
+            assert check_gx = 1;
+            // No square roots were found
+            // Note that Uint384(0, 0, 0) is not a square root here, but something needs to be returned
+            return (0, Uint384(0, 0, 0));
+        }
+    }
+
+}
+
+func test_field_arithmetics_extension_operations{range_check_ptr, bitwise_ptr: BitwiseBuiltin*}() {
+    // Test get_square
+
+    //Small prime
+    let p_a = Uint384(7, 0, 0);
+    let x_a = Uint384(2, 0, 0);
+    let generator_a = Uint384(3, 0, 0);
+    let (s_a, r_a) = field_arithmetic.get_square_root(x_a, p_a, generator_a);
+    assert s_a = 1;
+
+    assert r_a.d0 = 3;
+    assert r_a.d1 = 0;
+    assert r_a.d2 = 0;
+
+    // Goldilocks Prime
+    let p_b = Uint384(18446744069414584321, 0, 0); // Goldilocks Prime
+    let x_b = Uint384(25, 0, 0);
+    let generator_b = Uint384(7, 0, 0);
+    let (s_b, r_b) = field_arithmetic.get_square_root(x_b, p_b, generator_b);
+    assert s_b = 1;
+
+    assert r_b.d0 = 5;
+    assert r_b.d1 = 0;
+    assert r_b.d2 = 0;
+
+    // Prime 2**101-99
+    let p_c = Uint384(77371252455336267181195165, 32767, 0);
+    let x_c = Uint384(96059601, 0, 0);
+    let generator_c = Uint384(3, 0, 0);
+    let (s_c, r_c) = field_arithmetic.get_square_root(x_c, p_c, generator_c);
+    assert s_c = 1;
+
+    assert r_c.d0 = 9801;
+    assert r_c.d1 = 0;
+    assert r_c.d2 = 0;
+
+    return ();
+}
+
+func main{range_check_ptr: felt, bitwise_ptr: BitwiseBuiltin*}() {
+    test_field_arithmetics_extension_operations();
+    return ();
+}