|
| 1 | +//! zkVM AIR (Algebraic Intermediate Representation) |
| 2 | +//! |
| 3 | +//! This module defines the transition constraints for the custom zkVM's Instruction Set |
| 4 | +//! Architecture (ISA). It translates the VM's operational semantics into a set of polynomial |
| 5 | +//! equations that must hold true between consecutive states (rows) of the execution trace. |
| 6 | +
|
| 7 | +use std::borrow::Borrow; |
| 8 | + |
| 9 | +use constant::{ |
| 10 | + COL_INDEX_ADD, COL_INDEX_ADDR_A, COL_INDEX_ADDR_B, COL_INDEX_ADDR_C, COL_INDEX_AUX, |
| 11 | + COL_INDEX_DEREF, COL_INDEX_FLAG_A, COL_INDEX_FLAG_B, COL_INDEX_FLAG_C, COL_INDEX_FP, |
| 12 | + COL_INDEX_JUZ, COL_INDEX_MEM_VALUE_A, COL_INDEX_MEM_VALUE_B, COL_INDEX_MEM_VALUE_C, |
| 13 | + COL_INDEX_MUL, COL_INDEX_OPERAND_A, COL_INDEX_OPERAND_B, COL_INDEX_OPERAND_C, COL_INDEX_PC, |
| 14 | + N_EXEC_AIR_COLUMNS, |
| 15 | +}; |
| 16 | +use p3_air::{Air, AirBuilder, BaseAir}; |
| 17 | +use p3_field::PrimeCharacteristicRing; |
| 18 | +use p3_matrix::Matrix; |
| 19 | + |
| 20 | +pub mod constant; |
| 21 | + |
| 22 | +/// Virtual Machine AIR |
| 23 | +#[derive(Debug)] |
| 24 | +pub struct VMAir; |
| 25 | + |
| 26 | +impl<F> BaseAir<F> for VMAir { |
| 27 | + /// The total number of columns in the execution trace. |
| 28 | + fn width(&self) -> usize { |
| 29 | + N_EXEC_AIR_COLUMNS |
| 30 | + } |
| 31 | +} |
| 32 | + |
| 33 | +impl<AB: AirBuilder> Air<AB> for VMAir { |
| 34 | + #[inline] |
| 35 | + fn eval(&self, builder: &mut AB) { |
| 36 | + // Get a view of the main execution trace. |
| 37 | + let main = builder.main(); |
| 38 | + |
| 39 | + // Get the current row (`local`) and the next row (`next`) from the trace. |
| 40 | + let local = main.row_slice(0).unwrap(); |
| 41 | + let local = local.borrow(); |
| 42 | + |
| 43 | + let next = main.row_slice(1).unwrap(); |
| 44 | + let next = next.borrow(); |
| 45 | + |
| 46 | + // INSTRUCTION DECODING |
| 47 | + // |
| 48 | + // Extract instruction fields (operands, flags, and opcodes) from the local row. |
| 49 | + // |
| 50 | + // These are treated as constants for a given row, looked up from the bytecode. |
| 51 | + let operand_a = local[COL_INDEX_OPERAND_A].clone().into(); |
| 52 | + let operand_b = local[COL_INDEX_OPERAND_B].clone().into(); |
| 53 | + let operand_c = local[COL_INDEX_OPERAND_C].clone().into(); |
| 54 | + let flag_a = local[COL_INDEX_FLAG_A].clone().into(); |
| 55 | + let flag_b = local[COL_INDEX_FLAG_B].clone().into(); |
| 56 | + let flag_c = local[COL_INDEX_FLAG_C].clone().into(); |
| 57 | + let add = local[COL_INDEX_ADD].clone().into(); |
| 58 | + let mul = local[COL_INDEX_MUL].clone().into(); |
| 59 | + let deref = local[COL_INDEX_DEREF].clone().into(); |
| 60 | + let juz = local[COL_INDEX_JUZ].clone().into(); |
| 61 | + let aux = local[COL_INDEX_AUX].clone().into(); |
| 62 | + |
| 63 | + // REGISTER & MEMORY VALUES |
| 64 | + // |
| 65 | + // Extract register values and memory access data from the trace. |
| 66 | + let (pc, next_pc) = ( |
| 67 | + local[COL_INDEX_PC].clone().into(), |
| 68 | + next[COL_INDEX_PC].clone().into(), |
| 69 | + ); |
| 70 | + let (fp, next_fp) = ( |
| 71 | + local[COL_INDEX_FP].clone().into(), |
| 72 | + next[COL_INDEX_FP].clone().into(), |
| 73 | + ); |
| 74 | + let (addr_a, addr_b, addr_c) = ( |
| 75 | + local[COL_INDEX_ADDR_A].clone().into(), |
| 76 | + local[COL_INDEX_ADDR_B].clone().into(), |
| 77 | + local[COL_INDEX_ADDR_C].clone().into(), |
| 78 | + ); |
| 79 | + let (value_a, value_b, value_c) = ( |
| 80 | + local[COL_INDEX_MEM_VALUE_A].clone().into(), |
| 81 | + local[COL_INDEX_MEM_VALUE_B].clone().into(), |
| 82 | + local[COL_INDEX_MEM_VALUE_C].clone().into(), |
| 83 | + ); |
| 84 | + |
| 85 | + // OPERAND RECONSTRUCTION |
| 86 | + // |
| 87 | + // Compute the effective values of the operands (`nu_a`, `nu_b`, `nu_c`). |
| 88 | + // |
| 89 | + // Each operand can be either: |
| 90 | + // - an immediate value (from `operand_x`), |
| 91 | + // - a value loaded from memory (from `value_x`), selected by `flag_x`. |
| 92 | + // |
| 93 | + // Formula: nu_x = flag_x * operand_x + (1 - flag_x) * value_x |
| 94 | + let nu_a = |
| 95 | + flag_a.clone() * operand_a.clone() + (AB::Expr::ONE - flag_a.clone()) * value_a.clone(); |
| 96 | + let nu_b = flag_b.clone() * operand_b.clone() + (AB::Expr::ONE - flag_b.clone()) * value_b; |
| 97 | + // Operand `c` is special: its immediate value can be the frame pointer `fp` itself. |
| 98 | + let nu_c = flag_c.clone() * fp.clone() + (AB::Expr::ONE - flag_c.clone()) * value_c.clone(); |
| 99 | + |
| 100 | + // MEMORY ADDRESS CONSTRAINTS |
| 101 | + // |
| 102 | + // Enforce that if an operand is loaded from memory (`flag_x` = 0), its address |
| 103 | + // (`addr_x`) must correctly correspond to the frame pointer plus its offset (`fp + operand_x`). |
| 104 | + // |
| 105 | + // If `flag_x` = 1, the constraint is `0 * ... = 0`, so it is satisfied. |
| 106 | + builder.assert_zero((AB::Expr::ONE - flag_a) * (addr_a - (fp.clone() + operand_a))); |
| 107 | + builder.assert_zero((AB::Expr::ONE - flag_b) * (addr_b - (fp.clone() + operand_b))); |
| 108 | + builder.assert_zero( |
| 109 | + (AB::Expr::ONE - flag_c) * (addr_c.clone() - (fp.clone() + operand_c.clone())), |
| 110 | + ); |
| 111 | + |
| 112 | + // INSTRUCTION CONSTRAINTS |
| 113 | + |
| 114 | + // ADD Instruction Constraint |
| 115 | + // |
| 116 | + // If the `add` opcode is active, enforce `nu_a + nu_c = nu_b`. |
| 117 | + builder.assert_zero(add * (nu_b.clone() - (nu_a.clone() + nu_c.clone()))); |
| 118 | + |
| 119 | + // MUL Instruction Constraint |
| 120 | + // |
| 121 | + // If the `mul` opcode is active, enforce `nu_a * nu_c = nu_b`. |
| 122 | + builder.assert_zero(mul * (nu_b.clone() - nu_a.clone() * nu_c.clone())); |
| 123 | + |
| 124 | + // DEREF Instruction Constraints |
| 125 | + // |
| 126 | + // This instruction computes `m[m[fp + α] + β] = res`. |
| 127 | + // |
| 128 | + // 1. Enforce that the final address `addr_c` is computed correctly: `addr_c = value_a + operand_c`. |
| 129 | + // |
| 130 | + // Here, `value_a` is the pointer `m[fp + α]` and `operand_c` is the offset `β`. |
| 131 | + builder.assert_zero(deref.clone() * (addr_c - (value_a + operand_c))); |
| 132 | + // 2. If `aux` = 1, the result `res` is a normal value (`nu_b`). Enforce `value_c = nu_b`. |
| 133 | + builder.assert_zero(deref.clone() * aux.clone() * (value_c.clone() - nu_b.clone())); |
| 134 | + // 3. If `aux` = 0, the result `res` is the frame pointer itself. Enforce `value_c = fp`. |
| 135 | + builder.assert_zero(deref * (AB::Expr::ONE - aux) * (value_c - fp.clone())); |
| 136 | + |
| 137 | + // JUZ (Jump Unless Zero) and Program Flow Constraints |
| 138 | + |
| 139 | + // Default Program Flow (No Jump) |
| 140 | + // |
| 141 | + // If `juz` is not active, the program counter must increment by 1. |
| 142 | + builder.assert_zero( |
| 143 | + (AB::Expr::ONE - juz.clone()) * (next_pc.clone() - (pc.clone() + AB::Expr::ONE)), |
| 144 | + ); |
| 145 | + // If `juz` is not active, the frame pointer must remain unchanged. |
| 146 | + builder.assert_zero((AB::Expr::ONE - juz.clone()) * (next_fp.clone() - fp.clone())); |
| 147 | + |
| 148 | + // JUZ Active Program Flow |
| 149 | + // |
| 150 | + // The condition for the jump is `nu_a`. |
| 151 | + // 1. Enforce that the condition `nu_a` is boolean (0 or 1). |
| 152 | + builder.assert_zero(juz.clone() * nu_a.clone() * (AB::Expr::ONE - nu_a.clone())); |
| 153 | + // 2. If jump is taken (`nu_a` = 1), the next `pc` must be the destination `nu_b`. |
| 154 | + builder.assert_zero(juz.clone() * nu_a.clone() * (next_pc.clone() - nu_b)); |
| 155 | + // 3. If jump is taken (`nu_a` = 1), the next `fp` must be the new frame pointer `nu_c`. |
| 156 | + builder.assert_zero(juz.clone() * nu_a.clone() * (next_fp.clone() - nu_c)); |
| 157 | + // 4. If jump is NOT taken (`nu_a` = 0), the next `pc` must be `pc + 1`. |
| 158 | + builder.assert_zero( |
| 159 | + juz.clone() * (AB::Expr::ONE - nu_a.clone()) * (next_pc - (pc + AB::Expr::ONE)), |
| 160 | + ); |
| 161 | + // 5. If jump is NOT taken (`nu_a` = 0), the next `fp` must be the same as the current `fp`. |
| 162 | + builder.assert_zero(juz * (AB::Expr::ONE - nu_a) * (next_fp - fp)); |
| 163 | + } |
| 164 | +} |
0 commit comments