Skip to main content

sp1_core_executor/
tracing.rs

1use std::{
2    marker::PhantomData,
3    ops::{Deref, DerefMut},
4    sync::Arc,
5};
6
7use hashbrown::HashMap;
8use sp1_hypercube::air::{PublicValues, PROOF_NONCE_NUM_WORDS};
9use sp1_jit::MinimalTrace;
10use sp1_primitives::consts::PAGE_SIZE;
11
12use crate::{
13    events::{
14        AluEvent, BranchEvent, InstructionDecodeEvent, InstructionFetchEvent, IntoMemoryRecord,
15        JumpEvent, MemInstrEvent, MemoryAccessPosition, MemoryLocalEvent, MemoryReadRecord,
16        MemoryRecord, MemoryRecordEnum, MemoryWriteRecord, PageProtLocalEvent, PageProtRecord,
17        PrecompileEvent, SyscallEvent, TrapExecEvent, TrapMemInstrEvent, UTypeEvent,
18    },
19    vm::{
20        results::{
21            AluResult, BranchResult, CycleResult, EcallResult, FetchResult, JumpResult, LoadResult,
22            LoadResultSupervisor, MaybeImmediate, StoreResult, StoreResultSupervisor, TrapResult,
23            UTypeResult,
24        },
25        syscall::SyscallRuntime,
26        CoreVM,
27    },
28    ALUTypeRecord, ExecutionError, ExecutionMode, ExecutionRecord, ITypeRecord, Instruction,
29    JTypeRecord, MemoryAccessRecord, Opcode, Program, RTypeRecord, Register, SP1CoreOpts,
30    SupervisorMode, SyscallCode, TrapError, UserMode,
31};
32
33/// A RISC-V VM that uses a [`MinimalTrace`] to create a [`ExecutionRecord`].
34///
35/// The type parameter `M` determines whether page protection checks are enabled.
36pub struct TracingVM<'a, M: ExecutionMode> {
37    /// The core VM.
38    pub core: CoreVM<'a, M>,
39    /// The local memory access for the CPU.
40    pub local_memory_access: LocalMemoryAccess,
41    /// The local page prot access for the CPU.
42    pub local_page_prot_access: LocalPageProtAccess,
43    /// The local memory access for any deferred precompiles.
44    pub precompile_local_memory_access: Option<LocalMemoryAccess>,
45    /// The local page prot access for any deferred precompiles.
46    pub precompile_local_page_prot_access: Option<LocalPageProtAccess>,
47    /// Decoded instruction events.
48    pub decoded_instruction_events: HashMap<u32, InstructionDecodeEvent>,
49    /// The execution record were populating.
50    pub record: &'a mut ExecutionRecord,
51    /// Phantom data for the execution mode.
52    _mode: PhantomData<M>,
53}
54
55impl TracingVM<'_, SupervisorMode> {
56    /// Execute the program until it halts.
57    pub fn execute(&mut self) -> Result<CycleResult, ExecutionError> {
58        if self.core.is_done() {
59            return Ok(CycleResult::Done(true));
60        }
61
62        loop {
63            match self.execute_instruction()? {
64                // Continue executing the program.
65                CycleResult::Done(false) => {}
66                CycleResult::TraceEnd => {
67                    self.register_refresh();
68                    self.postprocess();
69                    return Ok(CycleResult::ShardBoundary);
70                }
71                CycleResult::Done(true) => {
72                    self.postprocess();
73                    return Ok(CycleResult::Done(true));
74                }
75                CycleResult::ShardBoundary => {
76                    unreachable!("Shard boundary should never be returned for tracing VM")
77                }
78            }
79        }
80    }
81
82    /// Execute the next instruction at the current PC.
83    fn execute_instruction(&mut self) -> Result<CycleResult, ExecutionError> {
84        let pc = self.core.pc();
85        let instruction = self.core.fetch();
86
87        let mr_record = None;
88
89        match &instruction.opcode {
90            Opcode::ADD
91            | Opcode::ADDI
92            | Opcode::SUB
93            | Opcode::XOR
94            | Opcode::OR
95            | Opcode::AND
96            | Opcode::SLL
97            | Opcode::SLLW
98            | Opcode::SRL
99            | Opcode::SRA
100            | Opcode::SRLW
101            | Opcode::SRAW
102            | Opcode::SLT
103            | Opcode::SLTU
104            | Opcode::MUL
105            | Opcode::MULHU
106            | Opcode::MULHSU
107            | Opcode::MULH
108            | Opcode::MULW
109            | Opcode::DIVU
110            | Opcode::REMU
111            | Opcode::DIV
112            | Opcode::REM
113            | Opcode::DIVW
114            | Opcode::ADDW
115            | Opcode::SUBW
116            | Opcode::DIVUW
117            | Opcode::REMUW
118            | Opcode::REMW => {
119                self.execute_alu(&instruction, mr_record.as_ref(), pc);
120            }
121            Opcode::LB
122            | Opcode::LBU
123            | Opcode::LH
124            | Opcode::LHU
125            | Opcode::LW
126            | Opcode::LWU
127            | Opcode::LD => self.execute_load(&instruction, mr_record.as_ref(), pc)?,
128            Opcode::SB | Opcode::SH | Opcode::SW | Opcode::SD => {
129                self.execute_store(&instruction, mr_record.as_ref(), pc)?;
130            }
131            Opcode::JAL | Opcode::JALR => {
132                self.execute_jump(&instruction, mr_record.as_ref(), pc);
133            }
134            Opcode::BEQ | Opcode::BNE | Opcode::BLT | Opcode::BGE | Opcode::BLTU | Opcode::BGEU => {
135                self.execute_branch(&instruction, mr_record.as_ref(), pc);
136            }
137            Opcode::LUI | Opcode::AUIPC => {
138                self.execute_utype(&instruction, mr_record.as_ref(), pc);
139            }
140            Opcode::ECALL => self.execute_ecall(&instruction, mr_record.as_ref(), pc)?,
141            Opcode::EBREAK | Opcode::UNIMP => {
142                unreachable!("Invalid opcode for `execute_instruction`: {:?}", instruction.opcode)
143            }
144        }
145
146        Ok(self.core.advance())
147    }
148}
149
150impl TracingVM<'_, SupervisorMode> {
151    /// Execute a load instruction.
152    pub fn execute_load(
153        &mut self,
154        instruction: &Instruction,
155        untrusted_instruction_record: Option<&MemoryReadRecord>,
156        pc: u64,
157    ) -> Result<(), ExecutionError> {
158        let LoadResultSupervisor { mut a, b, c, rs1, rd, addr, rr_record, rw_record, mr_record } =
159            self.core.execute_load(instruction)?;
160
161        let mem_access_record = MemoryAccessRecord {
162            a: Some(MemoryRecordEnum::Write(rw_record)),
163            b: Some(MemoryRecordEnum::Read(rr_record)),
164            c: None,
165            memory: Some(MemoryRecordEnum::Read(mr_record)),
166            untrusted_instruction: untrusted_instruction_record
167                .map(|&record| (record.into(), (record.value >> (pc % 8 * 8)) as u32)),
168        };
169
170        let op_a_0 = instruction.op_a == 0;
171        if op_a_0 {
172            a = 0;
173        }
174
175        self.local_memory_access.insert_record(rd as u64, rw_record);
176        self.local_memory_access.insert_record(rs1 as u64, rr_record);
177        self.local_memory_access.insert_record(addr & !0b111, mr_record);
178
179        if let Some(untrusted_instruction_record) = untrusted_instruction_record {
180            self.local_memory_access.insert_record(pc & !0b111, *untrusted_instruction_record);
181            self.local_page_prot_access.insert_record(
182                pc / PAGE_SIZE as u64,
183                untrusted_instruction_record.prev_page_prot_record.unwrap(),
184                self.core.clk(),
185                untrusted_instruction_record.prev_page_prot_record.unwrap().page_prot,
186            );
187        }
188
189        self.emit_mem_instr_event(instruction, a, b, c, &mem_access_record, op_a_0);
190
191        self.emit_events(
192            self.core.clk(),
193            self.core.next_pc(),
194            false,
195            false,
196            instruction,
197            &mem_access_record,
198            0,
199        );
200
201        Ok(())
202    }
203
204    /// Execute a store instruction.
205    fn execute_store(
206        &mut self,
207        instruction: &Instruction,
208        untrusted_instruction_record: Option<&MemoryReadRecord>,
209        pc: u64,
210    ) -> Result<(), ExecutionError> {
211        let StoreResultSupervisor {
212            mut a,
213            b,
214            c,
215            rs1,
216            rs2,
217            addr,
218            rs1_record,
219            rs2_record,
220            mw_record,
221        } = self.core.execute_store(instruction)?;
222
223        let mem_access_record = MemoryAccessRecord {
224            a: Some(MemoryRecordEnum::Read(rs1_record)),
225            b: Some(MemoryRecordEnum::Read(rs2_record)),
226            c: None,
227            memory: Some(MemoryRecordEnum::Write(mw_record)),
228            untrusted_instruction: untrusted_instruction_record
229                .map(|&record| (record.into(), (record.value >> (pc % 8 * 8)) as u32)),
230        };
231
232        let op_a_0 = instruction.op_a == 0;
233        if op_a_0 {
234            a = 0;
235        }
236
237        self.local_memory_access.insert_record(addr & !0b111, mw_record);
238        self.local_memory_access.insert_record(rs1 as u64, rs1_record);
239        self.local_memory_access.insert_record(rs2 as u64, rs2_record);
240        if let Some(untrusted_instruction_record) = untrusted_instruction_record {
241            self.local_memory_access.insert_record(pc & !0b111, *untrusted_instruction_record);
242            self.local_page_prot_access.insert_record(
243                pc / PAGE_SIZE as u64,
244                untrusted_instruction_record.prev_page_prot_record.unwrap(),
245                self.core.clk(),
246                untrusted_instruction_record.prev_page_prot_record.unwrap().page_prot,
247            );
248        }
249
250        self.emit_mem_instr_event(instruction, a, b, c, &mem_access_record, op_a_0);
251
252        self.emit_events(
253            self.core.clk(),
254            self.core.next_pc(),
255            false,
256            false,
257            instruction,
258            &mem_access_record,
259            0,
260        );
261
262        Ok(())
263    }
264}
265
266impl TracingVM<'_, UserMode> {
267    /// Execute the program until it halts.
268    pub fn execute(&mut self) -> Result<CycleResult, ExecutionError> {
269        if self.core.is_done() {
270            return Ok(CycleResult::Done(true));
271        }
272
273        loop {
274            match self.execute_instruction()? {
275                CycleResult::Done(false) => {}
276                CycleResult::TraceEnd => {
277                    self.register_refresh();
278                    self.postprocess();
279                    return Ok(CycleResult::ShardBoundary);
280                }
281                CycleResult::Done(true) => {
282                    self.postprocess();
283                    return Ok(CycleResult::Done(true));
284                }
285                CycleResult::ShardBoundary => {
286                    unreachable!("Shard boundary should never be returned for tracing VM")
287                }
288            }
289        }
290    }
291
292    /// Execute the next instruction at the current PC.
293    fn execute_instruction(&mut self) -> Result<CycleResult, ExecutionError> {
294        let FetchResult { instruction, mr_record, pc, error } = self.core.fetch()?;
295
296        if let Some(error) = error {
297            let trap_result = self.handle_error(error)?;
298            assert!(
299                mr_record.is_some(),
300                "if an error occurred fetching an instruction, it must be an untrusted instruction"
301            );
302            let page_prot_record = mr_record.unwrap().prev_page_prot_record.unwrap();
303            self.local_page_prot_access.insert_record(
304                pc / PAGE_SIZE as u64,
305                page_prot_record,
306                self.core.clk(),
307                page_prot_record.page_prot,
308            );
309
310            self.emit_trap_exec_event(trap_result, page_prot_record);
311            self.emit_trap_events(self.core.clk(), self.core.next_pc());
312            return Ok(self.core.advance());
313        }
314
315        if instruction.is_none() {
316            unreachable!("Fetching the next instruction failed");
317        }
318
319        // SAFETY: The instruction is guaranteed to be valid as we checked for `is_none` above.
320        let instruction = unsafe { instruction.unwrap_unchecked() };
321
322        match &instruction.opcode {
323            Opcode::ADD
324            | Opcode::ADDI
325            | Opcode::SUB
326            | Opcode::XOR
327            | Opcode::OR
328            | Opcode::AND
329            | Opcode::SLL
330            | Opcode::SLLW
331            | Opcode::SRL
332            | Opcode::SRA
333            | Opcode::SRLW
334            | Opcode::SRAW
335            | Opcode::SLT
336            | Opcode::SLTU
337            | Opcode::MUL
338            | Opcode::MULHU
339            | Opcode::MULHSU
340            | Opcode::MULH
341            | Opcode::MULW
342            | Opcode::DIVU
343            | Opcode::REMU
344            | Opcode::DIV
345            | Opcode::REM
346            | Opcode::DIVW
347            | Opcode::ADDW
348            | Opcode::SUBW
349            | Opcode::DIVUW
350            | Opcode::REMUW
351            | Opcode::REMW => {
352                self.execute_alu(&instruction, mr_record.as_ref(), pc);
353            }
354            Opcode::LB
355            | Opcode::LBU
356            | Opcode::LH
357            | Opcode::LHU
358            | Opcode::LW
359            | Opcode::LWU
360            | Opcode::LD => self.execute_load(&instruction, mr_record.as_ref(), pc)?,
361            Opcode::SB | Opcode::SH | Opcode::SW | Opcode::SD => {
362                self.execute_store(&instruction, mr_record.as_ref(), pc)?;
363            }
364            Opcode::JAL | Opcode::JALR => {
365                self.execute_jump(&instruction, mr_record.as_ref(), pc);
366            }
367            Opcode::BEQ | Opcode::BNE | Opcode::BLT | Opcode::BGE | Opcode::BLTU | Opcode::BGEU => {
368                self.execute_branch(&instruction, mr_record.as_ref(), pc);
369            }
370            Opcode::LUI | Opcode::AUIPC => {
371                self.execute_utype(&instruction, mr_record.as_ref(), pc);
372            }
373            Opcode::ECALL => self.execute_ecall(&instruction, mr_record.as_ref(), pc)?,
374            Opcode::EBREAK | Opcode::UNIMP => {
375                unreachable!("Invalid opcode for `execute_instruction`: {:?}", instruction.opcode)
376            }
377        }
378
379        Ok(self.core.advance())
380    }
381}
382
383impl TracingVM<'_, UserMode> {
384    /// Execute a load instruction.
385    ///
386    /// This method will update the local memory access for the memory read, the register read,
387    /// and the register write.
388    ///
389    /// It will also emit the memory instruction event and the events for the load instruction.
390    pub fn execute_load(
391        &mut self,
392        instruction: &Instruction,
393        untrusted_instruction_record: Option<&MemoryReadRecord>,
394        pc: u64,
395    ) -> Result<(), ExecutionError> {
396        let LoadResult { mut a, b, c, rs1, rd, addr, rr_record, rw_record, mr_record, error } =
397            self.core.execute_load(instruction)?;
398
399        let mem_access_record = MemoryAccessRecord {
400            a: Some(MemoryRecordEnum::Write(rw_record)),
401            b: Some(MemoryRecordEnum::Read(rr_record)),
402            c: None,
403            memory: Some(MemoryRecordEnum::Read(mr_record)),
404            untrusted_instruction: untrusted_instruction_record
405                .map(|&record| (record.into(), (record.value >> (pc % 8 * 8)) as u32)),
406        };
407
408        let op_a_0 = instruction.op_a == 0;
409        if op_a_0 {
410            a = 0;
411        }
412
413        self.local_memory_access.insert_record(rd as u64, rw_record);
414        self.local_memory_access.insert_record(rs1 as u64, rr_record);
415        if error.is_none() {
416            self.local_memory_access.insert_record(addr & !0b111, mr_record);
417        }
418
419        if let Some(untrusted_instruction_record) = untrusted_instruction_record {
420            self.local_memory_access.insert_record(pc & !0b111, *untrusted_instruction_record);
421            self.local_page_prot_access.insert_record(
422                pc / PAGE_SIZE as u64,
423                untrusted_instruction_record.prev_page_prot_record.unwrap(),
424                self.core.clk(),
425                untrusted_instruction_record.prev_page_prot_record.unwrap().page_prot,
426            );
427        }
428
429        self.local_page_prot_access.insert_record(
430            addr / PAGE_SIZE as u64,
431            mr_record.prev_page_prot_record.unwrap(),
432            self.core.clk() + MemoryAccessPosition::Memory as u64,
433            mr_record.prev_page_prot_record.unwrap().page_prot,
434        );
435
436        let is_trap = error.is_some();
437        if let Some(error) = error {
438            let trap_result = self.handle_error(error)?;
439            self.emit_trap_mem_instr_event(
440                instruction,
441                a,
442                b,
443                c,
444                &mem_access_record,
445                trap_result,
446                op_a_0,
447            );
448        } else {
449            self.emit_mem_instr_event(instruction, a, b, c, &mem_access_record, op_a_0);
450        }
451
452        self.emit_events(
453            self.core.clk(),
454            self.core.next_pc(),
455            is_trap,
456            false,
457            instruction,
458            &mem_access_record,
459            0,
460        );
461
462        Ok(())
463    }
464
465    /// Execute a store instruction.
466    ///
467    /// This method will update the local memory access for the memory read, the register read,
468    /// and the register write.
469    ///
470    /// It will also emit the memory instruction event and the events for the store instruction.
471    fn execute_store(
472        &mut self,
473        instruction: &Instruction,
474        untrusted_instruction_record: Option<&MemoryReadRecord>,
475        pc: u64,
476    ) -> Result<(), ExecutionError> {
477        let StoreResult { mut a, b, c, rs1, rs2, addr, rs1_record, rs2_record, mw_record, error } =
478            self.core.execute_store(instruction)?;
479
480        let mem_access_record = MemoryAccessRecord {
481            a: Some(MemoryRecordEnum::Read(rs1_record)),
482            b: Some(MemoryRecordEnum::Read(rs2_record)),
483            c: None,
484            memory: Some(MemoryRecordEnum::Write(mw_record)),
485            untrusted_instruction: untrusted_instruction_record
486                .map(|&record| (record.into(), (record.value >> (pc % 8 * 8)) as u32)),
487        };
488
489        let op_a_0 = instruction.op_a == 0;
490        if op_a_0 {
491            a = 0;
492        }
493
494        if error.is_none() {
495            self.local_memory_access.insert_record(addr & !0b111, mw_record);
496        }
497        self.local_memory_access.insert_record(rs1 as u64, rs1_record);
498        self.local_memory_access.insert_record(rs2 as u64, rs2_record);
499        if let Some(untrusted_instruction_record) = untrusted_instruction_record {
500            self.local_memory_access.insert_record(pc & !0b111, *untrusted_instruction_record);
501            self.local_page_prot_access.insert_record(
502                pc / PAGE_SIZE as u64,
503                untrusted_instruction_record.prev_page_prot_record.unwrap(),
504                self.core.clk(),
505                untrusted_instruction_record.prev_page_prot_record.unwrap().page_prot,
506            );
507        }
508
509        self.local_page_prot_access.insert_record(
510            addr / PAGE_SIZE as u64,
511            mw_record.prev_page_prot_record.unwrap(),
512            self.core.clk() + MemoryAccessPosition::Memory as u64,
513            mw_record.prev_page_prot_record.unwrap().page_prot,
514        );
515
516        let is_trap = error.is_some();
517
518        if let Some(error) = error {
519            let trap_result = self.handle_error(error)?;
520            self.emit_trap_mem_instr_event(
521                instruction,
522                a,
523                b,
524                c,
525                &mem_access_record,
526                trap_result,
527                op_a_0,
528            );
529        } else {
530            self.emit_mem_instr_event(instruction, a, b, c, &mem_access_record, op_a_0);
531        }
532
533        self.emit_events(
534            self.core.clk(),
535            self.core.next_pc(),
536            is_trap,
537            false,
538            instruction,
539            &mem_access_record,
540            0,
541        );
542
543        Ok(())
544    }
545}
546
547impl<M: ExecutionMode> TracingVM<'_, M> {
548    fn postprocess(&mut self) {
549        if self.record.last_timestamp == 0 {
550            self.record.last_timestamp = self.core.clk();
551        }
552
553        self.record.program = self.core.program.clone();
554        self.record.public_values.is_untrusted_programs_enabled = M::PAGE_PROTECTION_ENABLED as u32;
555
556        if self.record.contains_cpu() {
557            self.record.public_values.pc_start = self.record.pc_start.unwrap();
558            self.record.public_values.next_pc = self.record.next_pc;
559            self.record.public_values.exit_code = self.record.exit_code;
560            self.record.public_values.last_timestamp = self.record.last_timestamp;
561            self.record.public_values.initial_timestamp = self.record.initial_timestamp;
562        }
563
564        for (_, event) in self.local_memory_access.inner.drain() {
565            self.record.cpu_local_memory_access.push(event);
566        }
567
568        if M::PAGE_PROTECTION_ENABLED {
569            for (_, event) in self.local_page_prot_access.inner.drain() {
570                self.record.cpu_local_page_prot_access.push(event);
571            }
572
573            let decoded_events =
574                std::mem::replace(&mut self.decoded_instruction_events, HashMap::new());
575            self.record.instruction_decode_events.extend(decoded_events.into_values());
576        }
577    }
578
579    fn register_refresh(&mut self) {
580        for (addr, record) in self.core.register_refresh().into_iter().enumerate() {
581            self.local_memory_access.insert_record(addr as u64, record);
582
583            self.record.bump_memory_events.push((
584                MemoryRecordEnum::Read(record),
585                addr as u64,
586                true,
587            ));
588        }
589    }
590
591    /// Get the current registers (immutable).
592    #[must_use]
593    pub fn registers(&self) -> &[MemoryRecord; 32] {
594        self.core.registers()
595    }
596
597    /// This object is used to read and write memory in a precompile.
598    #[must_use]
599    pub fn registers_mut(&mut self) -> &mut [MemoryRecord; 32] {
600        self.core.registers_mut()
601    }
602}
603
604impl<'a, M: ExecutionMode> TracingVM<'a, M> {
605    /// Create a new full-tracing VM from a minimal trace.
606    pub fn new<T: MinimalTrace>(
607        trace: &'a T,
608        program: Arc<Program>,
609        opts: SP1CoreOpts,
610        proof_nonce: [u32; PROOF_NONCE_NUM_WORDS],
611        record: &'a mut ExecutionRecord,
612    ) -> Self {
613        record.initial_timestamp = trace.clk_start();
614
615        Self {
616            core: CoreVM::new(trace, program, opts, proof_nonce),
617            record,
618            local_memory_access: LocalMemoryAccess::default(),
619            local_page_prot_access: LocalPageProtAccess::default(),
620            precompile_local_memory_access: None,
621            precompile_local_page_prot_access: None,
622            decoded_instruction_events: HashMap::new(),
623            _mode: PhantomData,
624        }
625    }
626
627    /// Get the public values from the record.
628    #[must_use]
629    pub fn public_values(&self) -> &PublicValues<u32, u64, u64, u32> {
630        &self.record.public_values
631    }
632
633    /// Handle a trap.
634    pub fn handle_error(&mut self, e: TrapError) -> Result<TrapResult, ExecutionError> {
635        let trap_result = self.core.handle_error(e)?;
636
637        self.local_memory_access.insert_record(trap_result.context, trap_result.handler_record);
638        self.local_memory_access.insert_record(trap_result.context + 8, trap_result.code_record);
639        self.local_memory_access.insert_record(trap_result.context + 16, trap_result.pc_record);
640
641        Ok(trap_result)
642    }
643
644    /// Execute an ALU instruction and emit the events.
645    fn execute_alu(
646        &mut self,
647        instruction: &Instruction,
648        untrusted_instruction_record: Option<&MemoryReadRecord>,
649        pc: u64,
650    ) {
651        let AluResult { rd, rw_record, mut a, b, c, rs1, rs2 } = self.core.execute_alu(instruction);
652
653        if let MaybeImmediate::Register(rs2, rs2_record) = rs2 {
654            self.local_memory_access.insert_record(rs2 as u64, rs2_record);
655        }
656
657        if let MaybeImmediate::Register(rs1, rs1_record) = rs1 {
658            self.local_memory_access.insert_record(rs1 as u64, rs1_record);
659        }
660
661        self.local_memory_access.insert_record(rd as u64, rw_record);
662        if let Some(untrusted_instruction_record) = untrusted_instruction_record {
663            self.local_memory_access.insert_record(pc & !0b111, *untrusted_instruction_record);
664            self.local_page_prot_access.insert_record(
665                pc / PAGE_SIZE as u64,
666                untrusted_instruction_record.prev_page_prot_record.unwrap(),
667                self.core.clk(),
668                untrusted_instruction_record.prev_page_prot_record.unwrap().page_prot,
669            );
670        }
671
672        let mem_access_record = MemoryAccessRecord {
673            a: Some(MemoryRecordEnum::Write(rw_record)),
674            b: rs1.record().map(|r| MemoryRecordEnum::Read(*r)),
675            c: rs2.record().map(|r| MemoryRecordEnum::Read(*r)),
676            memory: None,
677            untrusted_instruction: untrusted_instruction_record
678                .map(|&record| (record.into(), (record.value >> (pc % 8 * 8)) as u32)),
679        };
680
681        let op_a_0 = instruction.op_a == 0;
682        if op_a_0 {
683            a = 0;
684        }
685
686        self.emit_events(
687            self.core.clk(),
688            self.core.next_pc(),
689            false,
690            false,
691            instruction,
692            &mem_access_record,
693            0,
694        );
695        self.emit_alu_event(instruction, a, b, c, &mem_access_record, op_a_0);
696    }
697
698    /// Execute a jump instruction and emit the events.
699    fn execute_jump(
700        &mut self,
701        instruction: &Instruction,
702        untrusted_instruction_record: Option<&MemoryReadRecord>,
703        pc: u64,
704    ) {
705        let JumpResult { mut a, b, c, rd, rd_record, rs1 } = self.core.execute_jump(instruction);
706
707        if let MaybeImmediate::Register(rs1, rs1_record) = rs1 {
708            self.local_memory_access.insert_record(rs1 as u64, rs1_record);
709        }
710
711        self.local_memory_access.insert_record(rd as u64, rd_record);
712        if let Some(untrusted_instruction_record) = untrusted_instruction_record {
713            self.local_memory_access.insert_record(pc & !0b111, *untrusted_instruction_record);
714            self.local_page_prot_access.insert_record(
715                pc / PAGE_SIZE as u64,
716                untrusted_instruction_record.prev_page_prot_record.unwrap(),
717                self.core.clk(),
718                untrusted_instruction_record.prev_page_prot_record.unwrap().page_prot,
719            );
720        }
721
722        let mem_access_record = MemoryAccessRecord {
723            a: Some(MemoryRecordEnum::Write(rd_record)),
724            b: rs1.record().map(|r| MemoryRecordEnum::Read(*r)),
725            c: None,
726            memory: None,
727            untrusted_instruction: untrusted_instruction_record
728                .map(|&record| (record.into(), (record.value >> (pc % 8 * 8)) as u32)),
729        };
730
731        let op_a_0 = instruction.op_a == 0;
732        if op_a_0 {
733            a = 0;
734        }
735
736        self.emit_events(
737            self.core.clk(),
738            self.core.next_pc(),
739            false,
740            false,
741            instruction,
742            &mem_access_record,
743            0,
744        );
745        match instruction.opcode {
746            Opcode::JAL => self.emit_jal_event(
747                instruction,
748                a,
749                b,
750                c,
751                &mem_access_record,
752                op_a_0,
753                self.core.next_pc(),
754            ),
755            Opcode::JALR => self.emit_jalr_event(
756                instruction,
757                a,
758                b,
759                c,
760                &mem_access_record,
761                op_a_0,
762                self.core.next_pc(),
763            ),
764            _ => unreachable!("Invalid opcode for `execute_jump`: {:?}", instruction.opcode),
765        }
766    }
767
768    /// Execute a branch instruction and emit the events.
769    fn execute_branch(
770        &mut self,
771        instruction: &Instruction,
772        untrusted_instruction_record: Option<&MemoryReadRecord>,
773        pc: u64,
774    ) {
775        let BranchResult { mut a, rs1, a_record, b, rs2, b_record, c } =
776            self.core.execute_branch(instruction);
777
778        self.local_memory_access.insert_record(rs2 as u64, b_record);
779        self.local_memory_access.insert_record(rs1 as u64, a_record);
780        if let Some(untrusted_instruction_record) = untrusted_instruction_record {
781            self.local_memory_access.insert_record(pc & !0b111, *untrusted_instruction_record);
782            self.local_page_prot_access.insert_record(
783                pc / PAGE_SIZE as u64,
784                untrusted_instruction_record.prev_page_prot_record.unwrap(),
785                self.core.clk(),
786                untrusted_instruction_record.prev_page_prot_record.unwrap().page_prot,
787            );
788        }
789
790        let mem_access_record = MemoryAccessRecord {
791            a: Some(MemoryRecordEnum::Read(a_record)),
792            b: Some(MemoryRecordEnum::Read(b_record)),
793            c: None,
794            memory: None,
795            untrusted_instruction: untrusted_instruction_record
796                .map(|&record| (record.into(), (record.value >> (pc % 8 * 8)) as u32)),
797        };
798
799        let op_a_0 = instruction.op_a == 0;
800        if op_a_0 {
801            a = 0;
802        }
803
804        self.emit_events(
805            self.core.clk(),
806            self.core.next_pc(),
807            false,
808            false,
809            instruction,
810            &mem_access_record,
811            0,
812        );
813        self.emit_branch_event(
814            instruction,
815            a,
816            b,
817            c,
818            &mem_access_record,
819            op_a_0,
820            self.core.next_pc(),
821        );
822    }
823
824    /// Execute a U-type instruction and emit the events.   
825    fn execute_utype(
826        &mut self,
827        instruction: &Instruction,
828        untrusted_instruction_record: Option<&MemoryReadRecord>,
829        pc: u64,
830    ) {
831        let UTypeResult { mut a, b, c, rd, rw_record } = self.core.execute_utype(instruction);
832
833        self.local_memory_access.insert_record(rd as u64, rw_record);
834        if let Some(untrusted_instruction_record) = untrusted_instruction_record {
835            self.local_memory_access.insert_record(pc & !0b111, *untrusted_instruction_record);
836            self.local_page_prot_access.insert_record(
837                pc / PAGE_SIZE as u64,
838                untrusted_instruction_record.prev_page_prot_record.unwrap(),
839                self.core.clk(),
840                untrusted_instruction_record.prev_page_prot_record.unwrap().page_prot,
841            );
842        }
843
844        let mem_access_record = MemoryAccessRecord {
845            a: Some(MemoryRecordEnum::Write(rw_record)),
846            b: None,
847            c: None,
848            memory: None,
849            untrusted_instruction: untrusted_instruction_record
850                .map(|&record| (record.into(), (record.value >> (pc % 8 * 8)) as u32)),
851        };
852
853        let op_a_0 = instruction.op_a == 0;
854        if op_a_0 {
855            a = 0;
856        }
857
858        self.emit_events(
859            self.core.clk(),
860            self.core.next_pc(),
861            false,
862            false,
863            instruction,
864            &mem_access_record,
865            0,
866        );
867        self.emit_utype_event(instruction, a, b, c, &mem_access_record, op_a_0);
868    }
869
870    /// Execute an ecall instruction and emit the events.
871    fn execute_ecall(
872        &mut self,
873        instruction: &Instruction,
874        untrusted_instruction_record: Option<&MemoryReadRecord>,
875        pc: u64,
876    ) -> Result<(), ExecutionError> {
877        let code = self.core.read_code();
878        let is_sigreturn = code == SyscallCode::SIG_RETURN;
879
880        // If the syscall is not retained, we need to track the local memory access separately.
881        //
882        // Note that the `precompile_local_memory_access` is set to `None` in the
883        // `postprocess_precompile` method.
884        if !self.core().is_retained_syscall(code) && code.should_send() == 1 {
885            self.precompile_local_memory_access = Some(LocalMemoryAccess::default());
886            if M::PAGE_PROTECTION_ENABLED {
887                self.precompile_local_page_prot_access = Some(LocalPageProtAccess::default());
888            } else {
889                self.precompile_local_page_prot_access = None;
890            }
891        } else {
892            self.precompile_local_page_prot_access = None;
893            self.precompile_local_memory_access = None;
894        }
895
896        if is_sigreturn {
897            let c_record_peek = self.core().rr_peek(Register::X11, MemoryAccessPosition::C);
898            let b_record_peek = self.core().rr_peek(Register::X10, MemoryAccessPosition::B);
899            let a_record_peek = self.core().rr_peek(Register::X5, MemoryAccessPosition::A);
900            self.local_memory_access.insert_record(Register::X11 as u64, c_record_peek);
901            self.local_memory_access.insert_record(Register::X10 as u64, b_record_peek);
902            self.local_memory_access.insert_record(Register::X5 as u64, a_record_peek);
903        }
904
905        // Actually execute the ecall.
906        let EcallResult { a: _, a_record, b, b_record, c, c_record, error, sig_return_pc_record } =
907            CoreVM::<'a>::execute_ecall(&mut PrecompileMemory::new(self), instruction, code)?;
908
909        if let Some(record) = sig_return_pc_record {
910            self.local_memory_access.insert_record(b, record);
911        }
912
913        if !is_sigreturn {
914            self.local_memory_access.insert_record(Register::X11 as u64, c_record);
915            self.local_memory_access.insert_record(Register::X10 as u64, b_record);
916            self.local_memory_access.insert_record(Register::X5 as u64, a_record);
917        }
918
919        let mem_access_record = MemoryAccessRecord {
920            a: Some(MemoryRecordEnum::Write(a_record)),
921            b: Some(MemoryRecordEnum::Read(b_record)),
922            c: Some(MemoryRecordEnum::Read(c_record)),
923            memory: None,
924            untrusted_instruction: untrusted_instruction_record
925                .map(|&record| (record.into(), (record.value >> (pc % 8 * 8)) as u32)),
926        };
927
928        let is_trap = error.is_some();
929        let trap_result =
930            if let Some(ref err) = error { Some(self.handle_error(*err)?) } else { None };
931
932        self.emit_events(
933            self.core.clk(),
934            self.core.next_pc(),
935            is_trap,
936            is_sigreturn,
937            instruction,
938            &mem_access_record,
939            self.core.exit_code(),
940        );
941
942        self.emit_syscall_event(
943            self.core.clk(),
944            code,
945            b,
946            c,
947            &mem_access_record,
948            self.core.next_pc(),
949            self.core.exit_code(),
950            instruction,
951            sig_return_pc_record,
952            trap_result,
953            error,
954        );
955
956        Ok(())
957    }
958}
959
960impl<M: ExecutionMode> TracingVM<'_, M> {
961    /// Emit events for a trap.
962    fn emit_trap_events(&mut self, clk: u64, next_pc: u64) {
963        self.record.pc_start.get_or_insert(self.core.pc());
964        self.record.next_pc = next_pc;
965        self.record.cpu_event_count += 1;
966
967        let increment = self.core.next_clk() - clk;
968
969        let bump1 = clk % (1 << 24) + increment >= (1 << 24);
970        if bump1 {
971            self.record.bump_state_events.push((clk, increment, false, next_pc));
972        }
973    }
974
975    /// Emit events for this cycle.
976    #[allow(clippy::too_many_arguments)]
977    fn emit_events(
978        &mut self,
979        clk: u64,
980        next_pc: u64,
981        is_trap: bool,
982        is_sig_return: bool,
983        instruction: &Instruction,
984        record: &MemoryAccessRecord,
985        exit_code: u32,
986    ) {
987        self.record.pc_start.get_or_insert(self.core.pc());
988        self.record.next_pc = next_pc;
989        self.record.exit_code = exit_code;
990        self.record.cpu_event_count += 1;
991
992        let increment = self.core.next_clk() - clk;
993
994        let bump1 = clk % (1 << 24) + increment >= (1 << 24);
995        let bump2 = !instruction.is_with_correct_next_pc()
996            && !is_trap
997            && !is_sig_return
998            && next_pc == self.core.pc().wrapping_add(4)
999            && (next_pc >> 16) != (self.core.pc() >> 16);
1000
1001        if bump1 || bump2 {
1002            self.record.bump_state_events.push((clk, increment, bump2, next_pc));
1003        }
1004
1005        if let Some(x) = record.a {
1006            if x.current_record().timestamp >> 24 != x.previous_record().timestamp >> 24 {
1007                self.record.bump_memory_events.push((x, instruction.op_a as u64, false));
1008            }
1009        }
1010        if let Some(x) = record.b {
1011            if x.current_record().timestamp >> 24 != x.previous_record().timestamp >> 24 {
1012                self.record.bump_memory_events.push((x, instruction.op_b, false));
1013            }
1014        }
1015        if let Some(x) = record.c {
1016            if x.current_record().timestamp >> 24 != x.previous_record().timestamp >> 24 {
1017                self.record.bump_memory_events.push((x, instruction.op_c, false));
1018            }
1019        }
1020
1021        if let Some((_record, instruction_value)) = record.untrusted_instruction {
1022            let encoded_instruction = instruction_value;
1023
1024            self.emit_instruction_fetch_event(instruction, encoded_instruction, record);
1025
1026            self.decoded_instruction_events
1027                .entry(encoded_instruction)
1028                .and_modify(|e| e.multiplicity += 1)
1029                .or_insert_with(|| InstructionDecodeEvent {
1030                    instruction: *instruction,
1031                    encoded_instruction,
1032                    multiplicity: 1,
1033                });
1034        }
1035    }
1036
1037    // Emit an event for a trap due to an untrusted instruction not having permission.
1038    fn emit_trap_exec_event(&mut self, trap_result: TrapResult, page_prot_record: PageProtRecord) {
1039        let event = TrapExecEvent {
1040            clk: self.core.clk(),
1041            pc: self.core.pc(),
1042            trap_result,
1043            page_prot_record,
1044        };
1045        self.record.trap_exec_events.push(event);
1046    }
1047
1048    /// Emit a instruction fetch event.
1049    #[allow(clippy::too_many_arguments)]
1050    fn emit_instruction_fetch_event(
1051        &mut self,
1052        instruction: &Instruction,
1053        encoded_instruction: u32,
1054        record: &MemoryAccessRecord,
1055    ) {
1056        let event = InstructionFetchEvent {
1057            clk: self.core.clk(),
1058            pc: self.core.pc(),
1059            instruction: *instruction,
1060            encoded_instruction,
1061        };
1062        self.record.instruction_fetch_events.push((event, *record));
1063    }
1064
1065    #[allow(clippy::too_many_arguments)]
1066    #[inline]
1067    fn emit_trap_mem_instr_event(
1068        &mut self,
1069        instruction: &Instruction,
1070        a: u64,
1071        b: u64,
1072        c: u64,
1073        record: &MemoryAccessRecord,
1074        trap_result: TrapResult,
1075        op_a_0: bool,
1076    ) {
1077        let opcode = instruction.opcode;
1078        let event = TrapMemInstrEvent {
1079            clk: self.core.clk(),
1080            pc: self.core.pc(),
1081            opcode,
1082            a,
1083            b,
1084            c,
1085            op_a_0,
1086            page_prot_access: record.memory.unwrap().previous_page_prot_record().unwrap(),
1087            trap_result,
1088        };
1089        let record = ITypeRecord::new(record, instruction);
1090        self.record.trap_load_store_events.push((event, record));
1091    }
1092
1093    /// Emit a memory instruction event.
1094    #[allow(clippy::too_many_arguments)]
1095    #[inline]
1096    fn emit_mem_instr_event(
1097        &mut self,
1098        instruction: &Instruction,
1099        a: u64,
1100        b: u64,
1101        c: u64,
1102        record: &MemoryAccessRecord,
1103        op_a_0: bool,
1104    ) {
1105        let opcode = instruction.opcode;
1106        let event = MemInstrEvent {
1107            clk: self.core.clk(),
1108            pc: self.core.pc(),
1109            opcode,
1110            a,
1111            b,
1112            c,
1113            op_a_0,
1114            // SAFETY: We explicity populate the memory of the record on the following callsites:
1115            // - `execute_load`
1116            // - `execute_store`
1117            mem_access: unsafe { record.memory.unwrap_unchecked() },
1118        };
1119
1120        let record = ITypeRecord::new(record, instruction);
1121        if matches!(
1122            opcode,
1123            Opcode::LB
1124                | Opcode::LBU
1125                | Opcode::LH
1126                | Opcode::LHU
1127                | Opcode::LW
1128                | Opcode::LWU
1129                | Opcode::LD
1130        ) && op_a_0
1131        {
1132            self.record.memory_load_x0_events.push((event, record));
1133        } else if matches!(opcode, Opcode::LB | Opcode::LBU) {
1134            self.record.memory_load_byte_events.push((event, record));
1135        } else if matches!(opcode, Opcode::LH | Opcode::LHU) {
1136            self.record.memory_load_half_events.push((event, record));
1137        } else if matches!(opcode, Opcode::LW | Opcode::LWU) {
1138            self.record.memory_load_word_events.push((event, record));
1139        } else if opcode == Opcode::LD {
1140            self.record.memory_load_double_events.push((event, record));
1141        } else if opcode == Opcode::SB {
1142            self.record.memory_store_byte_events.push((event, record));
1143        } else if opcode == Opcode::SH {
1144            self.record.memory_store_half_events.push((event, record));
1145        } else if opcode == Opcode::SW {
1146            self.record.memory_store_word_events.push((event, record));
1147        } else if opcode == Opcode::SD {
1148            self.record.memory_store_double_events.push((event, record));
1149        }
1150    }
1151
1152    /// Emit an ALU event.
1153    fn emit_alu_event(
1154        &mut self,
1155        instruction: &Instruction,
1156        a: u64,
1157        b: u64,
1158        c: u64,
1159        record: &MemoryAccessRecord,
1160        op_a_0: bool,
1161    ) {
1162        let opcode = instruction.opcode;
1163        let event = AluEvent { clk: self.core.clk(), pc: self.core.pc(), opcode, a, b, c, op_a_0 };
1164
1165        if op_a_0 {
1166            let record = ALUTypeRecord::new(record, instruction);
1167            self.record.alu_x0_events.push((event, record));
1168            return;
1169        }
1170
1171        match opcode {
1172            Opcode::ADD => {
1173                let record = RTypeRecord::new(record, instruction);
1174                self.record.add_events.push((event, record));
1175            }
1176            Opcode::ADDW => {
1177                let record = ALUTypeRecord::new(record, instruction);
1178                self.record.addw_events.push((event, record));
1179            }
1180            Opcode::ADDI => {
1181                let record = ITypeRecord::new(record, instruction);
1182                self.record.addi_events.push((event, record));
1183            }
1184            Opcode::SUB => {
1185                let record = RTypeRecord::new(record, instruction);
1186                self.record.sub_events.push((event, record));
1187            }
1188            Opcode::SUBW => {
1189                let record = RTypeRecord::new(record, instruction);
1190                self.record.subw_events.push((event, record));
1191            }
1192            Opcode::XOR | Opcode::OR | Opcode::AND => {
1193                let record = ALUTypeRecord::new(record, instruction);
1194                self.record.bitwise_events.push((event, record));
1195            }
1196            Opcode::SLL | Opcode::SLLW => {
1197                let record = ALUTypeRecord::new(record, instruction);
1198                self.record.shift_left_events.push((event, record));
1199            }
1200            Opcode::SRL | Opcode::SRA | Opcode::SRLW | Opcode::SRAW => {
1201                let record = ALUTypeRecord::new(record, instruction);
1202                self.record.shift_right_events.push((event, record));
1203            }
1204            Opcode::SLT | Opcode::SLTU => {
1205                let record = ALUTypeRecord::new(record, instruction);
1206                self.record.lt_events.push((event, record));
1207            }
1208            Opcode::MUL | Opcode::MULHU | Opcode::MULHSU | Opcode::MULH | Opcode::MULW => {
1209                let record = RTypeRecord::new(record, instruction);
1210                self.record.mul_events.push((event, record));
1211            }
1212            Opcode::DIVU
1213            | Opcode::REMU
1214            | Opcode::DIV
1215            | Opcode::REM
1216            | Opcode::DIVW
1217            | Opcode::DIVUW
1218            | Opcode::REMUW
1219            | Opcode::REMW => {
1220                let record = RTypeRecord::new(record, instruction);
1221                self.record.divrem_events.push((event, record));
1222            }
1223            _ => unreachable!(),
1224        }
1225    }
1226
1227    /// Emit a jal event.
1228    #[inline]
1229    #[allow(clippy::too_many_arguments)]
1230    fn emit_jal_event(
1231        &mut self,
1232        instruction: &Instruction,
1233        a: u64,
1234        b: u64,
1235        c: u64,
1236        record: &MemoryAccessRecord,
1237        op_a_0: bool,
1238        next_pc: u64,
1239    ) {
1240        let event = JumpEvent {
1241            clk: self.core.clk(),
1242            pc: self.core.pc(),
1243            next_pc,
1244            opcode: instruction.opcode,
1245            a,
1246            b,
1247            c,
1248            op_a_0,
1249        };
1250        let record = JTypeRecord::new(record, instruction);
1251        self.record.jal_events.push((event, record));
1252    }
1253
1254    /// Emit a jalr event.
1255    #[inline]
1256    #[allow(clippy::too_many_arguments)]
1257    fn emit_jalr_event(
1258        &mut self,
1259        instruction: &Instruction,
1260        a: u64,
1261        b: u64,
1262        c: u64,
1263        record: &MemoryAccessRecord,
1264        op_a_0: bool,
1265        next_pc: u64,
1266    ) {
1267        let event = JumpEvent {
1268            clk: self.core.clk(),
1269            pc: self.core.pc(),
1270            next_pc,
1271            opcode: instruction.opcode,
1272            a,
1273            b,
1274            c,
1275            op_a_0,
1276        };
1277        let record = ITypeRecord::new(record, instruction);
1278        self.record.jalr_events.push((event, record));
1279    }
1280
1281    /// Emit a branch event.
1282    #[inline]
1283    #[allow(clippy::too_many_arguments)]
1284    fn emit_branch_event(
1285        &mut self,
1286        instruction: &Instruction,
1287        a: u64,
1288        b: u64,
1289        c: u64,
1290        record: &MemoryAccessRecord,
1291        op_a_0: bool,
1292        next_pc: u64,
1293    ) {
1294        let event = BranchEvent {
1295            clk: self.core.clk(),
1296            pc: self.core.pc(),
1297            next_pc,
1298            opcode: instruction.opcode,
1299            a,
1300            b,
1301            c,
1302            op_a_0,
1303        };
1304        let record = ITypeRecord::new(record, instruction);
1305        self.record.branch_events.push((event, record));
1306    }
1307
1308    /// Emit a `UType` event.
1309    #[inline]
1310    fn emit_utype_event(
1311        &mut self,
1312        instruction: &Instruction,
1313        a: u64,
1314        b: u64,
1315        c: u64,
1316        record: &MemoryAccessRecord,
1317        op_a_0: bool,
1318    ) {
1319        let event = UTypeEvent {
1320            clk: self.core.clk(),
1321            pc: self.core.pc(),
1322            opcode: instruction.opcode,
1323            a,
1324            b,
1325            c,
1326            op_a_0,
1327        };
1328        let record = JTypeRecord::new(record, instruction);
1329        self.record.utype_events.push((event, record));
1330    }
1331
1332    /// Emit a syscall event.
1333    #[allow(clippy::too_many_arguments)]
1334    fn emit_syscall_event(
1335        &mut self,
1336        clk: u64,
1337        syscall_code: SyscallCode,
1338        arg1: u64,
1339        arg2: u64,
1340        record: &MemoryAccessRecord,
1341        next_pc: u64,
1342        exit_code: u32,
1343        instruction: &Instruction,
1344        sig_return_pc_record: Option<MemoryReadRecord>,
1345        trap_result: Option<TrapResult>,
1346        trap_error: Option<TrapError>,
1347    ) {
1348        let syscall_event = self.syscall_event(
1349            clk,
1350            syscall_code,
1351            arg1,
1352            arg2,
1353            next_pc,
1354            exit_code,
1355            sig_return_pc_record,
1356            trap_result,
1357            trap_error,
1358        );
1359
1360        let record = RTypeRecord::new(record, instruction);
1361        self.record.syscall_events.push((syscall_event, record));
1362    }
1363}
1364
1365impl<'a, M: ExecutionMode> SyscallRuntime<'a, M> for TracingVM<'a, M> {
1366    const TRACING: bool = true;
1367
1368    fn core(&self) -> &CoreVM<'a, M> {
1369        &self.core
1370    }
1371
1372    fn core_mut(&mut self) -> &mut CoreVM<'a, M> {
1373        &mut self.core
1374    }
1375
1376    /// Create a syscall event.
1377    #[inline]
1378    fn syscall_event(
1379        &self,
1380        clk: u64,
1381        syscall_code: SyscallCode,
1382        arg1: u64,
1383        arg2: u64,
1384        next_pc: u64,
1385        exit_code: u32,
1386        sig_return_pc_record: Option<MemoryReadRecord>,
1387        trap_result: Option<TrapResult>,
1388        trap_error: Option<TrapError>,
1389    ) -> SyscallEvent {
1390        // should_send: if the syscall is usually sent and it is not manually set as internal.
1391        let should_send =
1392            syscall_code.should_send() != 0 && !self.core.is_retained_syscall(syscall_code);
1393
1394        SyscallEvent {
1395            pc: self.core.pc(),
1396            next_pc,
1397            clk,
1398            should_send,
1399            syscall_code,
1400            syscall_id: syscall_code.syscall_id(),
1401            arg1,
1402            arg2,
1403            exit_code,
1404            sig_return_pc_record,
1405            trap_result,
1406            trap_error,
1407        }
1408    }
1409
1410    fn add_precompile_event(
1411        &mut self,
1412        syscall_code: SyscallCode,
1413        syscall_event: SyscallEvent,
1414        event: PrecompileEvent,
1415    ) {
1416        self.record.precompile_events.add_event(syscall_code, syscall_event, event);
1417    }
1418
1419    fn record_mut(&mut self) -> &mut ExecutionRecord {
1420        self.record
1421    }
1422
1423    fn rr(&mut self, register: usize) -> MemoryReadRecord {
1424        let record = SyscallRuntime::rr(self.core_mut(), register);
1425
1426        if let Some(local_memory_access) = &mut self.precompile_local_memory_access {
1427            local_memory_access.insert_record(register as u64, record);
1428        } else {
1429            self.local_memory_access.insert_record(register as u64, record);
1430        }
1431
1432        record
1433    }
1434
1435    fn rw(&mut self, register: usize, value: u64) -> MemoryWriteRecord {
1436        let record = SyscallRuntime::rw(self.core_mut(), register, value);
1437
1438        if let Some(local_memory_access) = &mut self.precompile_local_memory_access {
1439            local_memory_access.insert_record(register as u64, record);
1440        } else {
1441            self.local_memory_access.insert_record(register as u64, record);
1442        }
1443
1444        record
1445    }
1446
1447    fn mr_without_prot(&mut self, addr: u64) -> MemoryReadRecord {
1448        let record = SyscallRuntime::mr_without_prot(self.core_mut(), addr);
1449        if let Some(local_memory_access) = &mut self.precompile_local_memory_access {
1450            local_memory_access.insert_record(addr, record);
1451        } else {
1452            self.local_memory_access.insert_record(addr, record);
1453        }
1454
1455        record
1456    }
1457
1458    fn mw_without_prot(&mut self, addr: u64) -> MemoryWriteRecord {
1459        let record = SyscallRuntime::mw_without_prot(self.core_mut(), addr);
1460        if let Some(local_memory_access) = &mut self.precompile_local_memory_access {
1461            local_memory_access.insert_record(addr, record);
1462        } else {
1463            self.local_memory_access.insert_record(addr, record);
1464        }
1465        record
1466    }
1467
1468    fn page_prot_write(&mut self, page_idx: u64, prot: u8) -> PageProtRecord {
1469        let record = SyscallRuntime::page_prot_write(self.core_mut(), page_idx, prot);
1470
1471        let clk = self.core().clk();
1472        if let Some(local_page_prot_access) = &mut self.precompile_local_page_prot_access {
1473            local_page_prot_access.insert_record(page_idx, record, clk, prot);
1474        } else {
1475            self.local_page_prot_access.insert_record(page_idx, record, clk, prot);
1476        }
1477
1478        record
1479    }
1480}
1481
1482#[derive(Debug, Default)]
1483pub struct LocalMemoryAccess {
1484    pub inner: HashMap<u64, MemoryLocalEvent>,
1485}
1486
1487impl LocalMemoryAccess {
1488    #[inline]
1489    #[allow(clippy::needless_pass_by_value)]
1490    pub(crate) fn insert_record(&mut self, addr: u64, event: impl IntoMemoryRecord) {
1491        self.inner
1492            .entry(addr)
1493            .and_modify(|e| {
1494                let current_record = event.current_record();
1495                let previous_record = event.previous_record();
1496
1497                // The latest record is the one with the highest timestamp.
1498                if current_record.timestamp > e.final_mem_access.timestamp {
1499                    e.final_mem_access = current_record;
1500                }
1501
1502                // The initial record is the one with the lowest timestamp.
1503                if previous_record.timestamp < e.initial_mem_access.timestamp {
1504                    e.initial_mem_access = previous_record;
1505                }
1506            })
1507            .or_insert_with(|| MemoryLocalEvent {
1508                addr,
1509                initial_mem_access: event.previous_record(),
1510                final_mem_access: event.current_record(),
1511            });
1512    }
1513}
1514
1515impl Deref for LocalMemoryAccess {
1516    type Target = HashMap<u64, MemoryLocalEvent>;
1517    fn deref(&self) -> &Self::Target {
1518        &self.inner
1519    }
1520}
1521
1522impl DerefMut for LocalMemoryAccess {
1523    fn deref_mut(&mut self) -> &mut Self::Target {
1524        &mut self.inner
1525    }
1526}
1527
1528#[derive(Debug, Default)]
1529pub struct LocalPageProtAccess {
1530    pub inner: HashMap<u64, PageProtLocalEvent>,
1531}
1532
1533impl LocalPageProtAccess {
1534    #[inline]
1535    #[allow(clippy::needless_pass_by_value)]
1536    pub(crate) fn insert_record(
1537        &mut self,
1538        page_idx: u64,
1539        previous_record: PageProtRecord,
1540        clk: u64,
1541        value: u8,
1542    ) {
1543        self.inner
1544            .entry(page_idx)
1545            .and_modify(|e| e.final_page_prot_access.timestamp = clk)
1546            .or_insert(PageProtLocalEvent {
1547                page_idx,
1548                initial_page_prot_access: previous_record,
1549                final_page_prot_access: {
1550                    let mut final_prev_page_prot_record = previous_record;
1551                    final_prev_page_prot_record.timestamp = clk;
1552                    final_prev_page_prot_record.page_prot = value;
1553                    final_prev_page_prot_record
1554                },
1555            });
1556    }
1557}
1558
1559pub struct PrecompileMemory<'a, 'b, M: ExecutionMode> {
1560    inner: &'b mut TracingVM<'a, M>,
1561}
1562
1563impl<'a, 'b, M: ExecutionMode> PrecompileMemory<'a, 'b, M> {
1564    pub(crate) fn new(inner: &'b mut TracingVM<'a, M>) -> Self {
1565        Self { inner }
1566    }
1567}
1568
1569impl<'a, M: ExecutionMode> SyscallRuntime<'a, M> for PrecompileMemory<'a, '_, M> {
1570    const TRACING: bool = true;
1571
1572    fn core(&self) -> &CoreVM<'a, M> {
1573        self.inner.core()
1574    }
1575
1576    fn core_mut(&mut self) -> &mut CoreVM<'a, M> {
1577        self.inner.core_mut()
1578    }
1579
1580    #[allow(clippy::too_many_arguments)]
1581    fn syscall_event(
1582        &self,
1583        clk: u64,
1584        syscall_code: SyscallCode,
1585        arg1: u64,
1586        arg2: u64,
1587        next_pc: u64,
1588        exit_code: u32,
1589        sig_return_pc_record: Option<MemoryReadRecord>,
1590        trap_result: Option<TrapResult>,
1591        trap_error: Option<TrapError>,
1592    ) -> SyscallEvent {
1593        self.inner.syscall_event(
1594            clk,
1595            syscall_code,
1596            arg1,
1597            arg2,
1598            next_pc,
1599            exit_code,
1600            sig_return_pc_record,
1601            trap_result,
1602            trap_error,
1603        )
1604    }
1605
1606    fn add_precompile_event(
1607        &mut self,
1608        syscall_code: SyscallCode,
1609        syscall_event: SyscallEvent,
1610        event: PrecompileEvent,
1611    ) {
1612        self.inner.add_precompile_event(syscall_code, syscall_event, event);
1613    }
1614
1615    fn record_mut(&mut self) -> &mut ExecutionRecord {
1616        self.inner.record_mut()
1617    }
1618
1619    fn postprocess_precompile(&mut self) -> (Vec<MemoryLocalEvent>, Vec<PageProtLocalEvent>) {
1620        let mut precompile_local_memory_access = Vec::new();
1621        if let Some(local_memory_access) = &mut self.inner.precompile_local_memory_access {
1622            for (addr, event) in local_memory_access.drain() {
1623                if let Some(cpu_mem_access) = self.inner.local_memory_access.remove(&addr) {
1624                    self.inner.record.cpu_local_memory_access.push(cpu_mem_access);
1625                }
1626                precompile_local_memory_access.push(event);
1627            }
1628        }
1629
1630        let mut precompile_local_page_prot_access = Vec::new();
1631        if let Some(local_page_prot_access) = &mut self.inner.precompile_local_page_prot_access {
1632            for (page_idx, event) in local_page_prot_access.inner.drain() {
1633                if let Some(cpu_page_prot_access) =
1634                    self.inner.local_page_prot_access.inner.remove(&page_idx)
1635                {
1636                    self.inner.record.cpu_local_page_prot_access.push(cpu_page_prot_access);
1637                }
1638
1639                precompile_local_page_prot_access.push(event);
1640            }
1641        }
1642
1643        (precompile_local_memory_access, precompile_local_page_prot_access)
1644    }
1645
1646    fn page_prot_write(&mut self, page_idx: u64, prot: u8) -> PageProtRecord {
1647        let prev_page_prot_record = self.inner.page_prot_write(page_idx, prot);
1648
1649        let clk = self.inner.core.clk();
1650        if let Some(local_page_prot_access) = &mut self.inner.precompile_local_page_prot_access {
1651            local_page_prot_access.insert_record(page_idx, prev_page_prot_record, clk, prot);
1652        }
1653        assert!(self.inner.precompile_local_page_prot_access.is_some());
1654        assert!(self.inner.precompile_local_page_prot_access.as_ref().unwrap().inner.len() == 1);
1655
1656        prev_page_prot_record
1657    }
1658
1659    fn page_prot_range_check(
1660        &mut self,
1661        start_page_idx: u64,
1662        end_page_idx: u64,
1663        page_prot_bitmap: u8,
1664    ) -> (Vec<PageProtRecord>, Option<TrapError>) {
1665        let (records, error) =
1666            self.inner.page_prot_range_check(start_page_idx, end_page_idx, page_prot_bitmap);
1667        let clk = self.inner.core.clk();
1668        for record in &records {
1669            if let Some(local_page_prot_access) = &mut self.inner.precompile_local_page_prot_access
1670            {
1671                local_page_prot_access.insert_record(
1672                    record.page_idx,
1673                    *record,
1674                    clk,
1675                    record.page_prot,
1676                );
1677            } else {
1678                self.inner.local_page_prot_access.insert_record(
1679                    record.page_idx,
1680                    *record,
1681                    clk,
1682                    record.page_prot,
1683                );
1684            }
1685        }
1686        (records, error)
1687    }
1688
1689    fn rr(&mut self, reg_no: usize) -> MemoryReadRecord {
1690        debug_assert!(reg_no < 32, "out of bounds register: {reg_no}");
1691
1692        let current_clk = self.inner.core.clk();
1693        let registers = self.inner.core.registers_mut();
1694        let old_record = registers[reg_no];
1695        let new_record = MemoryRecord { timestamp: current_clk, value: old_record.value };
1696        registers[reg_no] = new_record;
1697
1698        let record = MemoryReadRecord {
1699            value: old_record.value,
1700            timestamp: self.inner.core.clk(),
1701            prev_timestamp: old_record.timestamp,
1702            prev_page_prot_record: None,
1703        };
1704
1705        let reg_no = reg_no as u64;
1706        if let Some(local_memory_access) = &mut self.inner.precompile_local_memory_access {
1707            local_memory_access.insert_record(reg_no, record);
1708        } else {
1709            self.inner.local_memory_access.insert_record(reg_no, record);
1710        }
1711
1712        record
1713    }
1714
1715    fn rw(&mut self, reg_no: usize, value: u64) -> MemoryWriteRecord {
1716        debug_assert!(reg_no < 32, "out of bounds register: {reg_no}");
1717
1718        let current_clk = self.inner.core.timestamp(MemoryAccessPosition::A);
1719        let registers = self.inner.core.registers_mut();
1720        let old_record = registers[reg_no];
1721        let new_record = MemoryRecord { timestamp: current_clk, value };
1722        registers[reg_no] = new_record;
1723
1724        let record = MemoryWriteRecord {
1725            value,
1726            timestamp: current_clk,
1727            prev_timestamp: old_record.timestamp,
1728            prev_value: old_record.value,
1729            prev_page_prot_record: None,
1730        };
1731
1732        let reg_no = reg_no as u64;
1733        if let Some(local_memory_access) = &mut self.inner.precompile_local_memory_access {
1734            local_memory_access.insert_record(reg_no, record);
1735        } else {
1736            self.inner.local_memory_access.insert_record(reg_no, record);
1737        }
1738
1739        record
1740    }
1741
1742    fn mr_without_prot(&mut self, addr: u64) -> MemoryReadRecord {
1743        let record = self.inner.mr_without_prot(addr);
1744
1745        if let Some(local_memory_access) = &mut self.inner.precompile_local_memory_access {
1746            local_memory_access.insert_record(addr, record);
1747        } else {
1748            self.inner.local_memory_access.insert_record(addr, record);
1749        }
1750
1751        record
1752    }
1753
1754    fn mr_slice_without_prot(&mut self, addr: u64, len: usize) -> Vec<MemoryReadRecord> {
1755        let records = self.inner.mr_slice_without_prot(addr, len);
1756        for (i, record) in records.iter().enumerate() {
1757            if let Some(local_memory_access) = &mut self.inner.precompile_local_memory_access {
1758                local_memory_access.insert_record(addr + i as u64 * 8, *record);
1759            } else {
1760                self.inner.local_memory_access.insert_record(addr + i as u64 * 8, *record);
1761            }
1762        }
1763        records
1764    }
1765
1766    fn mw_slice_without_prot(&mut self, addr: u64, len: usize) -> Vec<MemoryWriteRecord> {
1767        let records = self.inner.mw_slice_without_prot(addr, len);
1768        for (i, record) in records.iter().enumerate() {
1769            if let Some(local_memory_access) = &mut self.inner.precompile_local_memory_access {
1770                local_memory_access.insert_record(addr + i as u64 * 8, *record);
1771            } else {
1772                self.inner.local_memory_access.insert_record(addr + i as u64 * 8, *record);
1773            }
1774        }
1775        records
1776    }
1777
1778    fn mw_without_prot(&mut self, addr: u64) -> MemoryWriteRecord {
1779        let record = self.inner.mw_without_prot(addr);
1780
1781        if let Some(local_memory_access) = &mut self.inner.precompile_local_memory_access {
1782            local_memory_access.insert_record(addr, record);
1783        } else {
1784            self.inner.local_memory_access.insert_record(addr, record);
1785        }
1786
1787        record
1788    }
1789}
1790
1791/// Wrapper enum to handle `TracingVM` with different execution modes at runtime.
1792pub enum TracingVMEnum<'a> {
1793    /// `TracingVM` for `SupervisorMode`.
1794    Supervisor(TracingVM<'a, SupervisorMode>),
1795    /// `TracingVM` for `UserMode`.
1796    User(TracingVM<'a, UserMode>),
1797}
1798
1799impl<'a> TracingVMEnum<'a> {
1800    /// Create a new `TracingVMEnum` based on program's `enable_untrusted_programs` flag.
1801    pub fn new<T: MinimalTrace>(
1802        trace: &'a T,
1803        program: Arc<Program>,
1804        opts: SP1CoreOpts,
1805        proof_nonce: [u32; PROOF_NONCE_NUM_WORDS],
1806        record: &'a mut ExecutionRecord,
1807    ) -> Self {
1808        if program.enable_untrusted_programs {
1809            Self::User(TracingVM::<UserMode>::new(trace, program, opts, proof_nonce, record))
1810        } else {
1811            Self::Supervisor(TracingVM::<SupervisorMode>::new(
1812                trace,
1813                program,
1814                opts,
1815                proof_nonce,
1816                record,
1817            ))
1818        }
1819    }
1820
1821    /// Execute the program until it halts or reaches a shard boundary.
1822    pub fn execute(&mut self) -> Result<CycleResult, ExecutionError> {
1823        match self {
1824            Self::Supervisor(vm) => vm.execute(),
1825            Self::User(vm) => vm.execute(),
1826        }
1827    }
1828
1829    /// Get the public values.
1830    #[must_use]
1831    pub fn public_values(&self) -> &PublicValues<u32, u64, u64, u32> {
1832        match self {
1833            Self::Supervisor(vm) => vm.public_values(),
1834            Self::User(vm) => vm.public_values(),
1835        }
1836    }
1837
1838    /// Get the current clock.
1839    #[must_use]
1840    pub fn clk(&self) -> u64 {
1841        match self {
1842            Self::Supervisor(vm) => vm.core.clk(),
1843            Self::User(vm) => vm.core.clk(),
1844        }
1845    }
1846
1847    /// Get the current PC.
1848    #[must_use]
1849    pub fn pc(&self) -> u64 {
1850        match self {
1851            Self::Supervisor(vm) => vm.core.pc(),
1852            Self::User(vm) => vm.core.pc(),
1853        }
1854    }
1855
1856    /// Get the registers.
1857    #[must_use]
1858    pub fn registers(&self) -> &[MemoryRecord; 32] {
1859        match self {
1860            Self::Supervisor(vm) => vm.core.registers(),
1861            Self::User(vm) => vm.core.registers(),
1862        }
1863    }
1864
1865    /// Get the record.
1866    #[must_use]
1867    pub fn record(&self) -> &ExecutionRecord {
1868        match self {
1869            Self::Supervisor(vm) => vm.record,
1870            Self::User(vm) => vm.record,
1871        }
1872    }
1873
1874    /// Get the record mutably.
1875    #[must_use]
1876    pub fn record_mut(&mut self) -> &mut ExecutionRecord {
1877        match self {
1878            Self::Supervisor(vm) => vm.record,
1879            Self::User(vm) => vm.record,
1880        }
1881    }
1882}