diff --git a/Cargo.toml b/Cargo.toml index da03a4b5..b04347f3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,6 +40,8 @@ p3-field = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } p3-symmetric = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } p3-koala-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } +p3-air = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } +p3-matrix = { git = "https://github.com/Plonky3/Plonky3.git", rev = "d0c4a36" } thiserror = "2.0" proptest = "1.7" diff --git a/crates/leanVm/Cargo.toml b/crates/leanVm/Cargo.toml index 43d085d7..70c1b3e7 100644 --- a/crates/leanVm/Cargo.toml +++ b/crates/leanVm/Cargo.toml @@ -13,6 +13,8 @@ p3-baby-bear.workspace = true p3-koala-bear.workspace = true p3-field.workspace = true p3-symmetric.workspace = true +p3-air.workspace = true +p3-matrix.workspace = true thiserror.workspace = true num-traits.workspace = true diff --git a/crates/leanVm/src/air/constant.rs b/crates/leanVm/src/air/constant.rs new file mode 100644 index 00000000..1ccc001b --- /dev/null +++ b/crates/leanVm/src/air/constant.rs @@ -0,0 +1,53 @@ +// Bytecode columns (looked up from the program ROM) + +/// The immediate value or memory offset for the first operand, `a`. +pub(crate) const COL_INDEX_OPERAND_A: usize = 0; +/// The immediate value or memory offset for the second operand, `b`. +pub(crate) const COL_INDEX_OPERAND_B: usize = 1; +/// The immediate value or memory offset for the third operand, `c`. +pub(crate) const COL_INDEX_OPERAND_C: usize = 2; +/// A boolean flag selecting between the immediate `OPERAND_A` (1) and a memory value (0). +pub(crate) const COL_INDEX_FLAG_A: usize = 3; +/// A boolean flag selecting between the immediate `OPERAND_B` (1) and a memory value (0). +pub(crate) const COL_INDEX_FLAG_B: usize = 4; +/// A boolean flag selecting between `fp` or an immediate `OPERAND_C` (1) and a memory value (0). +pub(crate) const COL_INDEX_FLAG_C: usize = 5; +/// The opcode flag that activates the ADD instruction's constraints. +pub(crate) const COL_INDEX_ADD: usize = 6; +/// The opcode flag that activates the MUL instruction's constraints. +pub(crate) const COL_INDEX_MUL: usize = 7; +/// The opcode flag that activates the DEREF instruction's constraints. +pub(crate) const COL_INDEX_DEREF: usize = 8; +/// The opcode flag that activates the JUZ (Jump Unless Zero) instruction's constraints. +pub(crate) const COL_INDEX_JUZ: usize = 9; +/// An auxiliary operand used for multi-purpose logic, like specifying the `res` type in DEREF. +pub(crate) const COL_INDEX_AUX: usize = 10; +/// The total number of columns representing a single decoded instruction from the bytecode. +pub(crate) const N_BYTECODE_COLUMNS: usize = 11; + +// Execution trace columns (committed by the prover) + +/// The column for the Program Counter (pc) register in the execution trace. +pub(crate) const COL_INDEX_PC: usize = N_BYTECODE_COLUMNS + 3; +/// The column for the Frame Pointer (fp) register in the execution trace. +pub(crate) const COL_INDEX_FP: usize = N_BYTECODE_COLUMNS + 4; +/// The column holding the memory address for operand `a` when `FLAG_A` is 0. +pub(crate) const COL_INDEX_ADDR_A: usize = N_BYTECODE_COLUMNS + 5; +/// The column holding the memory address for operand `b` when `FLAG_B` is 0. +pub(crate) const COL_INDEX_ADDR_B: usize = N_BYTECODE_COLUMNS + 6; +/// The column holding the memory address for operand `c` when `FLAG_C` is 0. +pub(crate) const COL_INDEX_ADDR_C: usize = N_BYTECODE_COLUMNS + 7; +/// The total count of columns in the execution trace that are directly committed by the prover. +pub(crate) const N_REAL_EXEC_COLUMNS: usize = 5; + +// Virtual columns (derived via lookups, not committed directly) + +/// The virtual column holding the value read from memory for operand `a`. +pub(crate) const COL_INDEX_MEM_VALUE_A: usize = N_BYTECODE_COLUMNS; +/// The virtual column holding the value read from memory for operand `b`. +pub(crate) const COL_INDEX_MEM_VALUE_B: usize = N_BYTECODE_COLUMNS + 1; +/// The virtual column holding the value read from memory for operand `c`. +pub(crate) const COL_INDEX_MEM_VALUE_C: usize = N_BYTECODE_COLUMNS + 2; + +/// The total number of columns in the AIR, including bytecode, real, and virtual columns. +pub(crate) const N_EXEC_AIR_COLUMNS: usize = N_BYTECODE_COLUMNS + 3 + N_REAL_EXEC_COLUMNS; diff --git a/crates/leanVm/src/air/mod.rs b/crates/leanVm/src/air/mod.rs new file mode 100644 index 00000000..5983319c --- /dev/null +++ b/crates/leanVm/src/air/mod.rs @@ -0,0 +1,164 @@ +//! zkVM AIR (Algebraic Intermediate Representation) +//! +//! This module defines the transition constraints for the custom zkVM's Instruction Set +//! Architecture (ISA). It translates the VM's operational semantics into a set of polynomial +//! equations that must hold true between consecutive states (rows) of the execution trace. + +use std::borrow::Borrow; + +use constant::{ + COL_INDEX_ADD, COL_INDEX_ADDR_A, COL_INDEX_ADDR_B, COL_INDEX_ADDR_C, COL_INDEX_AUX, + COL_INDEX_DEREF, COL_INDEX_FLAG_A, COL_INDEX_FLAG_B, COL_INDEX_FLAG_C, COL_INDEX_FP, + COL_INDEX_JUZ, COL_INDEX_MEM_VALUE_A, COL_INDEX_MEM_VALUE_B, COL_INDEX_MEM_VALUE_C, + COL_INDEX_MUL, COL_INDEX_OPERAND_A, COL_INDEX_OPERAND_B, COL_INDEX_OPERAND_C, COL_INDEX_PC, + N_EXEC_AIR_COLUMNS, +}; +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::PrimeCharacteristicRing; +use p3_matrix::Matrix; + +pub mod constant; + +/// Virtual Machine AIR +#[derive(Debug)] +pub struct VMAir; + +impl BaseAir for VMAir { + /// The total number of columns in the execution trace. + fn width(&self) -> usize { + N_EXEC_AIR_COLUMNS + } +} + +impl Air for VMAir { + #[inline] + fn eval(&self, builder: &mut AB) { + // Get a view of the main execution trace. + let main = builder.main(); + + // Get the current row (`local`) and the next row (`next`) from the trace. + let local = main.row_slice(0).unwrap(); + let local = local.borrow(); + + let next = main.row_slice(1).unwrap(); + let next = next.borrow(); + + // INSTRUCTION DECODING + // + // Extract instruction fields (operands, flags, and opcodes) from the local row. + // + // These are treated as constants for a given row, looked up from the bytecode. + let operand_a = local[COL_INDEX_OPERAND_A].clone().into(); + let operand_b = local[COL_INDEX_OPERAND_B].clone().into(); + let operand_c = local[COL_INDEX_OPERAND_C].clone().into(); + let flag_a = local[COL_INDEX_FLAG_A].clone().into(); + let flag_b = local[COL_INDEX_FLAG_B].clone().into(); + let flag_c = local[COL_INDEX_FLAG_C].clone().into(); + let add = local[COL_INDEX_ADD].clone().into(); + let mul = local[COL_INDEX_MUL].clone().into(); + let deref = local[COL_INDEX_DEREF].clone().into(); + let juz = local[COL_INDEX_JUZ].clone().into(); + let aux = local[COL_INDEX_AUX].clone().into(); + + // REGISTER & MEMORY VALUES + // + // Extract register values and memory access data from the trace. + let (pc, next_pc) = ( + local[COL_INDEX_PC].clone().into(), + next[COL_INDEX_PC].clone().into(), + ); + let (fp, next_fp) = ( + local[COL_INDEX_FP].clone().into(), + next[COL_INDEX_FP].clone().into(), + ); + let (addr_a, addr_b, addr_c) = ( + local[COL_INDEX_ADDR_A].clone().into(), + local[COL_INDEX_ADDR_B].clone().into(), + local[COL_INDEX_ADDR_C].clone().into(), + ); + let (value_a, value_b, value_c) = ( + local[COL_INDEX_MEM_VALUE_A].clone().into(), + local[COL_INDEX_MEM_VALUE_B].clone().into(), + local[COL_INDEX_MEM_VALUE_C].clone().into(), + ); + + // OPERAND RECONSTRUCTION + // + // Compute the effective values of the operands (`nu_a`, `nu_b`, `nu_c`). + // + // Each operand can be either: + // - an immediate value (from `operand_x`), + // - a value loaded from memory (from `value_x`), selected by `flag_x`. + // + // Formula: nu_x = flag_x * operand_x + (1 - flag_x) * value_x + let nu_a = + flag_a.clone() * operand_a.clone() + (AB::Expr::ONE - flag_a.clone()) * value_a.clone(); + let nu_b = flag_b.clone() * operand_b.clone() + (AB::Expr::ONE - flag_b.clone()) * value_b; + // Operand `c` is special: its immediate value can be the frame pointer `fp` itself. + let nu_c = flag_c.clone() * fp.clone() + (AB::Expr::ONE - flag_c.clone()) * value_c.clone(); + + // MEMORY ADDRESS CONSTRAINTS + // + // Enforce that if an operand is loaded from memory (`flag_x` = 0), its address + // (`addr_x`) must correctly correspond to the frame pointer plus its offset (`fp + operand_x`). + // + // If `flag_x` = 1, the constraint is `0 * ... = 0`, so it is satisfied. + builder.assert_zero((AB::Expr::ONE - flag_a) * (addr_a - (fp.clone() + operand_a))); + builder.assert_zero((AB::Expr::ONE - flag_b) * (addr_b - (fp.clone() + operand_b))); + builder.assert_zero( + (AB::Expr::ONE - flag_c) * (addr_c.clone() - (fp.clone() + operand_c.clone())), + ); + + // INSTRUCTION CONSTRAINTS + + // ADD Instruction Constraint + // + // If the `add` opcode is active, enforce `nu_a + nu_c = nu_b`. + builder.assert_zero(add * (nu_b.clone() - (nu_a.clone() + nu_c.clone()))); + + // MUL Instruction Constraint + // + // If the `mul` opcode is active, enforce `nu_a * nu_c = nu_b`. + builder.assert_zero(mul * (nu_b.clone() - nu_a.clone() * nu_c.clone())); + + // DEREF Instruction Constraints + // + // This instruction computes `m[m[fp + α] + β] = res`. + // + // 1. Enforce that the final address `addr_c` is computed correctly: `addr_c = value_a + operand_c`. + // + // Here, `value_a` is the pointer `m[fp + α]` and `operand_c` is the offset `β`. + builder.assert_zero(deref.clone() * (addr_c - (value_a + operand_c))); + // 2. If `aux` = 1, the result `res` is a normal value (`nu_b`). Enforce `value_c = nu_b`. + builder.assert_zero(deref.clone() * aux.clone() * (value_c.clone() - nu_b.clone())); + // 3. If `aux` = 0, the result `res` is the frame pointer itself. Enforce `value_c = fp`. + builder.assert_zero(deref * (AB::Expr::ONE - aux) * (value_c - fp.clone())); + + // JUZ (Jump Unless Zero) and Program Flow Constraints + + // Default Program Flow (No Jump) + // + // If `juz` is not active, the program counter must increment by 1. + builder.assert_zero( + (AB::Expr::ONE - juz.clone()) * (next_pc.clone() - (pc.clone() + AB::Expr::ONE)), + ); + // If `juz` is not active, the frame pointer must remain unchanged. + builder.assert_zero((AB::Expr::ONE - juz.clone()) * (next_fp.clone() - fp.clone())); + + // JUZ Active Program Flow + // + // The condition for the jump is `nu_a`. + // 1. Enforce that the condition `nu_a` is boolean (0 or 1). + builder.assert_zero(juz.clone() * nu_a.clone() * (AB::Expr::ONE - nu_a.clone())); + // 2. If jump is taken (`nu_a` = 1), the next `pc` must be the destination `nu_b`. + builder.assert_zero(juz.clone() * nu_a.clone() * (next_pc.clone() - nu_b)); + // 3. If jump is taken (`nu_a` = 1), the next `fp` must be the new frame pointer `nu_c`. + builder.assert_zero(juz.clone() * nu_a.clone() * (next_fp.clone() - nu_c)); + // 4. If jump is NOT taken (`nu_a` = 0), the next `pc` must be `pc + 1`. + builder.assert_zero( + juz.clone() * (AB::Expr::ONE - nu_a.clone()) * (next_pc - (pc + AB::Expr::ONE)), + ); + // 5. If jump is NOT taken (`nu_a` = 0), the next `fp` must be the same as the current `fp`. + builder.assert_zero(juz * (AB::Expr::ONE - nu_a) * (next_fp - fp)); + } +} diff --git a/crates/leanVm/src/lib.rs b/crates/leanVm/src/lib.rs index db2bafb5..80389c94 100644 --- a/crates/leanVm/src/lib.rs +++ b/crates/leanVm/src/lib.rs @@ -1,3 +1,4 @@ +pub mod air; pub mod bytecode; pub mod constant; pub mod context;