refactor: Guest library re-organization (Part 1)

Cargo.toml

@@ -52,11 +52,9 @@ members = [
     "extensions/keccak256/circuit",
     "extensions/keccak256/transpiler",
     "extensions/keccak256/guest",
-    "extensions/keccak256/tests",
     "extensions/sha256/circuit",
     "extensions/sha256/transpiler",
     "extensions/sha256/guest",
-    "extensions/sha256/tests",
     "extensions/ecc/circuit",
     "extensions/ecc/transpiler",
     "extensions/ecc/guest",

benchmarks/guest/ecrecover/src/main.rs

@@ -24,9 +24,6 @@ openvm_ecc_guest::sw_macros::sw_init! {
 }
 
 pub fn main() {
-    setup_all_moduli();
-    setup_all_curves();
-
     let expected_address = read_vec();
     for _ in 0..5 {
         let input = read_vec();

benchmarks/guest/kitchen-sink/src/main.rs

@@ -48,10 +48,8 @@ openvm_algebra_guest::complex_macros::complex_init! {
 }
 
 pub fn main() {
-    // Setup will materialize every chip
-    setup_all_moduli();
-    setup_all_complex_extensions();
-    setup_all_curves();
+    // TODO: Since we don't explicitly call setup functions anymore, we should rewrite this test
+    // to use every declared modulus and curve to ensure that every chip is materialized.
 
     let [one, six] = [1, 6].map(Seven::from_u32);
     assert_eq!(one + six, Seven::ZERO);

benchmarks/guest/pairing/src/main.rs

@@ -22,11 +22,6 @@ openvm_algebra_guest::complex_macros::complex_init! {
 const PAIR_ELEMENT_LEN: usize = 32 * (2 + 4); // 1 G1Affine (2 Fp), 1 G2Affine (4 Fp)
 
 pub fn main() {
-    setup_all_moduli();
-    setup_all_complex_extensions();
-    // Pairing doesn't need G1Affine intrinsics, but we trigger it anyways to test the chips
-    setup_all_curves();
-
     // copied from https://github.com/bluealloy/revm/blob/9e39df5dbc5fdc98779c644629b28b8bee75794a/crates/precompile/src/bn128.rs#L395
     let input = hex::decode(
         "\

book/src/custom-extensions/algebra.md

@@ -17,6 +17,9 @@ The functional part is provided by the `openvm-algebra-guest` crate, which is a
 - `Field` trait:
   Provides constants `ZERO` and `ONE` and methods for basic arithmetic operations within a field.
 
+- `Sqrt` trait:
+    Implements square root in a field using hinting.
+
 ## Modular arithmetic
 
 To [leverage](./overview.md) compile-time known moduli for performance, you declare, initialize, and then set up the arithmetic structures:
@@ -26,11 +29,15 @@ To [leverage](./overview.md) compile-time known moduli for performance, you decl
 ```rust
 moduli_declare! {
     Bls12_381Fp { modulus = "0x1a0111ea397fe69a4b1ba7b6434bacd764774b84f38512bf6730d2a0f6b0f6241eabfffeb153ffffb9feffffffffaaab" },
-    Bn254Fp { modulus = "21888242871839275222246405745257275088696311157297823662689037894645226208583" },
+    Bn254Fp {
+        modulus = "21888242871839275222246405745257275088696311157297823662689037894645226208583", impl_field = true
+    },
 }
 ```
 
-This creates `Bls12_381Fp` and `Bn254Fp` structs, each implementing the `IntMod` trait. The modulus parameter must be a string literal in decimal or hexadecimal format.
+This creates `Bls12_381Fp` and `Bn254Fp` structs, each implementing the `IntMod` trait.
+Since `impl_field = true` is specified for `Bn254Fp`, it also implements the `Field` and `Sqrt` traits.
+The modulus parameter must be a string literal in decimal or hexadecimal format.
 
 2. **Init**: Use the `init!` macro exactly once in the final binary:
 
@@ -46,13 +53,10 @@ moduli_init! {
 
 This step enumerates the declared moduli (e.g., `0` for the first one, `1` for the second one) and sets up internal linkage so the compiler can generate the appropriate RISC-V instructions associated with each modulus.
 
-3. **Setup**: At runtime, before performing arithmetic, a setup instruction must be sent to ensure security and correctness. For the \\(i\\)-th modulus, you call `setup_<i>()` (e.g., `setup_0()` or `setup_1()`). Alternatively, `setup_all_moduli()` can be used to handle all declared moduli.
-
 **Summary**:
 
 - `moduli_declare!`: Declares modular arithmetic structures and can be done multiple times.
 - `init!`: Called once in the final binary to assign and lock in the moduli.
-- `setup_<i>()`/`setup_all_moduli()`: Ensures at runtime that the correct modulus is in use, providing a security check and finalizing the environment for safe arithmetic operations.
 
 ## Complex field extension
 
@@ -83,8 +87,6 @@ complex_init! {
 */
 ```
 
-3. **Setup**: Similar to moduli, call `setup_complex_<i>()` or `setup_all_complex_extensions()` at runtime to secure the environment.
-
 ### Config parameters
 
 For the guest program to build successfully, all used moduli must be declared in the `.toml` config file in the following format:

book/src/custom-extensions/ecc.md

@@ -54,13 +54,10 @@ sw_init! {
 */
 ```
 
-3. **Setup**: Similar to the moduli and complex extensions, runtime setup instructions ensure that the correct curve parameters are being used, guaranteeing secure operation.
-
 **Summary**:
 
 - `sw_declare!`: Declares elliptic curve structures.
 - `init!`: Initializes them once, linking them to the underlying moduli.
-- `setup_sw_<i>()`/`setup_all_curves()`: Secures runtime correctness.
 
 To use elliptic curve operations on a struct defined with `sw_declare!`, it is expected that the struct for the curve's coordinate field was defined using `moduli_declare!`. In particular, the coordinate field needs to be initialized and set up as described in the [algebra extension](./algebra.md) chapter.
 

book/src/custom-extensions/pairing.md

@@ -36,14 +36,6 @@ Additionally, we'll need to initialize our moduli and `Fp2` struct via the follo
 {{ #include ../../../examples/pairing/src/main.rs:init }}
 ```
 
-And we'll run the required setup functions at the top of the guest program's `main()` function:
-
-```rust,no_run,noplayground
-{{ #include ../../../examples/pairing/src/main.rs:setup }}
-```
-
-There are two moduli defined internally in the `Bls12_381` feature. The `moduli_init!` macro thus requires both of them to be initialized. However, we do not need the scalar field of BLS12-381 (which is at index 1), and thus we only initialize the modulus from index 0, thus we only use `setup_0()` (as opposed to `setup_all_moduli()`, which will save us some columns when generating the trace).
-
 ## Input values
 
 The inputs to the pairing check are `AffinePoint`s in \\(\mathbb{F}\_p\\) and \\(\mathbb{F}\_{p^2}\\). They can be constructed via the `AffinePoint::new` function, with the inner `Fp` and `Fp2` values constructed via various `from_...` functions.

docs/specs/ISA.md

@@ -611,9 +611,7 @@ the same format that is congruent modulo `N` to the respective operation applied
 
 For each instruction, the operand `d` is fixed to be `1` and `e` is fixed to be `2`.
 Each instruction performs block accesses with block size `4` in address space `1` and block size `N::BLOCK_SIZE` in
-address space `2`, where `N::NUM_LIMBS` is divisible by `N::BLOCK_SIZE`. Recall that `N::BLOCK_SIZE` must be a power of
-
-2.
+address space `2`, where `N::NUM_LIMBS` is divisible by `N::BLOCK_SIZE`. Recall that `N::BLOCK_SIZE` must be a power of 2.
 
 | Name                      | Operands    | Description                                                                                                                                                                                                |
 | ------------------------- | ----------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
@@ -635,6 +633,16 @@ format with each limb having `LIMB_BITS` bits.
 | ISEQMOD_RV32\<N\>       | `a,b,c,1,2` | `[a:4]_1 = [r32{0}(b): N::NUM_LIMBS]_2 == [r32{0}(c): N::NUM_LIMBS]_2 (mod N) ? 1 : 0`. Enforces that `[r32{0}(b): N::NUM_LIMBS]_2, [r32{0}(c): N::NUM_LIMBS]_2` are less than `N` and then sets the register value of `[a:4]_1` to `1` or `0` depending on whether the two big integers are equal. |
 | SETUP_ISEQMOD_RV32\<N\> | `a,b,c,1,2` | `assert([r32{0}(b): N::NUM_LIMBS]_2 == N)` in the chip that handles modular equality. For the sake of implementation convenience it also writes something (can be anything) into register value of `[a:4]_1`                                                                                        |
 
+#### Phantom Sub-Instructions
+
+
+| Name           | Discriminant | Operands      | Description                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            |
+| -------------- | ------------ | ------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| HintNonQr\<N\>  | 0x50         | `_,_,c_upper` | Use `c_upper` to determine the index of the modulus from the list of supported moduli. Reset the hint stream to equal a quadratic nonresidue modulo `N`. |
+| HintSqrt\<N\>   | 0x51         | `a,_,c_upper` | Use `c_upper` to determine the index of the modulus from the list of supported moduli. Read from memory `x = [r32{0}(a): N::NUM_LIMBS]_2`.  If `x` is a quadratic residue modulo `N`, reset the hint stream to `[1u8, 0u8, 0u8, 0u8]` followed by a square root of `x`.  If `x` is not a quadratic residue, reset the hint stream to `[0u8; 4]` followed by a square root of `x * non_qr`, where `non_qr` is the quadratic nonresidue returned by `HintNonQr<N>`. |
+
+#
+
 #### Complex Extension Field
 
 A complex extension field `Fp2` is the quadratic extension of a prime field `Fp` with irreducible polynomial `X^2 + 1`.

docs/specs/RISCV.md

@@ -132,7 +132,9 @@ generates classes `Bls12381` and `Bn254` that represent the elements of the corr
 
 ### Field Arithmetic
 
-For each created modular class, one must call a corresponding `setup_*` function once at the beginning of the program. For example, for the structs above this would be `setup_0()` and `setup_1()`. This function generates the `setup` intrinsics which are distinguished by the `rs2` operand that specifies the chip this instruction is passed to..
+For each created modular class, one must call a corresponding `setup_*` function before using the intrinsics.
+For example, for the structs above this would be `setup_0()` and `setup_1()`. This function generates the `setup` intrinsics which are distinguished by the `rs2` operand that specifies the chip this instruction is passed to..
+For convenience, each modulus's `setup_*` function is automatically called on the first use of any of its intrinsics.
 
 We use `config.mod_idx(N)` to denote the index of `N` in this list. In the list below, `idx` denotes `config.mod_idx(N)`.
 
@@ -146,6 +148,8 @@ We use `config.mod_idx(N)` to denote the index of `N` in this list. In the list
 | divmod\<N\>  | R   | 0101011     | 000    | `idx*8+3` | `[rd: N::NUM_LIMBS]_2 = [rs1: N::NUM_LIMBS]_2 / [rs2: N::NUM_LIMBS]_2 (mod N)` (undefined when `gcd([rs2: N::NUM_LIMBS]_2, N) != 1`)                                                                                                                                                                                                                                                                                        |
 | iseqmod\<N\> | R   | 0101011     | 000    | `idx*8+4` | `rd = [rs1: N::NUM_LIMBS]_2 == [rs2: N::NUM_LIMBS]_2 (mod N) ? 1 : 0`. If `rd != x0`, enforces that `[rs1: N::NUM_LIMBS]_2` and `[rs2: N::NUM_LIMBS]_2` are both less than `N` and then sets `rd` equal to boolean comparison value. If `rd = x0`, this is a no-op.                                                                                                                                                         |
 | setup\<N\>   | R   | 0101011     | 000    | `idx*8+5` | `assert([rs1: N::NUM_LIMBS]_2 == N)` in the chip defined by the register index of `rs2`. For the sake of implementation convenience it also writes an unconstrained value into `[rd: N::NUM_LIMBS]_2` if `ind(rs2) = 0,1` (for add_sub, mul_div) or it overwrites the register value of `rd` with an unconstrained value if `ind(rs2) = 2` (for iseq). If `ind(rs2) = 2`, then the instruction is **invalid** if `rd = x0`. |
+| hint_non_qr\<N\> | R   | 0101011     | 001    | `idx*8+6` | Reset the hint stream to equal `non_qr` where `non_qr` is a quadratic nonresidue modulo `N`. The same `non_qr` is returned in each execution of this instruction. `rd`, `rs1`, and `rs2` should be `x0`. |
+| hint_sqrt\<N\> | R   | 0101011     | 001    | `idx*8+7` | Read `x = [rs1: N::NUM_LIMBS]_2`. If `x` is a quadratic residue modulo `N` then reset the hint stream to `[1u0, 0u8, 0u8, 0u8]` concatenated with a square root of `x`. If `x` is not a quadratic residue, then reset the hint stream to `[0u8; 4]` concatenated with a square root of `x * non_qr` where `non_qr` is the quadratic nonresidue returned by `hint_non_qr<N>`. `rd` and `rs2` should be `x0`. |
 
 Since `funct7` is 7-bits, up to 16 moduli can be supported simultaneously. We use `idx*8` to leave some room for future expansion.
 

docs/specs/transpiler.md

@@ -186,6 +186,8 @@ Each VM extension's behavior is specified below.
 | divmod\<N\>  | DIVMOD_RV32\<N\> `ind(rd), ind(rs1), ind(rs2), 1, 2`                                                                                                                                                                          |
 | iseqmod\<N\> | ISEQMOD_RV32\<N\> `ind(rd), ind(rs1), ind(rs2), 1, 2` if `rd != x0`, otherwise PHANTOM `_, _, disc(Nop)`                                                                                                                      |
 | setup\<N\>   | SETUP_ADDSUBMOD_RV32\<N\> `ind(rd), ind(rs1), x0, 1, 2` if `ind(rs2) = 0`, SETUP_MULDIVMOD_RV32\<N\> `ind(rd), ind(rs1), x0, 1, 2` if `ind(rs2) = 1`, SETUP_ISEQMOD_RV32\<N\> `ind(rd), ind(rs1), x0, 1, 2` if `ind(rs2) = 2` |
+| hint_non_qr  | PHANTOM `0, 0, phantom_c(curve_idx, HintNonQr)`                                                                                                |
+| hint_sqrt    | PHANTOM `ind(rs1), 0, phantom_c(curve_idx, HintSqrt)`                                                                                                |
 
 #### Complex Extension Field Arithmetic
 

examples/algebra/src/main.rs

@@ -31,10 +31,6 @@ openvm_algebra_complex_macros::complex_init! {
 */
 
 pub fn main() {
-    // Since we only use an arithmetic operation with `Mod1` and not `Mod2`,
-    // we only need to call `setup_0()` here.
-    setup_0();
-    setup_all_complex_extensions();
     let a = Complex1::new(Mod1::ZERO, Mod1::from_u32(0x3b8) * Mod1::from_u32(0x100000)); // a = -i in the corresponding field
     let b = Complex2::new(Mod2::ZERO, Mod2::from_u32(1000000006)); // b = -i in the corresponding field
     assert_eq!(a.clone() * &a * &a * &a * &a, a); // a^5 = a

examples/ecc/src/main.rs

@@ -23,8 +23,6 @@ openvm_ecc_guest::sw_macros::sw_init! {
 
 // ANCHOR: main
 pub fn main() {
-    setup_all_moduli();
-    setup_all_curves();
     let x1 = Secp256k1Coord::from_u32(1);
     let y1 = Secp256k1Coord::from_le_bytes(&hex!(
         "EEA7767E580D75BC6FDD7F58D2A84C2614FB22586068DB63B346C6E60AF21842"

examples/pairing/src/main.rs

@@ -26,11 +26,6 @@ openvm_algebra_complex_macros::complex_init! {
 
 // ANCHOR: main
 pub fn main() {
-    // ANCHOR: setup
-    setup_0();
-    setup_all_complex_extensions();
-    // ANCHOR_END: setup
-
     let p0 = AffinePoint::new(
         Fp::from_be_bytes(&hex!("17f1d3a73197d7942695638c4fa9ac0fc3688c4f9774b905a14e3a3f171bac586c55e83ff97a1aeffb3af00adb22c6bb")),
         Fp::from_be_bytes(&hex!("08b3f481e3aaa0f1a09e30ed741d8ae4fcf5e095d5d00af600db18cb2c04b3edd03cc744a2888ae40caa232946c5e7e1"))

extensions/algebra/circuit/Cargo.toml

@@ -30,6 +30,7 @@ derive-new = { workspace = true }
 serde.workspace = true
 serde_with = { workspace = true }
 serde-big-array = { workspace = true }
+eyre = { workspace = true }
 
 [dev-dependencies]
 halo2curves-axiom = { workspace = true }

extensions/algebra/circuit/src/fp2_extension.rs

@@ -36,7 +36,7 @@ impl Fp2Extension {
     pub fn generate_complex_init(&self, modular_config: &ModularExtension) -> String {
         fn get_index_of_modulus(modulus: &BigUint, modular_config: &ModularExtension) -> usize {
             modular_config
-                .supported_modulus
+                .supported_moduli
                 .iter()
                 .position(|m| m == modulus)
                 .expect("Modulus used in Fp2Extension not found in ModularExtension")