Skip to content

Commit 791464e

Browse files
committed
trace: generation wip
1 parent 9c6308f commit 791464e

File tree

15 files changed

+587
-37
lines changed

15 files changed

+587
-37
lines changed

Cargo.toml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,10 +42,13 @@ p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4"
4242
p3-koala-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4" }
4343
p3-air = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4" }
4444
p3-matrix = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4" }
45+
p3-poseidon2 = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4" }
46+
p3-poseidon2-air = { git = "https://github.com/Plonky3/Plonky3.git", rev = "5ebf8e4" }
4547

4648
whir-p3 = { git = "https://github.com/tcoratger/whir-p3.git", rev = "df2241d" }
4749

4850
thiserror = "2.0"
4951
proptest = "1.7"
5052
num-traits = "0.2"
5153
rand = "0.9"
54+
rayon = "1.11"

crates/leanIsa/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
1+
pub mod precompiles;

crates/leanIsa/src/precompiles.rs

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
use std::fmt;
2+
3+
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
4+
pub struct Precompile {
5+
pub name: PrecompileName,
6+
pub n_inputs: usize,
7+
pub n_outputs: usize,
8+
}
9+
10+
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
11+
pub enum PrecompileName {
12+
Poseidon16,
13+
Poseidon24,
14+
DotProduct,
15+
MultilinearEval,
16+
}
17+
18+
impl PrecompileName {
19+
#[must_use]
20+
pub const fn as_str(&self) -> &'static str {
21+
match self {
22+
Self::Poseidon16 => "poseidon16",
23+
Self::Poseidon24 => "poseidon24",
24+
Self::DotProduct => "dot_product",
25+
Self::MultilinearEval => "multilinear_eval",
26+
}
27+
}
28+
}
29+
30+
impl fmt::Display for PrecompileName {
31+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
32+
f.write_str(self.as_str())
33+
}
34+
}
35+
36+
pub const POSEIDON_16: Precompile = Precompile {
37+
name: PrecompileName::Poseidon16,
38+
n_inputs: 2,
39+
n_outputs: 1,
40+
};
41+
42+
pub const POSEIDON_24: Precompile = Precompile {
43+
name: PrecompileName::Poseidon24,
44+
n_inputs: 2,
45+
n_outputs: 1,
46+
};
47+
48+
pub const DOT_PRODUCT: Precompile = Precompile {
49+
name: PrecompileName::DotProduct,
50+
n_inputs: 4,
51+
n_outputs: 0,
52+
};
53+
54+
pub const MULTILINEAR_EVAL: Precompile = Precompile {
55+
name: PrecompileName::MultilinearEval,
56+
n_inputs: 4,
57+
n_outputs: 0,
58+
};
59+
60+
pub const PRECOMPILES: [Precompile; 4] = [POSEIDON_16, POSEIDON_24, DOT_PRODUCT, MULTILINEAR_EVAL];

crates/leanVm/Cargo.toml

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,17 @@ p3-field.workspace = true
1515
p3-symmetric.workspace = true
1616
p3-air.workspace = true
1717
p3-matrix.workspace = true
18+
p3-poseidon2.workspace = true
19+
p3-poseidon2-air.workspace = true
1820

1921
whir-p3.workspace = true
2022

2123
thiserror.workspace = true
2224
num-traits.workspace = true
25+
rayon.workspace = true
26+
rand.workspace = true
27+
28+
lean-isa.workspace = true
2329

2430
[dev-dependencies]
2531
proptest.workspace = true
26-
rand.workspace = true

