Revamp lean vm

crates/air/src/prove.rs

@@ -5,8 +5,8 @@ use p3_field::{ExtensionField, Field, cyclic_subgroup_known_order};
 use p3_util::{log2_ceil_usize, log2_strict_usize};
 use sumcheck::{MleGroup, MleGroupOwned, MleGroupRef, ProductComputation};
 use tracing::{info_span, instrument};
-use utils::PF;
-use utils::{FSProver, add_multilinears, multilinears_linear_combination};
+use utils::{FSProver, multilinears_linear_combination};
+use utils::{PF, add_multilinears_inplace};
 use whir_p3::fiat_shamir::FSChallenger;
 use whir_p3::poly::evals::{eval_eq, fold_multilinear, scale_poly};
 use whir_p3::poly::{evals::EvaluationsList, multilinear::MultilinearPoint};
@@ -195,8 +195,9 @@ fn open_structured_columns<EF: ExtensionField<PF<EF>> + ExtensionField<IF>, IF:
 
     let batched_column =
         multilinears_linear_combination(witness, &poly_eq_batching_scalars[..n_columns]);
-    let batched_column_mixed = add_multilinears(
-        &column_up(&batched_column),
+    let mut batched_column_mixed = column_up(&batched_column);
+    add_multilinears_inplace(
+        &mut batched_column_mixed,
         &scale_poly(&column_down(&batched_column), alpha),
     );
     // TODO do not recompute this (we can deduce it from already computed values)
@@ -217,13 +218,9 @@ fn open_structured_columns<EF: ExtensionField<PF<EF>> + ExtensionField<IF>, IF:
     // TODO do not recompute this (we can deduce it from already computed values)
     let inner_sum = batched_column_mixed.evaluate(&MultilinearPoint(point.clone()));
 
-    let inner_mle = MleGroupOwned::Extension(vec![
-        add_multilinears(
-            &matrix_up_folded(&point),
-            &scale_poly(&matrix_down_folded(&point), alpha),
-        ),
-        batched_column,
-    ]);
+    let mut mat_up = matrix_up_folded(&point);
+    add_multilinears_inplace(&mut mat_up, &scale_poly(&matrix_down_folded(&point), alpha));
+    let inner_mle = MleGroupOwned::Extension(vec![mat_up, batched_column]);
 
     let (inner_challenges, _, _) = sumcheck::prove::<EF, _, _, _>(
         1,

crates/lean_compiler/src/lib.rs

@@ -24,7 +24,7 @@ pub fn compile_program(program: &str) -> (Bytecode, BTreeMap<usize, String>) {
     let intermediate_bytecode = compile_to_intermediate_bytecode(simple_program).unwrap();
     // println!("Intermediate Bytecode:\n\n{}", intermediate_bytecode.to_string());
     let compiled = compile_to_low_level_bytecode(intermediate_bytecode).unwrap();
-    println!("Compiled Program:\n\n{}", compiled);
+    println!("Compiled Program:\n\n{compiled}");
     (compiled, function_locations)
 }
 

crates/lean_vm/src/core/constants.rs

@@ -0,0 +1,28 @@
+//! VM constants and configuration
+
+/// Vector dimension for field operations
+pub const DIMENSION: usize = 5;
+
+/// Logarithm of vector length
+pub const LOG_VECTOR_LEN: usize = 3;
+
+/// Vector length (2^LOG_VECTOR_LEN)
+pub const VECTOR_LEN: usize = 1 << LOG_VECTOR_LEN;
+
+/// Convention: vectorized pointer of size 2, pointing to 16 zeros
+pub const ZERO_VEC_PTR: usize = 0;
+
+/// Convention: vectorized pointer of size 1, pointing to 10000000
+pub const ONE_VEC_PTR: usize = 2;
+
+/// Convention: vectorized pointer of size 2, = the 16 elements of poseidon_16(0)
+pub const POSEIDON_16_NULL_HASH_PTR: usize = 3;
+
+/// Convention: vectorized pointer of size 1, = the last 8 elements of poseidon_24(0)
+pub const POSEIDON_24_NULL_HASH_PTR: usize = 5;
+
+/// Normal pointer to start of public input
+pub const PUBLIC_INPUT_START: usize = 6 * 8;
+
+/// Maximum memory size for VM runner
+pub const MAX_RUNNER_MEMORY_SIZE: usize = 1 << 20;

crates/lean_vm/src/core/mod.rs

@@ -0,0 +1,7 @@
+//! Core VM abstractions and type definitions
+
+pub mod constants;
+pub mod types;
+
+pub use constants::*;
+pub use types::*;

crates/lean_vm/src/core/types.rs

@@ -0,0 +1,12 @@
+//! Core type definitions for the VM
+
+use p3_koala_bear::{KoalaBear, QuinticExtensionFieldKB};
+
+/// Base field type for VM operations
+pub type F = KoalaBear;
+
+/// Extension field type for VM operations
+pub type EF = QuinticExtensionFieldKB;
+
+/// Location in source code for debugging
+pub type LocationInSourceCode = usize;

crates/lean_vm/src/diagnostics/error.rs

@@ -0,0 +1,62 @@
+//! VM error types and execution results
+
+use crate::core::F;
+use crate::execution::Memory;
+use crate::witness::{
+    WitnessDotProduct, WitnessMultilinearEval, WitnessPoseidon16, WitnessPoseidon24,
+};
+use thiserror::Error;
+
+/// Errors that can occur during VM execution
+#[derive(Debug, Clone, Error)]
+pub enum RunnerError {
+    #[error("Out of memory")]
+    OutOfMemory,
+
+    #[error("Memory already set")]
+    MemoryAlreadySet,
+
+    #[error("Not a pointer")]
+    NotAPointer,
+
+    #[error("Division by zero")]
+    DivByZero,
+
+    #[error("Computation invalid: {0} != {1}")]
+    NotEqual(F, F),
+
+    #[error("Undefined memory access")]
+    UndefinedMemory,
+
+    #[error("Program counter out of bounds")]
+    PCOutOfBounds,
+
+    #[error("Point for multilinear eval not padded with zeros")]
+    MultilinearEvalPointNotPadded,
+}
+
+/// Result type for VM operations
+pub type VMResult<T> = Result<T, RunnerError>;
+
+/// Execution result containing outputs and execution data
+#[derive(Debug)]
+pub struct ExecutionResult {
+    pub no_vec_runtime_memory: usize,
+    pub public_memory_size: usize,
+    pub memory: Memory,
+    pub pcs: Vec<usize>,
+    pub fps: Vec<usize>,
+    pub poseidons_16: Vec<WitnessPoseidon16>,
+    pub poseidons_24: Vec<WitnessPoseidon24>,
+    pub dot_products: Vec<WitnessDotProduct>,
+    pub multilinear_evals: Vec<WitnessMultilinearEval>,
+}
+
+impl ExecutionResult {
+    /// Check if execution was successful
+    ///
+    /// TODO: placeholder for now.
+    pub fn is_success(&self) -> bool {
+        true
+    }
+}

crates/lean_vm/src/diagnostics/mod.rs

@@ -0,0 +1,10 @@
+//! Diagnostics, error handling, and debugging utilities
+#![allow(unused_imports)]
+
+pub mod error;
+pub mod profiler;
+pub mod stack_trace;
+
+pub use error::*;
+pub use profiler::*;
+pub use stack_trace::*;

crates/lean_vm/src/execution/context.rs

@@ -0,0 +1,69 @@
+//! Execution context and state management
+
+use crate::core::LocationInSourceCode;
+use std::collections::BTreeMap;
+
+/// Execution history for profiling and debugging
+#[derive(Debug, Clone, Default)]
+pub struct ExecutionHistory {
+    pub lines: Vec<LocationInSourceCode>,
+    pub cycles: Vec<usize>, // for each line, how many cycles it took
+}
+
+impl ExecutionHistory {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn add_line(&mut self, location: LocationInSourceCode, cycles: usize) {
+        self.lines.push(location);
+        self.cycles.push(cycles);
+    }
+
+    pub fn total_cycles(&self) -> usize {
+        self.cycles.iter().sum()
+    }
+
+    pub const fn len(&self) -> usize {
+        self.lines.len()
+    }
+
+    pub const fn is_empty(&self) -> bool {
+        self.lines.is_empty()
+    }
+}
+
+/// VM execution context
+#[derive(Debug)]
+pub struct ExecutionContext<'a> {
+    pub source_code: &'a str,
+    pub function_locations: &'a BTreeMap<usize, String>,
+    pub profiler_enabled: bool,
+    pub std_out: String,
+    pub instruction_history: ExecutionHistory,
+}
+
+impl<'a> ExecutionContext<'a> {
+    pub fn new(
+        source_code: &'a str,
+        function_locations: &'a BTreeMap<usize, String>,
+        profiler_enabled: bool,
+    ) -> Self {
+        Self {
+            source_code,
+            function_locations,
+            profiler_enabled,
+            std_out: String::new(),
+            instruction_history: ExecutionHistory::new(),
+        }
+    }
+
+    pub fn print(&mut self, message: &str) {
+        self.std_out.push_str(message);
+    }
+
+    pub fn println(&mut self, message: &str) {
+        self.std_out.push_str(message);
+        self.std_out.push('\n');
+    }
+}