Skip to content

refactor(prover): add VM table trace writers #703

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MauroToscano merged 1 commit into from
Jun 23, 2026
Merged

refactor(prover): add VM table trace writers #703

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 20 additions & 21 deletions prover/src/tables/bitwise.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ use stark::trace::{TraceTable, columns2rows};
#[cfg(feature = "parallel")]
use rayon::prelude::*;

use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, alu_op};
use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, VmTable, alu_op};

// =========================================================================
// Column indices for BITWISE table
Expand Down Expand Up @@ -357,38 +357,37 @@ pub fn preprocessed_commitment(options: &ProofOptions) -> Commitment {
/// All output columns are precomputed. Multiplicity columns are initialized
/// to zero and will be updated when other tables send lookups.
pub fn generate_bitwise_trace() -> TraceTable<GoldilocksField, GoldilocksExtension> {
let mut data = vec![FE::zero(); NUM_ROWS * cols::NUM_COLUMNS];
let mut trace = TraceTable::new_main(
vec![FE::zero(); NUM_ROWS * cols::NUM_COLUMNS],
cols::NUM_COLUMNS,
1,
);
let table = &mut trace.main_table;

for x in 0u32..256 {
for y in 0u32..256 {
for z in 0u32..16 {
let row_idx = (x as usize) + (y as usize) * 256 + (z as usize) * 256 * 256;
let base = row_idx * cols::NUM_COLUMNS;

// Input columns
data[base + cols::X] = FE::from(x as u64);
data[base + cols::Y] = FE::from(y as u64);
data[base + cols::Z] = FE::from(z as u64);
table.set_byte(row_idx, cols::X, x as u8);
table.set_byte(row_idx, cols::Y, y as u8);
table.set_byte(row_idx, cols::Z, z as u8);

// Bitwise operation results
data[base + cols::AND] = FE::from((x & y) as u64);
data[base + cols::OR] = FE::from((x | y) as u64);
data[base + cols::XOR] = FE::from((x ^ y) as u64);
table.set_byte(row_idx, cols::AND, (x & y) as u8);
table.set_byte(row_idx, cols::OR, (x | y) as u8);
table.set_byte(row_idx, cols::XOR, (x ^ y) as u8);

// MSB extractions
let msb8 = (x >> 7) & 1;
let halfword = x + y * 256;
let msb16 = (halfword >> 15) & 1;
data[base + cols::MSB8] = FE::from(msb8 as u64);
data[base + cols::MSB16] = FE::from(msb16 as u64);
table.set_bool(row_idx, cols::MSB8, msb8 == 1);
table.set_bool(row_idx, cols::MSB16, msb16 == 1);

// Zero check (X + 256*Y + 65536*Z must be zero)
let is_zero = if x == 0 && y == 0 && z == 0 {
1u64
} else {
0u64
};
data[base + cols::ZERO] = FE::from(is_zero);
table.set_bool(row_idx, cols::ZERO, x == 0 && y == 0 && z == 0);

// Shift operations on halfword
let sll = if z == 0 {
Expand All @@ -397,16 +396,16 @@ pub fn generate_bitwise_trace() -> TraceTable<GoldilocksField, GoldilocksExtensi
(halfword << z) & 0xFFFF
};
let sllc = if z == 0 { 0 } else { halfword >> (16 - z) };
data[base + cols::SLL] = FE::from(sll as u64);
data[base + cols::SLLC] = FE::from(sllc as u64);
table.set_half(row_idx, cols::SLL, sll as u16);
table.set_half(row_idx, cols::SLLC, sllc as u16);

// Multiplicity columns start at zero
// They will be updated by update_multiplicities()
}
}
}

TraceTable::new_main(data, cols::NUM_COLUMNS, 1)
trace
}

/// Computes the row index for a given (X, Y, Z) tuple.
Expand Down Expand Up @@ -444,7 +443,7 @@ pub fn update_multiplicities(

// Increment multiplicity
let current = trace.main_table.get_row(row)[mu_col];
trace.set_main(row, mu_col, current + FE::one());
trace.main_table.set_fe(row, mu_col, current + FE::one());
}
}

Expand Down
55 changes: 24 additions & 31 deletions prover/src/tables/branch.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ use stark::lookup::{BusInteraction, BusValue, LinearTerm, Multiplicity, Packing}
use stark::table::TableView;
use stark::trace::TraceTable;

use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, SHIFT_16, alu_op};
use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, SHIFT_16, VmTable, alu_op};

