diff --git a/crates/leanVm/src/bytecode/hint.rs b/crates/leanVm/src/bytecode/hint.rs index 84ca1b41..c1b4bdbf 100644 --- a/crates/leanVm/src/bytecode/hint.rs +++ b/crates/leanVm/src/bytecode/hint.rs @@ -54,7 +54,7 @@ impl Hint { /// /// These are not part of the trace or AIR and are only used by the prover for state setup or inspection. /// The verifier does not need to observe these effects. - fn execute( + pub fn execute( &self, memory_manager: &mut MemoryManager, run_context: &mut RunContext, diff --git a/crates/leanVm/src/bytecode/program.rs b/crates/leanVm/src/bytecode/program.rs index 872434ea..1fe121e0 100644 --- a/crates/leanVm/src/bytecode/program.rs +++ b/crates/leanVm/src/bytecode/program.rs @@ -1,5 +1,8 @@ use super::Bytecode; -use crate::constant::F; +use crate::{ + constant::F, + memory::{address::MemoryAddress, manager::MemoryManager}, +}; /// Represents a program to be executed by the zkVM. #[derive(Debug, Clone, Default)] @@ -11,3 +14,29 @@ pub struct Program { /// The private (witness) inputs for this specific execution. pub private_input: Vec, } + +/// Represents the results of a successful program execution. +#[derive(Debug)] +pub struct ExecutionResult { + /// The final state of the memory manager. + pub memory_manager: MemoryManager, + /// The execution trace of the program counter (pc) values. + /// + /// TODO: in the future, not sure we need this. + pub pcs: Vec, + /// The execution trace of the frame pointer (fp) values. + /// + /// TODO: in the future, not sure we need this. + pub fps: Vec, +} + +/// An helper struct to hold the results of a single execution pass. +#[derive(Debug)] +pub struct ExecutionPassResult { + /// The result of the execution. + pub result: ExecutionResult, + /// The initial allocation pointers. + pub initial_ap: MemoryAddress, + /// The initial allocation pointers for vectorized memory. + pub initial_ap_vec: MemoryAddress, +} diff --git a/crates/leanVm/src/constant.rs b/crates/leanVm/src/constant.rs index 1949473a..aa63a4db 100644 --- a/crates/leanVm/src/constant.rs +++ b/crates/leanVm/src/constant.rs @@ -34,3 +34,6 @@ pub const POSEIDON_24_NULL_HASH_PTR: MemoryAddress = MemoryAddress::new(VEC_RUNT /// Start of the public input memory region within the PUBLIC_DATA_SEGMENT. pub const PUBLIC_INPUT_START: MemoryAddress = MemoryAddress::new(PUBLIC_DATA_SEGMENT, 0); + +/// The maximum size of the memory. +pub const MAX_MEMORY_SIZE: usize = 1 << 23; diff --git a/crates/leanVm/src/core.rs b/crates/leanVm/src/core.rs index ff2c5ada..b1b2e0c0 100644 --- a/crates/leanVm/src/core.rs +++ b/crates/leanVm/src/core.rs @@ -5,11 +5,11 @@ use crate::{ bytecode::{ instruction::{Instruction, jump::JumpIfNotZeroInstruction}, operand::MemOrFp, - program::Program, + program::{ExecutionPassResult, ExecutionResult, Program}, }, constant::{ - CODE_SEGMENT, DIMENSION, F, POSEIDON_16_NULL_HASH_PTR, POSEIDON_24_NULL_HASH_PTR, - PUBLIC_INPUT_START, ZERO_VEC_PTR, + CODE_SEGMENT, DIMENSION, F, MAX_MEMORY_SIZE, POSEIDON_16_NULL_HASH_PTR, + POSEIDON_24_NULL_HASH_PTR, PUBLIC_INPUT_START, ZERO_VEC_PTR, }, context::run_context::RunContext, errors::{memory::MemoryError, vm::VirtualMachineError}, @@ -23,6 +23,15 @@ pub struct VirtualMachine { pub poseidon2_16: Perm16, pub poseidon2_24: Perm24, pub(crate) program: Program, + // Fields for tracing execution and gathering statistics. + pub(crate) pcs: Vec, + pub(crate) fps: Vec, + pub(crate) cpu_cycles: u64, + pub(crate) poseidon16_calls: u64, + pub(crate) poseidon24_calls: u64, + pub(crate) ext_mul_calls: u64, + // A string buffer to hold output from the Print hint. + pub(crate) std_out: String, } impl VirtualMachine @@ -37,6 +46,14 @@ where program: Program::default(), poseidon2_16, poseidon2_24, + // Initialize new fields for tracing and stats. + pcs: Vec::new(), + fps: Vec::new(), + cpu_cycles: 0, + poseidon16_calls: 0, + poseidon24_calls: 0, + ext_mul_calls: 0, + std_out: String::new(), } } @@ -261,6 +278,145 @@ where // update the pc and fp registers to prepare for the next cycle. self.update_registers(instruction) } + + /// Resets the machine's state to prepare for a new execution run. + fn reset_state(&mut self) { + self.run_context = RunContext::default(); + self.memory_manager = MemoryManager::default(); + self.program = Program::default(); + self.pcs.clear(); + self.fps.clear(); + self.cpu_cycles = 0; + self.poseidon16_calls = 0; + self.poseidon24_calls = 0; + self.ext_mul_calls = 0; + self.std_out.clear(); + } + + /// Updates execution statistics based on the instruction that was just executed. + const fn update_statistics(&mut self, instruction: &Instruction) { + match instruction { + Instruction::Poseidon2_16(_) => self.poseidon16_calls += 1, + Instruction::Poseidon2_24(_) => self.poseidon24_calls += 1, + Instruction::ExtensionMul(_) => self.ext_mul_calls += 1, + _ => (), // Other instructions do not have special counters. + } + } + + /// Executes a program, returning the final machine state and execution trace. + /// + /// This function implements the two-pass execution strategy: + /// 1. A first pass ("dry run") is executed to determine the exact amount of non-vectorized + /// runtime memory (`no_vec_runtime_memory`) required. + /// 2. A second, final pass is executed with the correctly sized memory allocations. + pub fn run(&mut self, program: &Program) -> Result { + // First pass: execute with an initial guess for non-vector memory to determine the actual required size. + let first_pass_result = + self.execute_program(program.clone(), MAX_MEMORY_SIZE / 2, false)?; + let no_vec_runtime_memory = + self.run_context.ap.offset - first_pass_result.initial_ap.offset; + + // Second pass: execute with the exact memory size determined from the first run. + let final_result = self.execute_program(program.clone(), no_vec_runtime_memory, true)?; + + Ok(final_result.result) + } + + /// The core execution loop of the virtual machine. + fn execute_program( + &mut self, + program: Program, + no_vec_runtime_memory: usize, + final_execution: bool, + ) -> Result { + // Reset state for the current run. + self.reset_state(); + + // Setup the VM state: load program, inputs, and initialize memory/registers. + self.setup(program, no_vec_runtime_memory)?; + + // Store initial allocation pointers to calculate memory usage later. + let initial_ap = self.run_context.ap; + let initial_ap_vec = self.run_context.ap_vectorized; + + // Main execution loop: continues until the program counter (pc) reaches the designated end address. + while self.run_context.pc.offset != self.program.bytecode.ending_pc { + // Ensure the program counter is within the bounds of the loaded bytecode instructions. + if self.run_context.pc.offset >= self.program.bytecode.instructions.len() { + return Err(VirtualMachineError::PCOutOfBounds); + } + + // Record current pc and fp for the execution trace. + self.pcs.push(self.run_context.pc); + self.fps.push(self.run_context.fp); + + self.cpu_cycles += 1; + + // Execute Hints: process non-deterministic hints associated with the current pc. + if let Some(hints) = self.program.bytecode.hints.get(&self.run_context.pc.offset) { + for hint in hints { + hint.execute(&mut self.memory_manager, &mut self.run_context)?; + } + } + + // Fetch and Execute Instruction. + let instruction = + self.program.bytecode.instructions[self.run_context.pc.offset].clone(); + self.run_instruction(&instruction)?; + + // Update statistics based on instruction type. + self.update_statistics(&instruction); + } + + // Record the final state of pc and fp. + self.pcs.push(self.run_context.pc); + self.fps.push(self.run_context.fp); + + if final_execution { + self.log_summary(); + } + + Ok(ExecutionPassResult { + result: ExecutionResult { + memory_manager: self.memory_manager.clone(), + pcs: self.pcs.clone(), + fps: self.fps.clone(), + }, + initial_ap, + initial_ap_vec, + }) + } + + /// Logs a summary of the execution results to standard output. + /// + /// TODO: remove this in the future (helper for debugging). + fn log_summary(&self) { + if !self.std_out.is_empty() { + print!("{}", self.std_out); + } + + println!( + "\nBytecode size: {}", + self.program.bytecode.instructions.len() + ); + println!("Public input size: {}", self.program.public_input.len()); + println!("Private input size: {}", self.program.private_input.len()); + println!("Executed {} instructions", self.cpu_cycles); + + let total_poseidon_calls = self.poseidon16_calls + self.poseidon24_calls; + if total_poseidon_calls > 0 { + println!( + "Poseidon2_16 calls: {}, Poseidon2_24 calls: {} (1 poseidon per {} instructions)", + self.poseidon16_calls, + self.poseidon24_calls, + self.cpu_cycles / total_poseidon_calls + ); + } + + if self.ext_mul_calls > 0 { + println!("ExtensionMul calls: {}", self.ext_mul_calls,); + } + } } #[cfg(test)] diff --git a/crates/leanVm/src/errors/vm.rs b/crates/leanVm/src/errors/vm.rs index 287e6c4b..b3ced7e7 100644 --- a/crates/leanVm/src/errors/vm.rs +++ b/crates/leanVm/src/errors/vm.rs @@ -22,4 +22,6 @@ pub enum VirtualMachineError { }, #[error("Too many unknown operands.")] TooManyUnknownOperands, + #[error("Program counter (pc) is out of bounds.")] + PCOutOfBounds, } diff --git a/crates/leanVm/src/memory/manager.rs b/crates/leanVm/src/memory/manager.rs index 4feb6e20..21edb98d 100644 --- a/crates/leanVm/src/memory/manager.rs +++ b/crates/leanVm/src/memory/manager.rs @@ -2,7 +2,7 @@ use super::{address::MemoryAddress, mem::Memory, val::MemoryValue}; use crate::errors::memory::MemoryError; /// A high level manager for the memory. -#[derive(Debug, Default)] +#[derive(Debug, Default, Clone)] pub struct MemoryManager { pub memory: Memory, } diff --git a/crates/leanVm/src/memory/mem.rs b/crates/leanVm/src/memory/mem.rs index b6a3ce90..53a5d7f9 100644 --- a/crates/leanVm/src/memory/mem.rs +++ b/crates/leanVm/src/memory/mem.rs @@ -3,7 +3,7 @@ use std::mem::MaybeUninit; use super::{address::MemoryAddress, cell::MemoryCell, val::MemoryValue}; use crate::errors::memory::MemoryError; -#[derive(Debug, Default)] +#[derive(Debug, Default, Clone)] pub struct Memory { pub(crate) data: Vec>, }