add: CPU table

Cargo.toml

@@ -8,3 +8,4 @@ edition = "2021"
 plonky2 = { git = "https://github.com/0xPolygonZero/plonky2", rev = "76da138" }
 starky = { git = "https://github.com/0xPolygonZero/plonky2", rev = "76da138" }
 anyhow = "1.0.86"
+array-macro = "1.0.5"

src/e2e_tests.rs

@@ -64,6 +64,7 @@ mod tests {
         let mut config = StarkConfig::standard_fast_config();
         // This needs to be done for tables shorter than `1<<5`. We take
         // a performance hit though!
+        // Number of cap = 2^{cap_height}
         config
             .fri_config
             .cap_height = 1;

src/lib.rs

@@ -17,8 +17,8 @@ mod preflight_simulator;
 mod vm_specs;
 
 // STARK tables -------------
-//#[allow(dead_code)]
-//mod stark_cpu;
+#[allow(dead_code)]
+mod stark_cpu;
 #[allow(dead_code)]
 mod stark_program_instructions;
 

src/preflight_simulator.rs

@@ -18,21 +18,21 @@ use crate::vm_specs::{
 pub struct SimulationRow {
     /// Encodes the instruction executed during this "row". This would
     /// be useful when we go for SN/TARK constraining.
-    instruction: Instruction,
+    pub instruction: Instruction,
 
     /// Clock cycle during execution. Supports large cycle count till
     /// `u32::MAX`
-    clock: u32,
+    pub clock: u32,
 
     /// Address of the program instruction
-    program_counter: u8,
+    pub program_counter: u8,
 
     /// Whether at this row the execution halted. Should only be true
     /// for the last row in any `PreflightSimulation`
-    is_halted: bool,
+    pub is_halted: bool,
 
     /// Registers
-    registers: [u8; REGISTER_COUNT],
+    pub registers: [u8; REGISTER_COUNT],
 
     /// This ideally should be something like `im::HashMap`, see:
     /// https://crates.io/crates/im for immutable collections.
@@ -44,7 +44,7 @@ pub struct SimulationRow {
     ///
     /// However, that optimization is not used for simplicity's sake and
     /// since our VM is small, this is not a large performance hit.
-    memory_snapshot: HashMap<u8, u8>,
+    pub memory_snapshot: HashMap<u8, u8>,
 }
 
 impl SimulationRow {
@@ -186,6 +186,12 @@ impl PreflightSimulation {
     /// Entry point to simulate a program and generate a `PreflightSimulation`
     /// to be used to generate tables
     pub fn simulate(prog: &Program) -> Result<Self> {
+        if prog
+            .code
+            .is_empty()
+        {
+            return Ok(Self { trace_rows: vec![] });
+        }
         let mut trace_rows =
             Vec::with_capacity(Self::MAX_CPU_CYCLES_ALLOWED / 4);
         let first_row = SimulationRow::generate_first_row(prog)?;

src/stark_cpu.rs

@@ -3,6 +3,7 @@
 //! to be linked to the static code "Program" by having a cross-table
 //! -lookup with `ProgramInstructionsStark`.
 
+use array_macro::array;
 use core::marker::PhantomData;
 use plonky2::{
     field::{
@@ -30,19 +31,29 @@ use starky::{
     util::trace_rows_to_poly_values,
 };
 
-use crate::vm_specs::Program;
+use crate::{
+    preflight_simulator::PreflightSimulation,
+    vm_specs::{
+        Instruction,
+        Program,
+    },
+};
 
 // Table description:
-// +-----+----+--------+--------+--------------+---------+
-// | Clk | PC | Reg R0 | Reg R1 | MemoryAddr   | Opcode* |
-// +-----+----+--------+--------+--------------+---------+
-// | ..  | .. | ...    | ...    |  ....        |  ...    |
-// +-----+----+--------+--------+--------------+---------+
+// +-----+----+--------+--------+--------------+---------+-------------+
+// | Clk | PC | Reg R0 | Reg R1 | Location     | Opcode* | Is_Executed |
+// +-----+----+--------+--------+--------------+---------+-------------+
+// | ..  | .. | ...    | ...    |  ....        |  ...    |             |
+// +-----+----+--------+--------+--------------+---------+-------------+
 //
 // `Opcode*` means `Opcode` that is one-hot encoded
-// 5 Columns for `Clk`, `PC`, `Reg R0`, `Reg R1`, `MemoryAccess`
-// 10 Columns for opcodes. See `Instruction::get_opcode`.
-const NUMBER_OF_COLS: usize = 5 + 10;
+// `Location` can be either Memory or Instruction location.
+// 5 Columns for `Clk`, `PC`, `Reg R0`, `Reg R1`, `Location`
+// 11 Columns for opcodes. See `Instruction::get_opcode`.
+// 1 Column for `Is_Executed`
+const NUM_DYNAMIC_COLS: usize = 5;
+const NUM_OPCODE_ONEHOT: usize = 11;
+const NUMBER_OF_COLS: usize = NUM_DYNAMIC_COLS + NUM_OPCODE_ONEHOT + 1;
 const PUBLIC_INPUTS: usize = 0;
 
 #[derive(Clone, Copy)]
@@ -58,30 +69,57 @@ where
         Self { _f: PhantomData }
     }
 
-    pub fn generate_trace(prog: ) -> Vec<PolynomialValues<F>>
+    pub fn generate_trace(sim: &PreflightSimulation) -> Vec<PolynomialValues<F>>
     where
         F: RichField,
     {
-        let mut trace = prog
-            .code
+        let mut trace = sim
+            .trace_rows
             .iter()
-            .map(|(pc, inst)| {
-                [
-                    // Program Counter (ID = 0)
-                    F::from_canonical_u8(*pc),
-                    // Instruction Opcode (ID = 1)
-                    F::from_canonical_u8(inst.get_opcode()),
-                    // Filter, true if actual instructions (ID = 2)
-                    F::ONE,
-                ]
+            .map(|row| {
+                let dynamic_elems = [
+                    // Clock
+                    F::from_canonical_u32(row.clock),
+                    // Program Counter
+                    F::from_canonical_u8(row.program_counter),
+                    // Registers
+                    F::from_canonical_u8(row.registers[0]),
+                    F::from_canonical_u8(row.registers[1]),
+                    // Memory Address (if any accessed)
+                    F::from_canonical_u8(match row.instruction {
+                        Instruction::Jz(_, l) => l.0,
+                        Instruction::Jnz(_, l) => l.0,
+                        Instruction::Lb(_, l) => l.0,
+                        Instruction::Sb(_, l) => l.0,
+                        _ => 0,
+                    }),
+                ];
+                let opcode_one_hot = row
+                    .instruction
+                    .one_hot_encode_and_apply::<F>();
+
+                let mut table_row = [F::ZERO; NUMBER_OF_COLS];
+                let mut idx = 0;
+                for elem in dynamic_elems {
+                    table_row[idx] = elem;
+                    idx += 1;
+                }
+                for elem in opcode_one_hot {
+                    table_row[idx] = elem;
+                    idx += 1;
+                }
+                // `Is_Executed`
+                table_row[NUMBER_OF_COLS - 1] = F::ONE;
+
+                table_row
             })
             .collect::<Vec<[F; NUMBER_OF_COLS]>>();
 
         // Need to pad the trace to a len of some power of 2
         let pow2_len = trace
             .len()
             .next_power_of_two();
-        trace.resize(pow2_len, [F::ZERO, F::ZERO, F::ZERO]);
+        trace.resize(pow2_len, [F::ZERO; NUMBER_OF_COLS]);
 
         // Convert into polynomial values
         trace_rows_to_poly_values(trace)
@@ -166,8 +204,10 @@ mod tests {
             .fri_config
             .cap_height = 1;
         let program = Program::default();
-        let trace =
-            CPUStark::<F, D>::generate_program_instructions_trace(&program);
+        let simulation = PreflightSimulation::simulate(&program);
+        assert!(simulation.is_ok());
+        let simulation = simulation.unwrap();
+        let trace = CPUStark::<F, D>::generate_trace(&simulation);
         let proof: Result<PR, anyhow::Error> = prove(
             stark.clone(),
             &config,

src/vm_specs.rs

@@ -2,6 +2,8 @@
 
 use std::collections::HashMap;
 
+use plonky2::hash::hash_types::RichField;
+
 #[derive(Copy, Clone, Default, Debug, PartialEq)]
 pub enum Register {
     #[default]
@@ -63,6 +65,20 @@ impl Instruction {
             Instruction::Halt => 10,
         }
     }
+
+    /// One-hot encoded description of the Opcode
+    pub fn one_hot_encode(&self) -> [u8; 11] {
+        let mut one_hot_enc = [0; 11];
+        one_hot_enc[self.get_opcode() as usize] = 1;
+        one_hot_enc
+    }
+
+    /// One-hot encodes the opcode and applies a function to it
+    pub fn one_hot_encode_and_apply<F: RichField>(&self) -> [F; 11] {
+        let mut encode_f: [F; 11] = [F::ZERO; 11];
+        encode_f[self.get_opcode() as usize] = F::ONE;
+        encode_f
+    }
 }
 
 #[derive(Clone, Debug, Default, PartialEq)]