// =========================================================================
// Column indices for BRANCH table
Expand Down Expand Up @@ -168,23 +168,14 @@ pub fn generate_branch_trace(

let unique_ops: Vec<_> = op_map.into_iter().collect();
let num_rows = unique_ops.len().next_power_of_two().max(4);
let mut data = vec![FE::zero(); num_rows * cols::NUM_COLUMNS];
let mut trace = TraceTable::new_main(
vec![FE::zero(); num_rows * cols::NUM_COLUMNS],
cols::NUM_COLUMNS,
1,
);
let table = &mut trace.main_table;

for (row_idx, (op, multiplicity)) in unique_ops.iter().enumerate() {
let base = row_idx * cols::NUM_COLUMNS;

// Extract pc as DWordWL: [Word, Word]
let pc_0 = (op.pc & 0xFFFF_FFFF) as u32;
let pc_1 = (op.pc >> 32) as u32;

// Extract offset as DWordWL: [Word, Word]
let offset_0 = (op.offset & 0xFFFF_FFFF) as u32;
let offset_1 = (op.offset >> 32) as u32;

// Extract register as DWordWL: [Word, Word]
let register_0 = (op.register & 0xFFFF_FFFF) as u32;
let register_1 = (op.register >> 32) as u32;

// Compute next_pc
let next_pc_unmasked = op.compute_next_pc_unmasked();
let next_pc = op.compute_next_pc();
Expand All @@ -203,23 +194,25 @@ pub fn generate_branch_trace(
let next_pc_high_2 = ((next_pc >> 48) & 0xFFFF) as u16;

// Store columns
data[base + cols::PC_0] = FE::from(pc_0 as u64);
data[base + cols::PC_1] = FE::from(pc_1 as u64);
data[base + cols::OFFSET_0] = FE::from(offset_0 as u64);
data[base + cols::OFFSET_1] = FE::from(offset_1 as u64);
data[base + cols::REGISTER_0] = FE::from(register_0 as u64);
data[base + cols::REGISTER_1] = FE::from(register_1 as u64);
data[base + cols::JALR] = FE::from(if op.jalr { 1u64 } else { 0u64 });
data[base + cols::NEXT_PC_HIGH_0] = FE::from(next_pc_high_0 as u64);
data[base + cols::NEXT_PC_HIGH_1] = FE::from(next_pc_high_1 as u64);
data[base + cols::NEXT_PC_HIGH_2] = FE::from(next_pc_high_2 as u64);
data[base + cols::NEXT_PC_LOW_0] = FE::from(next_pc_low_0 as u64);
data[base + cols::NEXT_PC_LOW_1] = FE::from(next_pc_low_1 as u64);
data[base + cols::UNMASKED_LOW_BYTE] = FE::from(unmasked_low_byte as u64);
data[base + cols::MU] = FE::from(*multiplicity);
table.set_dword_wl(row_idx, cols::PC_0, op.pc);
table.set_dword_wl(row_idx, cols::OFFSET_0, op.offset);
table.set_dword_wl(row_idx, cols::REGISTER_0, op.register);
table.set_bool(row_idx, cols::JALR, op.jalr);
table.set_halves(
row_idx,
cols::NEXT_PC_HIGH_0,
&[next_pc_high_0, next_pc_high_1, next_pc_high_2],
);
table.set_bytes(
row_idx,
cols::NEXT_PC_LOW_0,
&[next_pc_low_0, next_pc_low_1],
);
table.set_byte(row_idx, cols::UNMASKED_LOW_BYTE, unmasked_low_byte);
table.set_u64(row_idx, cols::MU, *multiplicity);
}

TraceTable::new_main(data, cols::NUM_COLUMNS, 1)
trace
}

// =========================================================================
Expand Down
24 changes: 13 additions & 11 deletions prover/src/tables/bytewise.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
use stark::lookup::{BusInteraction, BusValue, Multiplicity, Packing};
use stark::trace::TraceTable;

use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, alu_op};
use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, VmTable, alu_op};

// =========================================================================
// Column indices for BYTEWISE table
Expand Down Expand Up @@ -106,22 +106,24 @@ pub fn generate_bytewise_trace(

let unique_ops: Vec<_> = op_map.into_iter().collect();
let num_rows = unique_ops.len().next_power_of_two().max(4);
let mut data = vec![FE::zero(); num_rows * cols::NUM_COLUMNS];
let mut trace = TraceTable::new_main(
vec![FE::zero(); num_rows * cols::NUM_COLUMNS],
cols::NUM_COLUMNS,
1,
);
let table = &mut trace.main_table;

for (row_idx, (op, multiplicity)) in unique_ops.iter().enumerate() {
let base = row_idx * cols::NUM_COLUMNS;
let res = op.compute_res();

for i in 0..8 {
data[base + cols::A[i]] = FE::from((op.a >> (8 * i)) & 0xFF);
data[base + cols::B[i]] = FE::from((op.b >> (8 * i)) & 0xFF);
data[base + cols::RES[i]] = FE::from((res >> (8 * i)) & 0xFF);
}
data[base + cols::OP] = FE::from(op.op as u64);
data[base + cols::MU] = FE::from(*multiplicity);
table.set_dword_bl(row_idx, cols::A[0], op.a);
table.set_dword_bl(row_idx, cols::B[0], op.b);
table.set_dword_bl(row_idx, cols::RES[0], res);
table.set_byte(row_idx, cols::OP, op.op);
table.set_u64(row_idx, cols::MU, *multiplicity);
}

TraceTable::new_main(data, cols::NUM_COLUMNS, 1)
trace
}

