vm: better building of the execution trace

crates/zk_vm/src/dot_product_air.rs

@@ -5,7 +5,7 @@ use p3_field::PrimeCharacteristicRing;
 use p3_matrix::Matrix;
 use vm::EF;
 
-use crate::execution_trace::WitnessDotProduct;
+use crate::witness::dot_product::WitnessDotProduct;
 
 /*
 | StartFlag | Len | IndexA | IndexB | IndexRes | ValueA | ValueB | Res           | Computation                   |

crates/zk_vm/src/lib.rs

@@ -10,6 +10,7 @@ mod instruction_encoder;
 mod poseidon_tables;
 pub mod prove_execution;
 pub mod verify_execution;
+pub mod witness;
 
 #[cfg(test)]
 pub mod recursion;

crates/zk_vm/src/poseidon_tables.rs

@@ -4,7 +4,7 @@ use utils::{
 };
 use vm::F;
 
-use crate::execution_trace::{WitnessPoseidon16, WitnessPoseidon24};
+use crate::witness::poseidon::{WitnessPoseidon16, WitnessPoseidon24};
 
 pub fn build_poseidon_columns(
     poseidons_16: &[WitnessPoseidon16],

crates/zk_vm/src/prove_execution.rs

@@ -35,7 +35,7 @@ use crate::{
     },
     dot_product_air::{DOT_PRODUCT_AIR_COLUMN_GROUPS, DotProductAir, build_dot_product_columns},
     exec_column_groups,
-    execution_trace::{ExecutionTrace, get_execution_trace},
+    execution_trace::ExecutionTrace,
     poseidon_tables::{all_poseidon_16_indexes, all_poseidon_24_indexes, build_poseidon_columns},
 };
 
@@ -57,7 +57,7 @@ pub fn prove_execution(
         memory,
     } = info_span!("Witness generation").in_scope(|| {
         let execution_result = bytecode.execute(public_input, private_input);
-        get_execution_trace(bytecode, &execution_result)
+        ExecutionTrace::new(bytecode, &execution_result)
     });
 
     let public_memory = &memory[..public_memory_size];

crates/zk_vm/src/witness/dot_product.rs

@@ -0,0 +1,22 @@
+use vm::EF;
+
+/// Holds the high-level witness data for a single dot product precompile execution.
+#[derive(Debug)]
+pub struct WitnessDotProduct {
+    /// The CPU cycle at which the dot product operation is initiated.
+    pub cycle: usize,
+    /// The starting memory address (vectorized pointer) of the first input slice.
+    pub addr_0: usize,
+    /// The starting memory address (vectorized pointer) of the second input slice.
+    pub addr_1: usize,
+    /// The memory address (vectorized pointer) where the final result is stored.
+    pub addr_res: usize,
+    /// The number of elements in each input slice.
+    pub len: usize,
+    /// The actual data values of the first input slice.
+    pub slice_0: Vec<EF>,
+    /// The actual data values of the second input slice.
+    pub slice_1: Vec<EF>,
+    /// The final computed result of the dot product.
+    pub res: EF,
+}

crates/zk_vm/src/witness/mod.rs

@@ -0,0 +1,16 @@
+use dot_product::WitnessDotProduct;
+use multilinear_eval::WitnessMultilinearEval;
+use poseidon::{WitnessPoseidon16, WitnessPoseidon24};
+
+pub mod dot_product;
+pub mod multilinear_eval;
+pub mod poseidon;
+
+/// An enum to encapsulate any possible precompile witness.
+#[derive(Debug)]
+pub enum Witness {
+    Poseidon16(WitnessPoseidon16),
+    Poseidon24(WitnessPoseidon24),
+    DotProduct(WitnessDotProduct),
+    MultilinearEval(WitnessMultilinearEval),
+}

crates/zk_vm/src/witness/multilinear_eval.rs

@@ -0,0 +1,20 @@
+use vm::EF;
+
+/// Holds the high-level witness data for a single multilinear evaluation precompile.
+#[derive(Debug)]
+pub struct WitnessMultilinearEval {
+    /// The CPU cycle at which this operation is initiated.
+    pub cycle: usize,
+    /// The memory address of the polynomial's coefficients.
+    pub addr_coeffs: usize,
+    /// The memory address of the evaluation point's coordinates.
+    pub addr_point: usize,
+    /// The memory address where the final result is stored.
+    pub addr_res: usize,
+    /// The number of variables in the multilinear polynomial.
+    pub n_vars: usize,
+    /// The coordinates of the evaluation point.
+    pub point: Vec<EF>,
+    /// The final computed result of the evaluation.
+    pub res: EF,
+}

crates/zk_vm/src/witness/poseidon.rs

@@ -0,0 +1,35 @@
+use vm::F;
+
+/// Holds the high-level witness data for a single Poseidon2 permutation over 16 elements.
+#[derive(Debug)]
+pub struct WitnessPoseidon16 {
+    /// The CPU cycle at which this operation is initiated, if applicable.
+    pub cycle: Option<usize>,
+    /// The memory address (vectorized pointer, of size 1) of the first 8-element input vector.
+    pub addr_input_a: usize,
+    /// The memory address (vectorized pointer, of size 1) of the second 8-element input vector.
+    pub addr_input_b: usize,
+    /// The memory address (vectorized pointer, of size 2) where the two 8-element output vectors are stored.
+    pub addr_output: usize,
+    /// The full 16-element input state for the permutation.
+    pub input: [F; 16],
+    /// The full 16-element output state resulting from the permutation.
+    pub output: [F; 16],
+}
+
+/// Holds the high-level witness data for a single Poseidon2 permutation over 24 elements.
+#[derive(Debug)]
+pub struct WitnessPoseidon24 {
+    /// The CPU cycle at which this operation is initiated, if applicable.
+    pub cycle: Option<usize>,
+    /// The memory address (vectorized pointer, of size 2) of the first two 8-element input vectors.
+    pub addr_input_a: usize,
+    /// The memory address (vectorized pointer, of size 1) of the third 8-element input vector.
+    pub addr_input_b: usize,
+    /// The memory address (vectorized pointer, of size 1) where the relevant 8-element output vector is stored.
+    pub addr_output: usize,
+    /// The full 24-element input state for the permutation.
+    pub input: [F; 24],
+    /// The last 8 elements of the 24-element output state from the permutation.
+    pub output: [F; 8],
+}