crates/leanVm/src/air/constant.rs

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,22 @@
1+
use lean_isa::precompiles::PRECOMPILES;
2+
3+
/// The total number of columns that represent a single instruction in the bytecode ROM.
4+
pub(crate) const N_INSTRUCTION_COLUMNS: usize = 15;
5+
/// The number of columns in the execution trace that are directly committed by the prover.
6+
pub(crate) const N_COMMITTED_EXEC_COLUMNS: usize = 5;
7+
/// The number of "virtual" columns that hold memory values, derived via lookups.
8+
///
9+
/// virtual (lookup into memory, with logup*)
10+
pub(crate) const N_MEMORY_VALUE_COLUMNS: usize = 3;
11+
/// The total number of columns related to the execution state (committed + virtual).
12+
pub(crate) const N_EXEC_COLUMNS: usize = N_COMMITTED_EXEC_COLUMNS + N_MEMORY_VALUE_COLUMNS;
13+
/// The number of instruction columns included in the AIR, excluding precompile-specific flags.
14+
pub(crate) const N_INSTRUCTION_COLUMNS_IN_AIR: usize = N_INSTRUCTION_COLUMNS - PRECOMPILES.len();
15+
/// The total width of the AIR table for the main CPU, combining instruction and execution columns.
16+
pub(crate) const N_EXEC_AIR_COLUMNS: usize = N_INSTRUCTION_COLUMNS_IN_AIR + N_EXEC_COLUMNS;
17+
/// The total number of columns in the combined trace (instruction + execution).
18+
pub(crate) const N_TOTAL_COLUMNS: usize = N_INSTRUCTION_COLUMNS + N_EXEC_COLUMNS;
19+
120
// Bytecode columns (looked up from the program ROM)
221

322
/// The immediate value or memory offset for the first operand, `a`.
@@ -32,11 +51,11 @@ pub(crate) const COL_INDEX_PC: usize = N_BYTECODE_COLUMNS + 3;
3251
/// The column for the Frame Pointer (fp) register in the execution trace.
3352
pub(crate) const COL_INDEX_FP: usize = N_BYTECODE_COLUMNS + 4;
3453
/// The column holding the memory address for operand `a` when `FLAG_A` is 0.
35-
pub(crate) const COL_INDEX_ADDR_A: usize = N_BYTECODE_COLUMNS + 5;
54+
pub(crate) const COL_INDEX_MEM_ADDRESS_A: usize = N_BYTECODE_COLUMNS + 5;
3655
/// The column holding the memory address for operand `b` when `FLAG_B` is 0.
37-
pub(crate) const COL_INDEX_ADDR_B: usize = N_BYTECODE_COLUMNS + 6;
56+
pub(crate) const COL_INDEX_MEM_ADDRESS_B: usize = N_BYTECODE_COLUMNS + 6;
3857
/// The column holding the memory address for operand `c` when `FLAG_C` is 0.
39-
pub(crate) const COL_INDEX_ADDR_C: usize = N_BYTECODE_COLUMNS + 7;
58+
pub(crate) const COL_INDEX_MEM_ADDRESS_C: usize = N_BYTECODE_COLUMNS + 7;
4059
/// The total count of columns in the execution trace that are directly committed by the prover.
4160
pub(crate) const N_REAL_EXEC_COLUMNS: usize = 5;
4261

@@ -48,6 +67,3 @@ pub(crate) const COL_INDEX_MEM_VALUE_A: usize = N_BYTECODE_COLUMNS;
4867
pub(crate) const COL_INDEX_MEM_VALUE_B: usize = N_BYTECODE_COLUMNS + 1;
4968
/// The virtual column holding the value read from memory for operand `c`.
5069
pub(crate) const COL_INDEX_MEM_VALUE_C: usize = N_BYTECODE_COLUMNS + 2;
51-
52-
/// The total number of columns in the AIR, including bytecode, real, and virtual columns.
53-
pub(crate) const N_EXEC_AIR_COLUMNS: usize = N_BYTECODE_COLUMNS + 3 + N_REAL_EXEC_COLUMNS;

crates/leanVm/src/air/vm.rs

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,11 @@ use p3_field::PrimeCharacteristicRing;
55
use p3_matrix::Matrix;
66

