openvm_circuit/arch/testing/
mod.rs

Help
1use std::{
2    cell::RefCell,
3    iter::zip,
4    rc::Rc,
5    sync::{Arc, Mutex},
6};
7
8use openvm_circuit_primitives::var_range::{
9    SharedVariableRangeCheckerChip, VariableRangeCheckerBus,
10};
11use openvm_instructions::instruction::Instruction;
12use openvm_stark_backend::{
13    config::{StarkGenericConfig, Val},
14    engine::VerificationData,
15    interaction::BusIndex,
16    p3_field::PrimeField32,
17    p3_matrix::dense::{DenseMatrix, RowMajorMatrix},
18    prover::types::AirProofInput,
19    verifier::VerificationError,
20    AirRef, Chip,
21};
22use openvm_stark_sdk::{
23    config::{
24        baby_bear_blake3::{BabyBearBlake3Config, BabyBearBlake3Engine},
25        baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine},
26        setup_tracing_with_log_level, FriParameters,
27    },
28    engine::{StarkEngine, StarkFriEngine},
29    p3_baby_bear::BabyBear,
30};
31use program::ProgramTester;
32use rand::{rngs::StdRng, RngCore, SeedableRng};
33use tracing::Level;
34
35use super::{ExecutionBus, InstructionExecutor, SystemPort};
36use crate::{
37    arch::{ExecutionState, MemoryConfig},
38    system::{
39        memory::{
40            offline_checker::{MemoryBridge, MemoryBus},
41            MemoryController, OfflineMemory,
42        },
43        poseidon2::Poseidon2PeripheryChip,
44        program::ProgramBus,
45    },
46};
47
48pub mod execution;
49pub mod memory;
50pub mod program;
51pub mod test_adapter;
52
53pub use execution::ExecutionTester;
54pub use memory::MemoryTester;
55pub use test_adapter::TestAdapterChip;
56
57pub const EXECUTION_BUS: BusIndex = 0;
58pub const MEMORY_BUS: BusIndex = 1;
59pub const POSEIDON2_DIRECT_BUS: BusIndex = 6;
60pub const READ_INSTRUCTION_BUS: BusIndex = 8;
61pub const BITWISE_OP_LOOKUP_BUS: BusIndex = 9;
62pub const BYTE_XOR_BUS: BusIndex = 10;
63pub const RANGE_TUPLE_CHECKER_BUS: BusIndex = 11;
64pub const MEMORY_MERKLE_BUS: BusIndex = 12;
65
66const RANGE_CHECKER_BUS: BusIndex = 4;
67
68pub struct VmChipTestBuilder<F: PrimeField32> {
69    pub memory: MemoryTester<F>,
70    pub execution: ExecutionTester<F>,
71    pub program: ProgramTester<F>,
72    rng: StdRng,
73    default_register: usize,
74    default_pointer: usize,
75}
76
77impl<F: PrimeField32> VmChipTestBuilder<F> {
78    pub fn new(
79        memory_controller: Rc<RefCell<MemoryController<F>>>,
80        execution_bus: ExecutionBus,
81        program_bus: ProgramBus,
82        rng: StdRng,
83    ) -> Self {
84        setup_tracing_with_log_level(Level::WARN);
85        Self {
86            memory: MemoryTester::new(memory_controller),
87            execution: ExecutionTester::new(execution_bus),
88            program: ProgramTester::new(program_bus),
89            rng,
90            default_register: 0,
91            default_pointer: 0,
92        }
93    }
94
95    // Passthrough functions from ExecutionTester and MemoryTester for better dev-ex
96    pub fn execute<E: InstructionExecutor<F>>(
97        &mut self,
98        executor: &mut E,
99        instruction: &Instruction<F>,
100    ) {
101        let initial_pc = self.next_elem_size_u32();
102        self.execute_with_pc(executor, instruction, initial_pc);
103    }
104
105    pub fn execute_with_pc<E: InstructionExecutor<F>>(
106        &mut self,
107        executor: &mut E,
108        instruction: &Instruction<F>,
109        initial_pc: u32,
110    ) {
111        let initial_state = ExecutionState {
112            pc: initial_pc,
113            timestamp: self.memory.controller.borrow().timestamp(),
114        };
115        tracing::debug!(?initial_state.timestamp);
116
117        let final_state = executor
118            .execute(
119                &mut *self.memory.controller.borrow_mut(),
120                instruction,
121                initial_state,
122            )
123            .expect("Expected the execution not to fail");
124
125        self.program.execute(instruction, &initial_state);
126        self.execution.execute(initial_state, final_state);
127    }
128
129    fn next_elem_size_u32(&mut self) -> u32 {
130        self.rng.next_u32() % (1 << (F::bits() - 2))
131    }
132
133    pub fn read_cell(&mut self, address_space: usize, pointer: usize) -> F {
134        self.memory.read_cell(address_space, pointer)
135    }
136
137    pub fn write_cell(&mut self, address_space: usize, pointer: usize, value: F) {
138        self.memory.write_cell(address_space, pointer, value);
139    }
140
141    pub fn read<const N: usize>(&mut self, address_space: usize, pointer: usize) -> [F; N] {
142        self.memory.read(address_space, pointer)
143    }
144
145    pub fn write<const N: usize>(&mut self, address_space: usize, pointer: usize, value: [F; N]) {
146        self.memory.write(address_space, pointer, value);
147    }
148
149    pub fn write_usize<const N: usize>(
150        &mut self,
151        address_space: usize,
152        pointer: usize,
153        value: [usize; N],
154    ) {
155        self.memory
156            .write(address_space, pointer, value.map(F::from_canonical_usize));
157    }
158
159    pub fn write_heap<const NUM_LIMBS: usize>(
160        &mut self,
161        register: usize,
162        pointer: usize,
163        writes: Vec<[F; NUM_LIMBS]>,
164    ) {
165        self.write(1usize, register, [F::from_canonical_usize(pointer)]);
166        for (i, &write) in writes.iter().enumerate() {
167            self.write(2usize, pointer + i * NUM_LIMBS, write);
168        }
169    }
170
171    pub fn system_port(&self) -> SystemPort {
172        SystemPort {
173            execution_bus: self.execution.bus,
174            program_bus: self.program.bus,
175            memory_bridge: self.memory_bridge(),
176        }
177    }
178
179    pub fn execution_bus(&self) -> ExecutionBus {
180        self.execution.bus
181    }
182
183    pub fn program_bus(&self) -> ProgramBus {
184        self.program.bus
185    }
186
187    pub fn memory_bus(&self) -> MemoryBus {
188        self.memory.bus
189    }
190
191    pub fn memory_controller(&self) -> Rc<RefCell<MemoryController<F>>> {
192        self.memory.controller.clone()
193    }
194
195    pub fn range_checker(&self) -> SharedVariableRangeCheckerChip {
196        self.memory.controller.borrow().range_checker.clone()
197    }
198
199    pub fn memory_bridge(&self) -> MemoryBridge {
200        self.memory.controller.borrow().memory_bridge()
201    }
202
203    pub fn address_bits(&self) -> usize {
204        self.memory.controller.borrow().mem_config.pointer_max_bits
205    }
206
207    pub fn offline_memory_mutex_arc(&self) -> Arc<Mutex<OfflineMemory<F>>> {
208        self.memory_controller().borrow().offline_memory().clone()
209    }
210
211    pub fn get_default_register(&mut self, increment: usize) -> usize {
212        self.default_register += increment;
213        self.default_register - increment
214    }
215
216    pub fn get_default_pointer(&mut self, increment: usize) -> usize {
217        self.default_pointer += increment;
218        self.default_pointer - increment
219    }
220
221    pub fn write_heap_pointer_default(
222        &mut self,
223        reg_increment: usize,
224        pointer_increment: usize,
225    ) -> (usize, usize) {
226        let register = self.get_default_register(reg_increment);
227        let pointer = self.get_default_pointer(pointer_increment);
228        self.write(1, register, pointer.to_le_bytes().map(F::from_canonical_u8));
229        (register, pointer)
230    }
231
232    pub fn write_heap_default<const NUM_LIMBS: usize>(
233        &mut self,
234        reg_increment: usize,
235        pointer_increment: usize,
236        writes: Vec<[F; NUM_LIMBS]>,
237    ) -> (usize, usize) {
238        let register = self.get_default_register(reg_increment);
239        let pointer = self.get_default_pointer(pointer_increment);
240        self.write_heap(register, pointer, writes);
241        (register, pointer)
242    }
243}
244
245// Use Blake3 as hash for faster tests.
246type TestSC = BabyBearBlake3Config;
247
248impl VmChipTestBuilder<BabyBear> {
249    pub fn build(self) -> VmChipTester<TestSC> {
250        self.memory
251            .controller
252            .borrow_mut()
253            .finalize(None::<&mut Poseidon2PeripheryChip<BabyBear>>);
254        let tester = VmChipTester {
255            memory: Some(self.memory),
256            ..Default::default()
257        };
258        let tester = tester.load(self.execution);
259        tester.load(self.program)
260    }
261    pub fn build_babybear_poseidon2(self) -> VmChipTester<BabyBearPoseidon2Config> {
262        self.memory
263            .controller
264            .borrow_mut()
265            .finalize(None::<&mut Poseidon2PeripheryChip<BabyBear>>);
266        let tester = VmChipTester {
267            memory: Some(self.memory),
268            ..Default::default()
269        };
270        let tester = tester.load(self.execution);
271        tester.load(self.program)
272    }
273}
274
275impl<F: PrimeField32> Default for VmChipTestBuilder<F> {
276    fn default() -> Self {
277        let mem_config = MemoryConfig::default();
278        let range_checker = SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(
279            RANGE_CHECKER_BUS,
280            mem_config.decomp,
281        ));
282        let memory_controller = MemoryController::with_volatile_memory(
283            MemoryBus::new(MEMORY_BUS),
284            mem_config,
285            range_checker,
286        );
287        Self {
288            memory: MemoryTester::new(Rc::new(RefCell::new(memory_controller))),
289            execution: ExecutionTester::new(ExecutionBus::new(EXECUTION_BUS)),
290            program: ProgramTester::new(ProgramBus::new(READ_INSTRUCTION_BUS)),
291            rng: StdRng::seed_from_u64(0),
292            default_register: 0,
293            default_pointer: 0,
294        }
295    }
296}
297
298pub struct VmChipTester<SC: StarkGenericConfig> {
299    pub memory: Option<MemoryTester<Val<SC>>>,
300    pub air_proof_inputs: Vec<(AirRef<SC>, AirProofInput<SC>)>,
301}
302
303impl<SC: StarkGenericConfig> Default for VmChipTester<SC> {
304    fn default() -> Self {
305        Self {
306            memory: None,
307            air_proof_inputs: vec![],
308        }
309    }
310}
311
312impl<SC: StarkGenericConfig> VmChipTester<SC>
313where
314    Val<SC>: PrimeField32,
315{
316    pub fn load<C: Chip<SC>>(mut self, chip: C) -> Self {
317        if chip.current_trace_height() > 0 {
318            let air = chip.air();
319            let air_proof_input = chip.generate_air_proof_input();
320            tracing::debug!("Generated air proof input for {}", air.name());
321            self.air_proof_inputs.push((air, air_proof_input));
322        }
323
324        self
325    }
326
327    pub fn finalize(mut self) -> Self {
328        if let Some(memory_tester) = self.memory.take() {
329            let memory_controller = memory_tester.controller.clone();
330            let range_checker = memory_controller.borrow().range_checker.clone();
331            self = self.load(memory_tester); // dummy memory interactions
332            {
333                let airs = memory_controller.borrow().airs();
334                let air_proof_inputs = Rc::try_unwrap(memory_controller)
335                    .unwrap_or_else(|_| panic!("Memory controller was not dropped"))
336                    .into_inner()
337                    .generate_air_proof_inputs();
338                self.air_proof_inputs.extend(
339                    zip(airs, air_proof_inputs).filter(|(_, input)| input.main_trace_height() > 0),
340                );
341            }
342            self = self.load(range_checker); // this must be last because other trace generation mutates its state
343        }
344        self
345    }
346
347    pub fn load_air_proof_input(
348        mut self,
349        air_proof_input: (AirRef<SC>, AirProofInput<SC>),
350    ) -> Self {
351        self.air_proof_inputs.push(air_proof_input);
352        self
353    }
354
355    pub fn load_with_custom_trace<C: Chip<SC>>(
356        mut self,
357        chip: C,
358        trace: RowMajorMatrix<Val<SC>>,
359    ) -> Self {
360        let air = chip.air();
361        let mut air_proof_input = chip.generate_air_proof_input();
362        air_proof_input.raw.common_main = Some(trace);
363        self.air_proof_inputs.push((air, air_proof_input));
364        self
365    }
366
367    pub fn load_and_prank_trace<C: Chip<SC>, P>(mut self, chip: C, modify_trace: P) -> Self
368    where
369        P: Fn(&mut DenseMatrix<Val<SC>>),
370    {
371        let air = chip.air();
372        let mut air_proof_input = chip.generate_air_proof_input();
373        let trace = air_proof_input.raw.common_main.as_mut().unwrap();
374        modify_trace(trace);
375        self.air_proof_inputs.push((air, air_proof_input));
376        self
377    }
378
379    /// Given a function to produce an engine from the max trace height,
380    /// runs a simple test on that engine
381    pub fn test<E: StarkEngine<SC>, P: Fn() -> E>(
382        &self, // do no take ownership so it's easier to prank
383        engine_provider: P,
384    ) -> Result<VerificationData<SC>, VerificationError> {
385        assert!(self.memory.is_none(), "Memory must be finalized");
386        let (airs, air_proof_inputs) = self.air_proof_inputs.iter().cloned().unzip();
387        engine_provider().run_test_impl(airs, air_proof_inputs)
388    }
389}
390
391impl VmChipTester<BabyBearPoseidon2Config> {
392    pub fn simple_test(
393        &self,
394    ) -> Result<VerificationData<BabyBearPoseidon2Config>, VerificationError> {
395        self.test(|| BabyBearPoseidon2Engine::new(FriParameters::new_for_testing(1)))
396    }
397
398    pub fn simple_test_with_expected_error(&self, expected_error: VerificationError) {
399        let msg = format!(
400            "Expected verification to fail with {:?}, but it didn't",
401            &expected_error
402        );
403        let result = self.simple_test();
404        assert_eq!(result.err(), Some(expected_error), "{}", msg);
405    }
406}
407
408impl VmChipTester<BabyBearBlake3Config> {
409    pub fn simple_test(&self) -> Result<VerificationData<BabyBearBlake3Config>, VerificationError> {
410        self.test(|| BabyBearBlake3Engine::new(FriParameters::new_for_testing(1)))
411    }
412
413    pub fn simple_test_with_expected_error(&self, expected_error: VerificationError) {
414        let msg = format!(
415            "Expected verification to fail with {:?}, but it didn't",
416            &expected_error
417        );
418        let result = self.simple_test();
419        assert_eq!(result.err(), Some(expected_error), "{}", msg);
420    }
421}
openvm_circuit/arch/testing/mod.rs

openvm_circuit/arch/testing/
mod.rs
