moduli-macros

openvm-algebra-moduli-macros

Procedural macros for use in guest program to generate modular arithmetic struct with custom intrinsics for compile-time modulus.

Example

openvm_algebra_moduli_macros::moduli_declare! {
    Bls12381 { modulus = "4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559787" },
    Mod1e18 { modulus = "1000000000000000003" },
}

openvm_algebra_moduli_macros::moduli_declare! {
    Mersenne61 { modulus = "0x1fffffffffffffff" },
}

openvm_algebra_moduli_macros::moduli_init! {
    "4002409555221667393417789825735904156556882819939007885332058136124031650490837864442687629129015664037894272559787",
    "1000000000000000003",
    "0x1fffffffffffffff",
}

Full story

The crate provides two macros: moduli_declare! and moduli_init!. The signatures are:

• moduli_declare! receives comma-separated list of moduli classes descriptions. Each description looks like ModulusName { modulus = "modulus_value" }. Here ModulusName is the name of the struct, and modulus_value is the modulus value in decimal or hex format.

• moduli_init! receives comma-separated list of modulus values in decimal or hex format.

What happens under the hood:

1. The purpose of these macros is to generate a struct where the operations will be compiled into our custom assembly instructions designed for the modular extension. However, since there may be several modular arithmetic chips corresponding to different moduli, we need to find a way to generate different instructions for different structs.

2. Since we cannot preserve the information about already reserved opcodes between compilation units, we use moduli_declare! to create the structs with placeholders instead of actual instructions and then moduli_init! to replace the placeholders with actual instructions.

3. Every modulus description in moduli_declare! (say, Mod1e18 { modulus = "1000000000000000003" }) transforms into something like this:

#[derive(Clone, Eq, serde::Serialize, serde::Deserialize)]
#[repr(C, align(8))]
pub struct Mod1e18(#[serde(with = "openvm_algebra_guest::BigArray")] [u8; 32]);

extern "C" {
    fn add_extern_func_de0b6b3a7640003(rd: usize, rs1: usize, rs2: usize);
    fn sub_extern_func_de0b6b3a7640003(rd: usize, rs1: usize, rs2: usize);
    fn mul_extern_func_de0b6b3a7640003(rd: usize, rs1: usize, rs2: usize);
    fn div_extern_func_de0b6b3a7640003(rd: usize, rs1: usize, rs2: usize);
    fn is_eq_extern_func_de0b6b3a7640003(rs1: usize, rs2: usize) -> bool;
}

impl Mod1e18 {
    #[inline(always)]
    fn add_assign_impl(&mut self, other: &Self) {
        #[cfg(not(target_os = "zkvm"))]
        {
            *self = Self::from_biguint(
                (self.as_biguint() + other.as_biguint()) % Self::modulus_biguint(),
            );
        }
        #[cfg(target_os = "zkvm")]
        {
            unsafe {
                add_extern_func_de0b6b3a7640003(
                    self as *mut Self as usize,
                    self as *const Self as usize,
                    other as *const Self as usize,
                );
            }
        }
    }

    #[inline(always)]
    fn sub_assign_impl(&mut self, other: &Self) {
        #[cfg(not(target_os = "zkvm"))]
        {
            let modulus = Self::modulus_biguint();
            *self = Self::from_biguint(
                (self.as_biguint() + modulus.clone() - other.as_biguint()) % modulus,
            );
        }
        #[cfg(target_os = "zkvm")]
        {
            unsafe {
                sub_extern_func_de0b6b3a7640003(
                    self as *mut Self as usize,
                    self as *const Self as usize,
                    other as *const Self as usize,
                );
            }
        }
    }

    // ...
}

// Put trait implementations in a private module to avoid conflicts
mod algebra_impl_0 {
    use openvm_algebra_guest::IntMod;

    use super::Mod1e18;

    impl IntMod for Mod1e18 {
        impl<'a> core::ops::AddAssign<&'a Mod1e18> for Mod1e18 {
            #[inline(always)]
            fn add_assign(&mut self, other: &'a Mod1e18) {
                self.add_assign_impl(other);
            }
        }
        // ...
    }
}

Here add_extern_func_de0b6b3a7640003 is the name of the function that will be generated by the moduli_init! macro. Note that de0b6b3a7640003 is the hexadecimal representation of the modulus value 1000000000000000003. Therefore, the names of the functions corresponding to equal moduli would be the same.

4. The moduli_init! macro generates the functions with the names corresponding to the moduli values. More specifically, if 1000000000000000003 is the third value in the moduli_init! macro (that is, indexed with 2, because the enumeration starts with 0), the moduli_init! macro generates the following.

#[cfg(target_os = "zkvm")]
#[link_section = ".openvm"]
#[no_mangle]
#[used]
static OPENVM_SERIALIZED_MODULUS_2: [u8; 32] = [/* bytes of the modulus */];
#[cfg(target_os = "zkvm")]
mod openvm_intrinsics_ffi {
    fn add_extern_func_de0b6b3a7640003(rd: usize, rs1: usize, rs2: usize);
    fn sub_extern_func_de0b6b3a7640003(rd: usize, rs1: usize, rs2: usize);
    fn mul_extern_func_de0b6b3a7640003(rd: usize, rs1: usize, rs2: usize);
    fn div_extern_func_de0b6b3a7640003(rd: usize, rs1: usize, rs2: usize);
    fn is_eq_extern_func_de0b6b3a7640003(rs1: usize, rs2: usize) -> bool;
}
#[allow(non_snake_case, non_upper_case_globals)]
pub mod openvm_intrinsics_meta_do_not_type_this_by_yourself {
    // information about the bytes of all moduli
}
#[allow(non_snake_case)]
pub fn setup_2() {
    #[cfg(target_os = "zkvm")]
    {
        // send the setup instruction designed for the chip number 2
    }
}
pub fn setup_all_moduli() {
    setup_0();
    setup_1();
    setup_2();
    // setup functions for all the other moduli provided in the `moduli_init!` function
}

The setup operation (e.g., setup_2) consists of reading the value OPENVM_SERIALIZED_MODULUS_2 from memory and constraining that the read value is equal to the modulus the chip has been configured with. For each used modulus, its corresponding setup instruction must be called before all other operations -- this currently must be checked by inspecting the program code; it is not enforced by the virtual machine.

5. It follows from the above that the moduli_declare! invocations may be in multiple places in various compilation units, but all the declare!d moduli must be specified at least once in moduli_init! so that there will be no linker errors due to missing function implementations. Correspondingly, the moduli_init! macro should only be called once in the entire program (in the guest crate as the topmost compilation unit). Finally, the order of the moduli in moduli_init! has nothing to do with the moduli_declare! invocations, but it must match the order of the moduli in the chip configuration -- more specifically, in the modular extension parameters (the order of numbers in ModularExtension::supported_modulus, which is usually defined with the whole app_vm_config in the openvm.toml file. The plan is to obtain this value from the specific section of the ELF file at some point).