77
use crate::air::constant::{
8-
COL_INDEX_ADD, COL_INDEX_ADDR_A, COL_INDEX_ADDR_B, COL_INDEX_ADDR_C, COL_INDEX_AUX,
9-
COL_INDEX_DEREF, COL_INDEX_FLAG_A, COL_INDEX_FLAG_B, COL_INDEX_FLAG_C, COL_INDEX_FP,
10-
COL_INDEX_JUZ, COL_INDEX_MEM_VALUE_A, COL_INDEX_MEM_VALUE_B, COL_INDEX_MEM_VALUE_C,
11-
COL_INDEX_MUL, COL_INDEX_OPERAND_A, COL_INDEX_OPERAND_B, COL_INDEX_OPERAND_C, COL_INDEX_PC,
12-
N_EXEC_AIR_COLUMNS,
8+
COL_INDEX_ADD, COL_INDEX_AUX, COL_INDEX_DEREF, COL_INDEX_FLAG_A, COL_INDEX_FLAG_B,
9+
COL_INDEX_FLAG_C, COL_INDEX_FP, COL_INDEX_JUZ, COL_INDEX_MEM_ADDRESS_A,
10+
COL_INDEX_MEM_ADDRESS_B, COL_INDEX_MEM_ADDRESS_C, COL_INDEX_MEM_VALUE_A, COL_INDEX_MEM_VALUE_B,
11+
COL_INDEX_MEM_VALUE_C, COL_INDEX_MUL, COL_INDEX_OPERAND_A, COL_INDEX_OPERAND_B,
12+
COL_INDEX_OPERAND_C, COL_INDEX_PC, N_EXEC_AIR_COLUMNS,
1313
};
1414