// =========================================================================
Expand Down
47 changes: 20 additions & 27 deletions prover/src/tables/commit.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ use stark::trace::TraceTable;

use crate::constraints::templates::{AddConstraint, AddOperand};

use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField};
use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, VmTable};

// =========================================================================
// Column indices for COMMIT table
Expand Down Expand Up @@ -164,70 +164,63 @@ pub fn generate_commit_trace(
) -> TraceTable<GoldilocksField, GoldilocksExtension> {
let n = ops.len();
let num_rows = n.next_power_of_two().max(4);
let mut data = vec![FE::zero(); num_rows * cols::NUM_COLUMNS];
let mut trace = TraceTable::new_main(
vec![FE::zero(); num_rows * cols::NUM_COLUMNS],
cols::NUM_COLUMNS,
1,
);
let table = &mut trace.main_table;

for (row_idx, op) in ops.iter().enumerate() {
let base = row_idx * cols::NUM_COLUMNS;

// Timestamp (DWordWL)
data[base + cols::TIMESTAMP_0] = FE::from(op.timestamp & 0xFFFF_FFFF);
data[base + cols::TIMESTAMP_1] = FE::from(op.timestamp >> 32);
table.set_dword_wl(row_idx, cols::TIMESTAMP_0, op.timestamp);

// Index (BaseField)
data[base + cols::INDEX] = FE::from(op.index);
table.set_u64(row_idx, cols::INDEX, op.index);

// Address (DWordWL)
data[base + cols::ADDRESS_0] = FE::from(op.address & 0xFFFF_FFFF);
data[base + cols::ADDRESS_1] = FE::from(op.address >> 32);
table.set_dword_wl(row_idx, cols::ADDRESS_0, op.address);

// address_incr = address + 1 (DWordHL: 4 halfwords)
let address_incr = op.address.wrapping_add(1);
data[base + cols::ADDRESS_INCR_0] = FE::from(address_incr & 0xFFFF);
data[base + cols::ADDRESS_INCR_1] = FE::from((address_incr >> 16) & 0xFFFF);
data[base + cols::ADDRESS_INCR_2] = FE::from((address_incr >> 32) & 0xFFFF);
data[base + cols::ADDRESS_INCR_3] = FE::from((address_incr >> 48) & 0xFFFF);
table.set_dword_hl(row_idx, cols::ADDRESS_INCR_0, address_incr);

// Count (DWordWL)
data[base + cols::COUNT_0] = FE::from(op.count & 0xFFFF_FFFF);
data[base + cols::COUNT_1] = FE::from(op.count >> 32);
table.set_dword_wl(row_idx, cols::COUNT_0, op.count);

// count_decr: if count == 0, use 0xFFFF_FFFF_FFFF_FFFF; else count - 1
let count_decr = if op.count == 0 {
u64::MAX
} else {
op.count - 1
};
data[base + cols::COUNT_DECR_0] = FE::from(count_decr & 0xFFFF);
data[base + cols::COUNT_DECR_1] = FE::from((count_decr >> 16) & 0xFFFF);
data[base + cols::COUNT_DECR_2] = FE::from((count_decr >> 32) & 0xFFFF);
data[base + cols::COUNT_DECR_3] = FE::from((count_decr >> 48) & 0xFFFF);
table.set_dword_hl(row_idx, cols::COUNT_DECR_0, count_decr);

// Control bits
data[base + cols::FIRST] = FE::from(op.first as u64);
data[base + cols::END] = FE::from(op.end as u64);
table.set_bool(row_idx, cols::FIRST, op.first);
table.set_bool(row_idx, cols::END, op.end);

// Value
data[base + cols::VALUE] = FE::from(op.value as u64);
table.set_byte(row_idx, cols::VALUE, op.value);

// mu = 1 for all real rows (first, middle, and end rows)
data[base + cols::MU] = FE::one();
table.set_fe(row_idx, cols::MU, FE::one());
}

// Padding rows: spec requires count=1 and address_incr=[1,0,0,0] so
// the unconditional ADD/SUB templates have valid carry values.
// count=1 → count_decr=0 (all halfwords zero), address=0 → address_incr=1.
for row_idx in n..num_rows {
let base = row_idx * cols::NUM_COLUMNS;
// count = 1 (low word)
data[base + cols::COUNT_0] = FE::one();
table.set_fe(row_idx, cols::COUNT_0, FE::one());
// address_incr halfword 0 = 1 (address=0, so address+1 = 1)
data[base + cols::ADDRESS_INCR_0] = FE::one();
table.set_fe(row_idx, cols::ADDRESS_INCR_0, FE::one());
// All other fields remain zero: timestamp=0, address=0, count_1=0,
// count_decr=[0,0,0,0], first=0, end=0, value=0, mu=0,
// address_incr_1..3=0
}

TraceTable::new_main(data, cols::NUM_COLUMNS, 1)
trace
}

// =========================================================================
Expand Down
Loading
Loading