feat: Lazily call setup function for moduli and curves on first use (#1603)

We no longer need to manually call `setup_*` at the start of main for
the moduli and curves. Instead, setup is automatically done on first use
of each modulus or curve.

Depends on #1596

---------

Co-authored-by: Jonathan Wang <[email protected]>

Author: Avaneesh-axiom

Cargo.lock

@@ -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/ecrecover/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/kitchen-sink/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(
         "\

benchmarks/guest/pairing/src/main.rs

@@ -46,13 +46,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 +80,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/algebra.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/ecc.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.

book/src/custom-extensions/pairing.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 developer convenience, in the Rust function bindings for these intrinsics, 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)`.
 

docs/specs/RISCV.md

@@ -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/algebra/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/ecc/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"))

examples/pairing/src/main.rs

@@ -27,8 +27,6 @@ openvm_algebra_complex_macros::complex_init! {
 */
 
 pub fn main() {
-    setup_all_moduli();
-    setup_all_complex_extensions();
     // ...
 }
 ```
@@ -68,24 +66,23 @@ extern "C" {
 2. Again, `complex_init!` macro implements these extern functions and defines the setup functions for the complex arithmetic struct.
 
 ```rust
+#[allow(non_snake_case)]
 #[cfg(target_os = "zkvm")]
 mod openvm_intrinsics_ffi_complex {
-    fn complex_add_extern_func_Complex(rd: usize, rs1: usize, rs2: usize) {
+    #[no_mangle]
+    extern "C" fn complex_add_extern_func_Complex(rd: usize, rs1: usize, rs2: usize) {
         // send the instructions for the corresponding complex chip
         // If this struct was `init`ed k-th, these operations will be sent to the k-th complex chip
     }
-    // implement the other functions
-}
-pub fn setup_complex_0() {
-    // send the setup instructions
-}
-pub fn setup_all_complex_extensions() {
-    setup_complex_0();
-    // call all other setup_complex_* for all the items in the moduli_init! macro
+    // .. implement the other functions
+    #[no_mangle]
+    extern "C" fn complex_setup_extern_func_Complex() {
+        // send the setup instructions
+    }
 }
 ```
 
-3. Obviously, `mod_idx` in the `complex_init!` must match the position of the corresponding modulus in the `moduli_init!` macro. The order of the items in `complex_init!` affects what `setup_complex_*` function will correspond to what complex class. Also, it **must match** the order of the moduli in the chip configuration -- more specifically, in the modular extension parameters (the order of numbers in `Fp2Extension::supported_modulus`, which is usually defined with the whole `app_vm_config` in the `openvm.toml` file). However, it again imposes the restriction that we only can invoke `complex_init!` once. Again analogous to the moduli setups, we must call `setup_complex_*` for each used complex extension before doing anything with entities of that class (or one can call `setup_all_complex_extensions` to setup all of them, if all are used).
+3. Obviously, `mod_idx` in the `complex_init!` must match the position of the corresponding modulus in the `moduli_init!` macro. The order of the items in `complex_init!` affects what `setup_complex_*` function will correspond to what complex class. Also, it **must match** the order of the moduli in the chip configuration -- more specifically, in the modular extension parameters (the order of numbers in `Fp2Extension::supported_modulus`, which is usually defined with the whole `app_vm_config` in the `openvm.toml` file). However, it again imposes the restriction that we only can invoke `complex_init!` once. Again analogous to the moduli setups, the rust bindings will automatically call `complex_setup_extern_func_*` on each complex extension on first use of its intrinsics.
 
 4. Note that, due to the nature of function names, the name of the struct used in `complex_init!` must be the same as in `complex_declare!`. To illustrate, the following code will **fail** to compile:
 
@@ -107,4 +104,4 @@ The reason is that, for example, the function `complex_add_extern_func_Bn254Fp2`
 
 5. `cargo openvm build` will automatically generate a call to `complex_init!` based on `openvm.toml`.
 Note that `openvm.toml` must list the supported moduli as pairs `(name, modulus)` where `name` is the name of the struct created by `complex_declare!` as a string (in the example at the top of this document, its `"Complex"`).
-The SDK also supports this feature.
+The SDK also supports this feature.

extensions/algebra/complex-macros/README.md

@@ -58,13 +58,15 @@ pub fn complex_declare(input: TokenStream) -> TokenStream {
         create_extern_func!(complex_sub_extern_func);
         create_extern_func!(complex_mul_extern_func);
         create_extern_func!(complex_div_extern_func);
+        create_extern_func!(complex_setup_extern_func);
 
         let result = TokenStream::from(quote::quote_spanned! { span.into() =>
             extern "C" {
                 fn #complex_add_extern_func(rd: usize, rs1: usize, rs2: usize);
                 fn #complex_sub_extern_func(rd: usize, rs1: usize, rs2: usize);
                 fn #complex_mul_extern_func(rd: usize, rs1: usize, rs2: usize);
                 fn #complex_div_extern_func(rd: usize, rs1: usize, rs2: usize);
+                fn #complex_setup_extern_func();
             }
 
 
@@ -110,6 +112,7 @@ pub fn complex_declare(input: TokenStream) -> TokenStream {
                     }
                     #[cfg(target_os = "zkvm")]
                     {
+                        Self::assert_is_setup();
                         unsafe {
                             #complex_add_extern_func(
                                 self as *mut Self as usize,
@@ -130,6 +133,7 @@ pub fn complex_declare(input: TokenStream) -> TokenStream {
                     }
                     #[cfg(target_os = "zkvm")]
                     {
+                        Self::assert_is_setup();
                         unsafe {
                             #complex_sub_extern_func(
                                 self as *mut Self as usize,
@@ -154,6 +158,7 @@ pub fn complex_declare(input: TokenStream) -> TokenStream {
                     }
                     #[cfg(target_os = "zkvm")]
                     {
+                        Self::assert_is_setup();
                         unsafe {
                             #complex_mul_extern_func(
                                 self as *mut Self as usize,
@@ -179,6 +184,7 @@ pub fn complex_declare(input: TokenStream) -> TokenStream {
                     }
                     #[cfg(target_os = "zkvm")]
                     {
+                        Self::assert_is_setup();
                         unsafe {
                             #complex_div_extern_func(
                                 self as *mut Self as usize,
@@ -199,6 +205,7 @@ pub fn complex_declare(input: TokenStream) -> TokenStream {
                     }
                     #[cfg(target_os = "zkvm")]
                     {
+                        Self::assert_is_setup();
                         let mut uninit: core::mem::MaybeUninit<Self> = core::mem::MaybeUninit::uninit();
                         unsafe {
                             #complex_add_extern_func(
@@ -222,6 +229,7 @@ pub fn complex_declare(input: TokenStream) -> TokenStream {
                     }
                     #[cfg(target_os = "zkvm")]
                     {
+                        Self::assert_is_setup();
                         let mut uninit: core::mem::MaybeUninit<Self> = core::mem::MaybeUninit::uninit();
                         unsafe {
                             #complex_sub_extern_func(
@@ -249,6 +257,7 @@ pub fn complex_declare(input: TokenStream) -> TokenStream {
                     }
                     #[cfg(target_os = "zkvm")]
                     {
+                        Self::assert_is_setup();
                         unsafe {
                             #complex_mul_extern_func(
                                 dst_ptr as usize,
@@ -270,6 +279,7 @@ pub fn complex_declare(input: TokenStream) -> TokenStream {
                     }
                     #[cfg(target_os = "zkvm")]
                     {
+                        Self::assert_is_setup();
                         let mut uninit: core::mem::MaybeUninit<Self> = core::mem::MaybeUninit::uninit();
                         unsafe {
                             #complex_div_extern_func(
@@ -281,6 +291,15 @@ pub fn complex_declare(input: TokenStream) -> TokenStream {
                         unsafe { uninit.assume_init() }
                     }
                 }
+
+                // Helper function to call the setup instruction on first use
+                fn assert_is_setup() {
+                    static is_setup: ::openvm_algebra_guest::once_cell::race::OnceBool = ::openvm_algebra_guest::once_cell::race::OnceBool::new();
+                    is_setup.get_or_init(|| {
+                        unsafe { #complex_setup_extern_func(); }
+                        true
+                    });
+                }
             }
 
             impl openvm_algebra_guest::field::ComplexConjugate for #struct_name {
@@ -527,8 +546,6 @@ pub fn complex_init(input: TokenStream) -> TokenStream {
     let MacroArgs { items } = parse_macro_input!(input as MacroArgs);
 
     let mut externs = Vec::new();
-    let mut setups = Vec::new();
-    let mut setup_all_complex_extensions = Vec::new();
 
     let span = proc_macro::Span::call_site();
 
@@ -587,22 +604,22 @@ pub fn complex_init(input: TokenStream) -> TokenStream {
             });
         }
 
-        let setup_function =
-            syn::Ident::new(&format!("setup_complex_{}", complex_idx), span.into());
+        let setup_extern_func = syn::Ident::new(
+            &format!("complex_setup_extern_func_{}", struct_name),
+            span.into(),
+        );
 
-        setup_all_complex_extensions.push(quote::quote_spanned! { span.into() =>
-            #setup_function();
-        });
-        setups.push(quote::quote_spanned! { span.into() =>
-            #[allow(non_snake_case)]
-            pub fn #setup_function() {
+        externs.push(quote::quote_spanned! { span.into() =>
+            #[no_mangle]
+            extern "C" fn #setup_extern_func() {
                 #[cfg(target_os = "zkvm")]
                 {
-                    let two_modulus_bytes = &openvm_intrinsics_meta_do_not_type_this_by_yourself::two_modular_limbs_list[openvm_intrinsics_meta_do_not_type_this_by_yourself::limb_list_borders[#mod_idx]..openvm_intrinsics_meta_do_not_type_this_by_yourself::limb_list_borders[#mod_idx + 1]];
+                    use super::openvm_intrinsics_meta_do_not_type_this_by_yourself::{two_modular_limbs_list, limb_list_borders};
+                    let two_modulus_bytes = &two_modular_limbs_list[limb_list_borders[#mod_idx]..limb_list_borders[#mod_idx + 1]];
 
                     // We are going to use the numeric representation of the `rs2` register to distinguish the chip to setup.
                     // The transpiler will transform this instruction, based on whether `rs2` is `x0` or `x1`, into a `SETUP_ADDSUB` or `SETUP_MULDIV` instruction.
-                    let mut uninit: core::mem::MaybeUninit<[u8; openvm_intrinsics_meta_do_not_type_this_by_yourself::limb_list_borders[#mod_idx + 1] - openvm_intrinsics_meta_do_not_type_this_by_yourself::limb_list_borders[#mod_idx]]> = core::mem::MaybeUninit::uninit();
+                    let mut uninit: core::mem::MaybeUninit<[u8; limb_list_borders[#mod_idx + 1] - limb_list_borders[#mod_idx]]> = core::mem::MaybeUninit::uninit();
                     openvm::platform::custom_insn_r!(
                         opcode = ::openvm_algebra_guest::OPCODE,
                         funct3 = ::openvm_algebra_guest::COMPLEX_EXT_FIELD_FUNCT3,
@@ -629,14 +646,11 @@ pub fn complex_init(input: TokenStream) -> TokenStream {
     }
 
     TokenStream::from(quote::quote_spanned! { span.into() =>
+        #[allow(non_snake_case)]
         #[cfg(target_os = "zkvm")]
         mod openvm_intrinsics_ffi_complex {
             #(#externs)*
         }
-        #(#setups)*
-        pub fn setup_all_complex_extensions() {
-            #(#setup_all_complex_extensions)*
-        }
     })
 }
 

extensions/algebra/complex-macros/src/lib.rs

@@ -12,6 +12,7 @@ openvm-algebra-moduli-macros = { workspace = true }
 openvm-algebra-complex-macros = { workspace = true }
 serde-big-array.workspace = true
 strum_macros.workspace = true
+once_cell = { workspace = true, features = ["race", "alloc"] }
 
 [target.'cfg(not(target_os = "zkvm"))'.dependencies]
 num-bigint.workspace = true

extensions/algebra/guest/Cargo.toml

@@ -68,6 +68,7 @@ mod halo2curves;
 /// Exponentiation by bytes
 mod exp_bytes;
 pub use exp_bytes::*;
+pub use once_cell;
 
 /// Division operation that is undefined behavior when the denominator is not invertible.
 pub trait DivUnsafe<Rhs = Self>: Sized {