1515
/// Virtual Machine AIR
@@ -65,9 +65,9 @@ impl<AB: AirBuilder> Air<AB> for VMAir {
6565
next[COL_INDEX_FP].clone().into(),
6666
);
6767
let (addr_a, addr_b, addr_c) = (
68-
local[COL_INDEX_ADDR_A].clone().into(),
69-
local[COL_INDEX_ADDR_B].clone().into(),
70-
local[COL_INDEX_ADDR_C].clone().into(),
68+
local[COL_INDEX_MEM_ADDRESS_A].clone().into(),
69+
local[COL_INDEX_MEM_ADDRESS_B].clone().into(),
70+
local[COL_INDEX_MEM_ADDRESS_C].clone().into(),
7171
);
7272
let (value_a, value_b, value_c) = (
7373
local[COL_INDEX_MEM_VALUE_A].clone().into(),

crates/leanVm/src/bytecode/instruction/mod.rs

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,24 @@ use deref::DerefInstruction;
33
use dot_product::DotProductInstruction;
44
use jump::JumpIfNotZeroInstruction;
55
use multilinear_eval::MultilinearEvalInstruction;
6+
use p3_field::PrimeCharacteristicRing;
67
use p3_symmetric::Permutation;
78
use poseidon16::Poseidon2_16Instruction;
89
use poseidon24::Poseidon2_24Instruction;
910

11+
use super::operand::{MemOrConstant, MemOrFp, MemOrFpOrConstant};
1012
use crate::{
13+
air::constant::{
14+
COL_INDEX_ADD, COL_INDEX_AUX, COL_INDEX_DEREF, COL_INDEX_FLAG_A, COL_INDEX_FLAG_B,
15+
COL_INDEX_FLAG_C, COL_INDEX_JUZ, COL_INDEX_MUL, COL_INDEX_OPERAND_A, COL_INDEX_OPERAND_B,
16+
COL_INDEX_OPERAND_C, N_INSTRUCTION_COLUMNS,
17+
},
18+
bytecode::operation::Operation,
1119
constant::{DIMENSION, F},
1220
context::run_context::RunContext,
1321
errors::vm::VirtualMachineError,
1422
memory::manager::MemoryManager,
23+
witness::Witness,
1524
};
1625

1726
pub mod computation;
@@ -107,4 +116,142 @@ impl Instruction {
107116
Self::MultilinearEval(instruction) => instruction.execute(run_context, memory_manager),
108117
}
109118
}
119+
120+
/// Generates a witness for a precompile instruction, if applicable.
121+
///
122+
/// This function checks the instruction type and, if it's a precompile, it calls the
123+
/// corresponding `generate_witness` method to capture the necessary data for the proof.
124+
/// For standard instructions, it returns `Ok(None)`.
125+
pub fn generate_witness(
126+
&self,
127+
cycle: usize,
128+
run_context: &RunContext,
129+
memory_manager: &MemoryManager,
130+
) -> Result<Option<Witness>, VirtualMachineError> {
131+
match self {
132+
// If the instruction is Poseidon2_16, generate its witness.
133+
Self::Poseidon2_16(instruction) => Ok(Some(Witness::Poseidon16(
134+
instruction.generate_witness(cycle, run_context, memory_manager)?,
135+
))),
136+
// If the instruction is Poseidon2_24, generate its witness.
137+
Self::Poseidon2_24(instruction) => Ok(Some(Witness::Poseidon24(
138+
instruction.generate_witness(cycle, run_context, memory_manager)?,
139+
))),
140+
// If the instruction is a DotProduct, generate its witness.
141+
Self::DotProduct(instruction) => Ok(Some(Witness::DotProduct(
142+
instruction.generate_witness(cycle, run_context, memory_manager)?,
143+
))),
144+
// If the instruction is a MultilinearEval, generate its witness.
145+
Self::MultilinearEval(instruction) => Ok(Some(Witness::MultilinearEval(
146+
instruction.generate_witness(cycle, run_context, memory_manager)?,
147+
))),
148+
// For all other instruction types, no special witness is generated.
149+
_ => Ok(None),
150+
}
151+
}
152+
153+
/// Converts an instruction into its raw field representation for the execution trace.
154+
#[must_use]
155+
pub fn field_representation(&self) -> Vec<F> {
156+
// Initialize a vector of zeros for all instruction-related columns.
157+
let mut repr = vec![F::ZERO; N_INSTRUCTION_COLUMNS];
158+
159+
// Populate the vector based on the instruction type.
160+
match self {
161+
Self::Computation(ComputationInstruction {
162+
operation,
163+
arg_a,
164+
arg_b,
165+
res,
166+
}) => {
167+
// Set the appropriate opcode flag for ADD or MUL.
168+
match operation {
169+
Operation::Add => repr[COL_INDEX_ADD] = F::ONE,
170+
Operation::Mul => repr[COL_INDEX_MUL] = F::ONE,
171+
}
172+
// Convert operands to their field and flag representations.
173+
let (op_a, flag_a) = op_to_field(arg_a);
174+
let (op_b, flag_b) = op_to_field_fp(arg_b);
175+
let (op_c, flag_c) = op_to_field(res);
176+
177+
// Populate the representation vector.
178+
repr[COL_INDEX_OPERAND_A] = op_a;
179+
repr[COL_INDEX_FLAG_A] = flag_a;
180+
repr[COL_INDEX_OPERAND_B] = op_b;
181+
repr[COL_INDEX_FLAG_B] = flag_b;
182+
repr[COL_INDEX_OPERAND_C] = op_c;
183+
repr[COL_INDEX_FLAG_C] = flag_c;
184+
}
185+
Self::Deref(DerefInstruction {
186+
shift_0,
187+
shift_1,
188+
res,
189+
}) => {
190+
// Set the DEREF opcode flag.
191+
repr[COL_INDEX_DEREF] = F::ONE;
192+
// The first-level pointer is always a memory access from `fp + shift_0`.
193+
repr[COL_INDEX_OPERAND_A] = F::from_usize(*shift_0);
194+
repr[COL_INDEX_FLAG_A] = F::ZERO;
195+
// The second-level offset is always an immediate value.
196+
repr[COL_INDEX_OPERAND_C] = F::from_usize(*shift_1);
197+
repr[COL_INDEX_FLAG_C] = F::ONE;
198+
199+
// Handle the different types of the result operand `res`.
200+
match res {
201+
MemOrFpOrConstant::Constant(c) => {
202+
repr[COL_INDEX_OPERAND_B] = *c;
203+
repr[COL_INDEX_FLAG_B] = F::ONE;
204+
repr[COL_INDEX_AUX] = F::ONE;
205+
}
206+
MemOrFpOrConstant::MemoryAfterFp { offset } => {
207+
repr[COL_INDEX_OPERAND_B] = F::from_usize(*offset);
208+
repr[COL_INDEX_FLAG_B] = F::ZERO;
209+
repr[COL_INDEX_AUX] = F::ONE;
210+
}
211+
MemOrFpOrConstant::Fp => {
212+
repr[COL_INDEX_AUX] = F::ZERO;
213+
}
214+
}
215+
}
216+
Self::JumpIfNotZero(JumpIfNotZeroInstruction {
217+
condition,
218+
dest,
219+
updated_fp,
220+
}) => {
221+
// Set the JUZ opcode flag.
222+
repr[COL_INDEX_JUZ] = F::ONE;
223+
// Convert operands to their field and flag representations.
224+
let (op_a, flag_a) = op_to_field(condition);
225+
let (op_b, flag_b) = op_to_field_fp(updated_fp);
226+
let (op_c, flag_c) = op_to_field(dest);
227+
228+
// Populate the representation vector.
229+
repr[COL_INDEX_OPERAND_A] = op_a;
230+
repr[COL_INDEX_FLAG_A] = flag_a;
231+
repr[COL_INDEX_OPERAND_B] = op_b;
232+
repr[COL_INDEX_FLAG_B] = flag_b;
233+
repr[COL_INDEX_OPERAND_C] = op_c;
234+
repr[COL_INDEX_FLAG_C] = flag_c;
235+
}
236+
// Precompiles do not set flags in the main instruction columns.
237+
_ => {}
238+
}
239+
repr
240+
}
241+
}
242+
243+
/// Helper to convert a `MemOrConstant` operand into its (value, flag) pair.
244+
fn op_to_field(op: &MemOrConstant) -> (F, F) {
245+
match op {
246+
MemOrConstant::Constant(c) => (*c, F::ONE),
247+
MemOrConstant::MemoryAfterFp { offset } => (F::from_usize(*offset), F::ZERO),
248+
}
249+
}
250+
251+
/// Helper to convert a `MemOrFp` operand into its (value, flag) pair.
252+
fn op_to_field_fp(op: &MemOrFp) -> (F, F) {
253+
match op {
254+
MemOrFp::Fp => (F::ZERO, F::ONE), // Represents the `fp` register itself.
255+
MemOrFp::MemoryAfterFp { offset } => (F::from_usize(*offset), F::ZERO),
256+
}
110257
}

