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::{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.random_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_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        let value = value.unwrap_or(array::from_fn(|_| F::from_u32(rng.random())));
65        let ptr = gen_pointer(rng, N);
66        tester.write::<N>(AS::Native as usize, ptr, value);
67        (value, ptr)
68    }
69
70    // Besides taking in system_config, this also returns Result and the full
71    // (PreflightExecutionOutput, VirtualMachine) for more advanced testing needs.
72    #[allow(clippy::type_complexity)]
73    pub fn execute_program_with_config<E, VB>(
74        program: Program<BabyBear>,
75        input_stream: impl Into<Streams<BabyBear>>,
76        builder: VB,
77        config: VB::VmConfig,
78    ) -> Result<
79        (
80            PreflightExecutionOutput<BabyBear, <VB as VmBuilder<E>>::RecordArena>,
81            VirtualMachine<E, VB>,
82        ),
83        VirtualMachineError,
84    >
85    where
86        E: StarkFriEngine,
87        Domain<E::SC>: PolynomialSpace<Val = BabyBear>,
88        VB: VmBuilder<E, VmConfig = NativeConfig>,
89        <VB::VmConfig as VmExecutionConfig<BabyBear>>::Executor:
90            PreflightExecutor<BabyBear, VB::RecordArena>,
91    {
92        setup_tracing();
93        assert!(!config.as_ref().continuation_enabled);
94        let input = input_stream.into();
95
96        let engine = E::new(FriParameters::new_for_testing(1));
97        let exe = VmExe::new(program);
98        let (vm, _) = VirtualMachine::new_with_keygen(engine, builder, config)?;
99        let ctx = vm.build_metered_ctx(&exe);
100        let (mut segments, _) = vm
101            .metered_interpreter(&exe)?
102            .execute_metered(input.clone(), ctx)?;
103        assert_eq!(segments.len(), 1, "test only supports one segment");
104        let Segment {
105            instret_start,
106            num_insns,
107            trace_heights,
108        } = segments.pop().unwrap();
109        assert_eq!(instret_start, 0);
110        let state = vm.create_initial_state(&exe, input);
111        let mut preflight_interpreter = vm.preflight_interpreter(&exe)?;
112        let output = vm.execute_preflight(
113            &mut preflight_interpreter,
114            state,
115            Some(num_insns),
116            &trace_heights,
117        )?;
118
119        Ok((output, vm))
120    }
121
122    pub fn execute_program(
123        program: Program<BabyBear>,
124        input_stream: impl Into<Streams<BabyBear>>,
125    ) -> VmState<BabyBear> {
126        let mut config = test_native_config();
127        config.system.num_public_values = 4;
128        // we set max segment len large so it doesn't segment
129        let (output, _) = execute_program_with_config::<BabyBearPoseidon2Engine, _>(
130            program,
131            input_stream,
132            NativeCpuBuilder,
133            config,
134        )
135        .unwrap();
136        output.to_state
137    }
138
139    pub fn test_native_config() -> NativeConfig {
140        let mut system = test_system_config_without_continuations();
141        system.memory_config.addr_spaces[RV32_REGISTER_AS as usize].num_cells = 0;
142        system.memory_config.addr_spaces[RV32_MEMORY_AS as usize].num_cells = 0;
143        NativeConfig {
144            system,
145            native: Default::default(),
146        }
147    }
148
149    pub fn test_native_continuations_config() -> NativeConfig {
150        NativeConfig {
151            system: test_system_config_without_continuations().with_continuations(),
152            native: Default::default(),
153        }
154    }
155
156    pub fn test_rv32_with_kernels_config() -> Rv32WithKernelsConfig {
157        Rv32WithKernelsConfig {
158            system: test_system_config_without_continuations().with_continuations(),
159            ..Default::default()
160        }
161    }
162}