vm core: setup implementation

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 {

crates/leanVm/src/bytecode/mod.rs

@@ -26,7 +26,7 @@ pub struct Bytecode {
     pub hints: BTreeMap<usize, Vec<Hint>>,
 
     /// 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,

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<F, DIMENSION>;
+
+// 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);

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,16 +38,16 @@ 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 {
     #[must_use]
     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)

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<Perm16, Perm24> {
     pub(crate) program: Program,
 }
 
-impl<Perm16, Perm24> VirtualMachine<Perm16, Perm24> {
+impl<Perm16, Perm24> VirtualMachine<Perm16, Perm24>
+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<Perm16, Perm24> VirtualMachine<Perm16, Perm24> {
         }
     }
 
+    /// 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<Perm16, Perm24> VirtualMachine<Perm16, Perm24> {
     /// 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,