crates/leanVm/src/bytecode/operand.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ pub enum MemOrConstant {
1010
/// Represents the scalar value at `m[fp + shift]`.
1111
MemoryAfterFp {
1212
/// The offset from `fp` where the memory location is located.
13-
shift: usize,
13+
offset: usize,
1414
},
1515
}
1616

@@ -20,7 +20,7 @@ pub enum MemOrFpOrConstant {
2020
/// A memory location specified by a positive offset from `fp`. Represents `m[fp + shift]`.
2121
MemoryAfterFp {
2222
/// The offset from `fp` where the memory location is located.
23-
shift: usize,
23+
offset: usize,
2424
},
2525
/// The value of the frame pointer (`fp`) register itself.
2626
Fp,
@@ -34,7 +34,7 @@ pub enum MemOrFp {
3434
/// A memory location specified by a positive offset from `fp`. Represents `m[fp + shift]`.
3535
MemoryAfterFp {
3636
/// The offset from `fp` where the memory location is located.
37-
shift: usize,
37+
offset: usize,
3838
},
3939
/// The value of the frame pointer (`fp`) register itself.
4040
Fp,

crates/leanVm/src/bytecode/program.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ pub struct Program {
1818
/// Represents the results of a successful program execution.
1919
#[derive(Debug)]
2020
pub struct ExecutionResult {
21+
/// The size of the public memory region.
22+
pub public_memory_size: usize,
2123
/// The final state of the memory manager.
2224
pub memory_manager: MemoryManager,
2325
/// The execution trace of the program counter (pc) values.

0 commit comments

Comments
 (0)