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 (vm, _) = VirtualMachine::new_with_keygen(engine, builder, config)?;
101        let ctx = vm.build_metered_ctx();
102        let exe = VmExe::new(program);
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, num_insns,
119            "metered execution insn count doesn't match preflight execution"
120        );
121        Ok((output, vm))
122    }
123
124    pub fn execute_program(
125        program: Program<BabyBear>,
126        input_stream: impl Into<Streams<BabyBear>>,
127    ) -> VmState<BabyBear> {
128        let mut config = test_native_config();
129        config.system.num_public_values = 4;
130        // we set max segment len large so it doesn't segment
131        let (output, _) = execute_program_with_config::<BabyBearPoseidon2Engine, _>(
132            program,
133            input_stream,
134            NativeCpuBuilder,
135            config,
136        )
137        .unwrap();
138        output.to_state
139    }
140
141    pub fn test_native_config() -> NativeConfig {
142        let mut system = test_system_config_without_continuations();
143        system.memory_config.addr_spaces[RV32_REGISTER_AS as usize].num_cells = 0;
144        system.memory_config.addr_spaces[RV32_MEMORY_AS as usize].num_cells = 0;
145        NativeConfig {
146            system,
147            native: Default::default(),
148        }
149    }
150
151    pub fn test_native_continuations_config() -> NativeConfig {
152        NativeConfig {
153            system: test_system_config_without_continuations().with_continuations(),
154            native: Default::default(),
155        }
156    }
157
158    pub fn test_rv32_with_kernels_config() -> Rv32WithKernelsConfig {
159        Rv32WithKernelsConfig {
160            system: test_system_config_without_continuations().with_continuations(),
161            ..Default::default()
162        }
163    }
164}