diff --git a/crates/leanVm/src/bytecode/hint.rs b/crates/leanVm/src/bytecode/hint.rs index 87124fe0..84ca1b41 100644 --- a/crates/leanVm/src/bytecode/hint.rs +++ b/crates/leanVm/src/bytecode/hint.rs @@ -79,18 +79,16 @@ impl Hint { // Store the current vectorized allocation pointer (`ap_vectorized`) at `addr`. memory_manager .memory - .insert(addr, F::from_usize(run_context.ap_vectorized))?; + .insert(addr, run_context.ap_vectorized)?; // Increase the vectorized allocation pointer by `size` (number of vectors). - run_context.ap_vectorized += size; + run_context.ap_vectorized.offset += size; } else { // Store the current scalar allocation pointer (`ap`) at `addr`. - memory_manager - .memory - .insert(addr, F::from_usize(run_context.ap))?; + memory_manager.memory.insert(addr, run_context.ap)?; // Increase the scalar allocation pointer by `size` (number of scalars). - run_context.ap += size; + run_context.ap.offset += size; } } Self::DecomposeBits { diff --git a/crates/leanVm/src/bytecode/mod.rs b/crates/leanVm/src/bytecode/mod.rs index cb9fb1bf..90f14a20 100644 --- a/crates/leanVm/src/bytecode/mod.rs +++ b/crates/leanVm/src/bytecode/mod.rs @@ -26,7 +26,7 @@ pub struct Bytecode { pub hints: BTreeMap>, /// The memory offset from the frame pointer (fp) where the public input for the program begins. - pub public_input_start: usize, + pub starting_frame_memory: usize, /// The program counter (pc) value at which the program execution is considered complete. pub ending_pc: usize, diff --git a/crates/leanVm/src/constant.rs b/crates/leanVm/src/constant.rs index 7979f099..1949473a 100644 --- a/crates/leanVm/src/constant.rs +++ b/crates/leanVm/src/constant.rs @@ -1,6 +1,36 @@ use p3_field::extension::BinomialExtensionField; use p3_koala_bear::KoalaBear; -pub(crate) const DIMENSION: usize = 8; +use crate::memory::address::MemoryAddress; + +/// The degree of the extension field. +pub const DIMENSION: usize = 8; + +/// The base field of the zkVM. pub(crate) type F = KoalaBear; + +/// The extension field of the zkVM. pub(crate) type EF = BinomialExtensionField; + +// Memory segment IDs + +/// Segment for public inputs and global constants. +pub const PUBLIC_DATA_SEGMENT: usize = 0; +/// Segment for the main stack (used by fp and ap). +pub const MAIN_STACK_SEGMENT: usize = 1; +/// Segment for runtime vector memory (for Poseidon, EF multiplication, etc.). +pub const VEC_RUNTIME_SEGMENT: usize = 2; +/// Segment for the compiled bytecode (where `pc` points). +pub const CODE_SEGMENT: usize = 3; + +// Convention-based virtual memory pointers. + +/// Points to `[0; DIMENSION]` in the vectorized memory segment. +pub const ZERO_VEC_PTR: MemoryAddress = MemoryAddress::new(VEC_RUNTIME_SEGMENT, 0); +/// Points to the result of `Poseidon2([0; 16])`, stored as 2 vector elements. +pub const POSEIDON_16_NULL_HASH_PTR: MemoryAddress = MemoryAddress::new(VEC_RUNTIME_SEGMENT, 1); +/// Points to the last 8 elements of `Poseidon2([0; 24])`, stored as 1 vector element. +pub const POSEIDON_24_NULL_HASH_PTR: MemoryAddress = MemoryAddress::new(VEC_RUNTIME_SEGMENT, 3); + +/// Start of the public input memory region within the PUBLIC_DATA_SEGMENT. +pub const PUBLIC_INPUT_START: MemoryAddress = MemoryAddress::new(PUBLIC_DATA_SEGMENT, 0); diff --git a/crates/leanVm/src/context/run_context.rs b/crates/leanVm/src/context/run_context.rs index 743c9b72..c6d5b635 100644 --- a/crates/leanVm/src/context/run_context.rs +++ b/crates/leanVm/src/context/run_context.rs @@ -25,7 +25,7 @@ pub struct RunContext { /// While `ap` determines where memory is written at runtime, its usage is opaque to the verifier. /// Instead, memory layout is statically determined and reflected through hint instructions that /// record where allocations were made. - pub(crate) ap: usize, + pub(crate) ap: MemoryAddress, /// Runtime allocation pointer (vectorized, chunked by `DIMENSION` field elements). /// @@ -38,7 +38,7 @@ pub struct RunContext { /// /// Like `ap`, this value is **not exposed to the verifier**. Its sole role is to guide prover-side /// allocation logic during execution. - pub(crate) ap_vectorized: usize, + pub(crate) ap_vectorized: MemoryAddress, } impl RunContext { @@ -46,8 +46,8 @@ impl RunContext { pub const fn new( pc: MemoryAddress, fp: MemoryAddress, - ap: usize, - ap_vectorized: usize, + ap: MemoryAddress, + ap_vectorized: MemoryAddress, ) -> Self { Self { pc, @@ -150,8 +150,8 @@ mod tests { segment_index: 1, offset: 0, }, - 0, - 0, + MemoryAddress::default(), + MemoryAddress::default(), ); // A constant operand with field element 42. @@ -189,8 +189,8 @@ mod tests { offset: 0, }, // dummy pc fp, - 0, - 0, + MemoryAddress::default(), + MemoryAddress::default(), ); // The operand asks to read memory at fp + 2. @@ -220,8 +220,8 @@ mod tests { offset: 0, }, // dummy pc fp, - 0, - 0, + MemoryAddress::default(), + MemoryAddress::default(), ); // Calling value_from_mem_or_constant should return a VirtualMachineError::MemoryError::UninitializedMemory. @@ -240,7 +240,12 @@ mod tests { #[test] fn test_get_value_from_mem_or_fp_or_constant_is_constant() { - let ctx = RunContext::new(MemoryAddress::new(0, 0), MemoryAddress::new(1, 0), 0, 0); + let ctx = RunContext::new( + MemoryAddress::new(0, 0), + MemoryAddress::new(1, 0), + MemoryAddress::default(), + MemoryAddress::default(), + ); let operand = MemOrFpOrConstant::Constant(F::from_u64(123)); let memory = MemoryManager::default(); let result = ctx @@ -252,7 +257,12 @@ mod tests { #[test] fn test_get_value_from_mem_or_fp_or_constant_is_fp() { let fp_addr = MemoryAddress::new(1, 10); - let ctx = RunContext::new(MemoryAddress::new(0, 0), fp_addr, 0, 0); + let ctx = RunContext::new( + MemoryAddress::new(0, 0), + fp_addr, + MemoryAddress::default(), + MemoryAddress::default(), + ); let operand = MemOrFpOrConstant::Fp; let memory = MemoryManager::default(); let result = ctx @@ -269,7 +279,12 @@ mod tests { let expected_val = MemoryValue::Address(MemoryAddress::new(5, 5)); memory.memory.insert(addr_to_read, expected_val).unwrap(); - let ctx = RunContext::new(MemoryAddress::new(0, 0), fp, 0, 0); + let ctx = RunContext::new( + MemoryAddress::new(0, 0), + fp, + MemoryAddress::default(), + MemoryAddress::default(), + ); let operand = MemOrFpOrConstant::MemoryAfterFp { shift: 7 }; let result = ctx .value_from_mem_or_fp_or_constant(&operand, &memory) diff --git a/crates/leanVm/src/core.rs b/crates/leanVm/src/core.rs index cbe182f8..ff2c5ada 100644 --- a/crates/leanVm/src/core.rs +++ b/crates/leanVm/src/core.rs @@ -1,4 +1,4 @@ -use p3_field::{Field, PrimeField64}; +use p3_field::{Field, PrimeCharacteristicRing, PrimeField64}; use p3_symmetric::Permutation; use crate::{ @@ -7,10 +7,13 @@ use crate::{ operand::MemOrFp, program::Program, }, - constant::{DIMENSION, F}, + constant::{ + CODE_SEGMENT, DIMENSION, F, 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}, - memory::manager::MemoryManager, + memory::{address::MemoryAddress, manager::MemoryManager}, }; #[derive(Debug)] @@ -22,7 +25,11 @@ pub struct VirtualMachine { pub(crate) program: Program, } -impl VirtualMachine { +impl VirtualMachine +where + Perm16: Permutation<[F; 2 * DIMENSION]>, + Perm24: Permutation<[F; 3 * DIMENSION]>, +{ pub fn new(poseidon2_16: Perm16, poseidon2_24: Perm24) -> Self { Self { run_context: RunContext::default(), @@ -33,6 +40,97 @@ impl VirtualMachine { } } + /// Initializes the virtual machine with a given program. + /// + /// This function performs all the necessary setup before execution: + /// + /// - Allocates the required memory segments + /// - Loads public and private inputs into stack memory + /// - Fills in convention-based memory slots (e.g., zero vector, null hashes) + /// - Computes and aligns allocation pointers + /// - Sets up the initial execution context (program counter, frame pointer, etc.) + pub fn setup( + &mut self, + program: Program, + no_vec_runtime_memory: usize, + ) -> Result<(), VirtualMachineError> { + // Save the program internally. + self.program = program; + + // Extract the public and private inputs from the program. + let public_input = &self.program.public_input; + let private_input = &self.program.private_input; + + // Allocate required memory segments: PUBLIC, STACK, VEC, CODE. + // We ensure all necessary segments are present. + while self.memory_manager.num_segments() <= CODE_SEGMENT { + self.memory_manager.add(); + } + + // Convention-Based Memory Initialization + + // Write [0; DIMENSION] to the vector memory at ZERO_VEC_PTR. + self.memory_manager + .load_data(ZERO_VEC_PTR, &[F::ZERO; DIMENSION])?; + + // Write Poseidon2([0; 16]) to POSEIDON_16_NULL_HASH_PTR. + let hash16 = self.poseidon2_16.permute([F::ZERO; DIMENSION * 2]); + self.memory_manager + .load_data(POSEIDON_16_NULL_HASH_PTR, &hash16)?; + + // Write the last 8 elements of Poseidon2([0; 24]) to POSEIDON_24_NULL_HASH_PTR. + let hash24 = self.poseidon2_24.permute([F::ZERO; DIMENSION * 3]); + self.memory_manager + .load_data(POSEIDON_24_NULL_HASH_PTR, &hash24[16..])?; + + // Load Public Inputs + + // Place public input values starting at PUBLIC_INPUT_START in the public data segment. + self.memory_manager + .load_data(PUBLIC_INPUT_START, public_input)?; + + // Compute the initial `fp` (frame pointer) just after the public inputs. + let mut fp = (PUBLIC_INPUT_START + public_input.len())?; + // Align the `fp` offset to the next power of two. + fp.offset = fp.offset.next_power_of_two(); + + // Load Private Inputs + + // Write private inputs starting at the aligned `fp`. + self.memory_manager.load_data(fp, private_input)?; + + // Advance `fp` past the private inputs. + fp.offset += private_input.len(); + // Ensure `fp` is aligned to `DIMENSION` for vector operations. + fp.offset = fp.offset.next_multiple_of(DIMENSION); + + // Compute Allocation Pointers + + // Compute the initial allocation pointer for stack memory. + let initial_ap = (fp + self.program.bytecode.starting_frame_memory)?; + + // Compute the vectorized allocation pointer, skipping past the non-vector memory. + let mut initial_ap_vec = (initial_ap + no_vec_runtime_memory)?; + // Align the vector allocation to the next multiple of `DIMENSION`. + initial_ap_vec.offset = initial_ap_vec.offset.next_multiple_of(DIMENSION) / DIMENSION; + + // Set Initial Registers + + // Set the program counter to the start of the code segment. + self.run_context.pc = MemoryAddress::new(CODE_SEGMENT, 0); + + // Set the frame pointer to the aligned `fp`. + self.run_context.fp = fp; + + // Set the allocation pointer for non-vector memory. + self.run_context.ap = initial_ap; + + // Set the allocation pointer for vector memory (in vector units, not field elements). + self.run_context.ap_vectorized = initial_ap_vec; + + Ok(()) + } + /// Advances the program counter (`pc`) to the next instruction. /// /// This function embodies the control flow logic of the zkVM. For most instructions, @@ -147,11 +245,10 @@ impl VirtualMachine { /// 2. **Register Update:** If the execution phase completes successfully, this function then /// calls `update_registers` to advance the program counter (`pc`) and frame pointer (`fp`) /// to prepare for the next instruction. - pub fn run_instruction(&mut self, instruction: &Instruction) -> Result<(), VirtualMachineError> - where - Perm16: Permutation<[F; 2 * DIMENSION]>, - Perm24: Permutation<[F; 3 * DIMENSION]>, - { + pub fn run_instruction( + &mut self, + instruction: &Instruction, + ) -> Result<(), VirtualMachineError> { // Execute the instruction. instruction.execute( &self.run_context,