openvm_native_circuit/
utils.rs

1pub(crate) const CASTF_MAX_BITS: usize = 30;
2
3pub(crate) const fn const_max(a: usize, b: usize) -> usize {
4    [a, b][(a < b) as usize]
5}
6
7/// Testing framework
8#[cfg(any(test, feature = "test-utils"))]
9pub mod test_utils {
10    use std::array;
11
12    use openvm_circuit::{
13        arch::{
14            execution_mode::Segment,
15            testing::{memory::gen_pointer, TestBuilder},
16            PreflightExecutionOutput, PreflightExecutor, Streams, VirtualMachine,
17            VirtualMachineError, VmBuilder, VmExecutionConfig, VmState,
18        },
19        utils::test_system_config_without_continuations,
20    };
21    use openvm_instructions::{
22        exe::VmExe,
23        program::Program,
24        riscv::{RV32_MEMORY_AS, RV32_REGISTER_AS},
25    };
26    use openvm_native_compiler::conversion::AS;
27    use openvm_stark_backend::{
28        config::Domain, p3_commit::PolynomialSpace, p3_field::PrimeField32,
29    };
30    use openvm_stark_sdk::{
31        config::{baby_bear_poseidon2::BabyBearPoseidon2Engine, setup_tracing, FriParameters},
32        engine::StarkFriEngine,
33        p3_baby_bear::BabyBear,
34    };
35    use rand::{distributions::Standard, prelude::Distribution, rngs::StdRng, Rng};
36
37    use crate::{NativeConfig, NativeCpuBuilder, Rv32WithKernelsConfig};
38
39    // If immediate, returns (value, AS::Immediate). Otherwise, writes to native memory and returns
40    // (ptr, AS::Native). If is_imm is None, randomizes it.
41    pub fn write_native_or_imm<F: PrimeField32>(
42        tester: &mut impl TestBuilder<F>,
43        rng: &mut StdRng,
44        value: F,
45        is_imm: Option<bool>,
46    ) -> (F, usize) {
47        let is_imm = is_imm.unwrap_or(rng.gen_bool(0.5));
48        if is_imm {
49            (value, AS::Immediate as usize)
50        } else {
51            let ptr = gen_pointer(rng, 1);
52            tester.write::<1>(AS::Native as usize, ptr, [value]);
53            (F::from_canonical_usize(ptr), AS::Native as usize)
54        }
55    }
56
57    // Writes value to native memory and returns a pointer to the first element together with the
58    // value If `value` is None, randomizes it.
59    pub fn write_native_array<F: PrimeField32, const N: usize>(
60        tester: &mut impl TestBuilder<F>,
61        rng: &mut StdRng,
62        value: Option<[F; N]>,
63    ) -> ([F; N], usize)
64    where
65        Standard: Distribution<F>, // Needed for `rng.gen`
66    {
67        let value = value.unwrap_or(array::from_fn(|_| rng.gen()));
68        let ptr = gen_pointer(rng, N);
69        tester.write::<N>(AS::Native as usize, ptr, value);
70        (value, ptr)
71    }
72
73    // Besides taking in system_config, this also returns Result and the full
74    // (PreflightExecutionOutput, VirtualMachine) for more advanced testing needs.
75    #[allow(clippy::type_complexity)]
76    pub fn execute_program_with_config<E, VB>(
77        program: Program<BabyBear>,
78        input_stream: impl Into<Streams<BabyBear>>,
79        builder: VB,
80        config: VB::VmConfig,
81    ) -> Result<
82        (
83            PreflightExecutionOutput<BabyBear, <VB as VmBuilder<E>>::RecordArena>,
84            VirtualMachine<E, VB>,
85        ),
86        VirtualMachineError,
87    >
88    where
89        E: StarkFriEngine,
90        Domain<E::SC>: PolynomialSpace<Val = BabyBear>,
91        VB: VmBuilder<E, VmConfig = NativeConfig>,
92        <VB::VmConfig as VmExecutionConfig<BabyBear>>::Executor:
93            PreflightExecutor<BabyBear, VB::RecordArena>,
94    {
95        setup_tracing();
96        assert!(!config.as_ref().continuation_enabled);
97        let input = input_stream.into();
98
99        let engine = E::new(FriParameters::new_for_testing(1));
100        let exe = VmExe::new(program);
101        let (vm, _) = VirtualMachine::new_with_keygen(engine, builder, config)?;
102        let ctx = vm.build_metered_ctx(&exe);
103        let (mut segments, _) = vm
104            .metered_interpreter(&exe)?
105            .execute_metered(input.clone(), ctx)?;
106        assert_eq!(segments.len(), 1, "test only supports one segment");
107        let Segment {
108            instret_start,
109            num_insns,
110            trace_heights,
111        } = segments.pop().unwrap();
112        assert_eq!(instret_start, 0);
113        let state = vm.create_initial_state(&exe, input);
114        let mut preflight_interpreter = vm.preflight_interpreter(&exe)?;
115        let output =
116            vm.execute_preflight(&mut preflight_interpreter, state, None, &trace_heights)?;
117        assert_eq!(
118            output.to_state.instret(),
119            num_insns,
120            "metered execution insn count doesn't match preflight execution"
121        );
122        Ok((output, vm))
123    }
124
125    pub fn execute_program(
126        program: Program<BabyBear>,
127        input_stream: impl Into<Streams<BabyBear>>,
128    ) -> VmState<BabyBear> {
129        let mut config = test_native_config();
130        config.system.num_public_values = 4;
131        // we set max segment len large so it doesn't segment
132        let (output, _) = execute_program_with_config::<BabyBearPoseidon2Engine, _>(
133            program,
134            input_stream,
135            NativeCpuBuilder,
136            config,
137        )
138        .unwrap();
139        output.to_state
140    }
141
142    pub fn test_native_config() -> NativeConfig {
143        let mut system = test_system_config_without_continuations();
144        system.memory_config.addr_spaces[RV32_REGISTER_AS as usize].num_cells = 0;
145        system.memory_config.addr_spaces[RV32_MEMORY_AS as usize].num_cells = 0;
146        NativeConfig {
147            system,
148            native: Default::default(),
149        }
150    }
151
152    pub fn test_native_continuations_config() -> NativeConfig {
153        NativeConfig {
154            system: test_system_config_without_continuations().with_continuations(),
155            native: Default::default(),
156        }
157    }
158
159    pub fn test_rv32_with_kernels_config() -> Rv32WithKernelsConfig {
160        Rv32WithKernelsConfig {
161            system: test_system_config_without_continuations().with_continuations(),
162            ..Default::default()
163        }
164    }
165}