Morgan/issue/155

basic/src/lib.rs

@@ -47,7 +47,7 @@ use valida_machine::__internal::{
 };
 use valida_machine::{
     generate_permutation_trace, verify_constraints, AdviceProvider, BusArgument, Chip, ChipProof,
-    Commitments, Instruction, Machine, MachineProof, OpenedValues, ProgramROM, StoppingFlag,
+    Commitments, Instruction, Machine, MachineProof, OpenedValues, ProgramROM, StoppingFlag, ColumnVector, ColumnIndex,
     ValidaAirBuilder,
 };
 use valida_memory::{MachineWithMemoryChip, MemoryChip};
@@ -260,6 +260,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
         let alpha: SC::Challenge = challenger.sample_ext_element();
 
         let mut quotients: Vec<RowMajorMatrix<SC::Val>> = vec![];
+        let mut public_inputs: Vec<Vec<(ColumnIndex, ColumnVector<SC::Val>)>> = vec![];
 
         let mut i: usize = 0;
 
@@ -284,6 +285,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(vec![]);
         i += 1;
 
         let chip = self.program();
@@ -300,13 +302,16 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             config,
             chip,
             log_degrees[i],
-            Some(preprocessed_trace_ldes.remove(0)),
+            None::<RowMajorMatrix<SC::Val>>,
             main_trace_ldes.remove(0),
             perm_trace_ldes.remove(0),
             cumulative_sums[i],
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(Chip::generate_public_inputs(
+            chip as &dyn Chip<Self, SC>,
+        ));
         i += 1;
 
         let chip = self.mem();
@@ -330,6 +335,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(vec![]);
         i += 1;
 
         let chip = self.add_u32();
@@ -353,6 +359,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(vec![]);
         i += 1;
 
         let chip = self.sub_u32();
@@ -376,6 +383,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(vec![]);
         i += 1;
 
         let chip = self.mul_u32();
@@ -399,6 +407,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(vec![]);
         i += 1;
 
         let chip = self.div_u32();
@@ -422,6 +431,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(vec![]);
         i += 1;
 
         let chip = self.shift_u32();
@@ -445,6 +455,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(vec![]);
         i += 1;
 
         let chip = self.lt_u32();
@@ -468,6 +479,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(vec![]);
         i += 1;
 
         let chip = self.com_u32();
@@ -491,6 +503,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(vec![]);
         i += 1;
 
         let chip = self.bitwise_u32();
@@ -514,6 +527,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(vec![]);
         i += 1;
 
         let chip = self.output();
@@ -537,6 +551,9 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(Chip::generate_public_inputs(
+            chip as &dyn Chip<Self, SC>,
+        ));
         i += 1;
 
         let chip = self.range();
@@ -560,6 +577,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(vec![]);
         i += 1;
 
         let chip = self.static_data();
@@ -583,6 +601,9 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             &perm_challenges,
             alpha,
         ));
+        public_inputs.push(Chip::generate_public_inputs(
+            chip as &dyn Chip<Self, SC>,
+        ));
         i += 1;
 
         assert_eq!(quotients.len(), NUM_CHIPS);
@@ -633,7 +654,8 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
             .zip(perm_openings)
             .zip(quotient_openings)
             .zip(perm_traces)
-            .map(|((((log_degree, main), perm), quotient), perm_trace)| {
+            .zip(public_inputs)
+            .map(|(((((log_degree, main), perm), quotient), perm_trace), public_inputs)| {
                 // TODO: add preprocessed openings
                 let [preprocessed_local, preprocessed_next] = [vec![], vec![]];
 
@@ -657,6 +679,7 @@ impl<F: PrimeField32 + TwoAdicField> Machine<F> for BasicMachine<F> {
                     .unwrap()
                     .clone();
                 ChipProof {
+                    public_inputs,
                     log_degree: *log_degree,
                     opened_values,
                     cumulative_sum,

basic_macro/src/lib.rs

@@ -40,7 +40,7 @@ use valida_cpu::{CpuChip, MachineWithCpuChip};
 use valida_derive::Machine;
 use valida_machine::{
     AdviceProvider, BusArgument, Chip, ChipProof, Instruction, Machine, MachineProof, ProgramROM,
-    StoppingFlag, ValidaAirBuilder,
+    StoppingFlag, ValidaAirBuilder, ColumnVector, ColumnIndex,
 };
 use valida_memory::{MachineWithMemoryChip, MemoryChip};
 use valida_output::{MachineWithOutputChip, OutputChip, WriteInstruction};

derive/src/lib.rs

@@ -268,6 +268,9 @@ fn prove_method(chips: &[&Field]) -> TokenStream2 {
                     &perm_challenges,
                     alpha,
                 ));
+                public_inputs.push(Chip::generate_public_inputs(
+                  self.#chip_name() as &dyn Chip<Self, SC>,
+                ));
             }
         })
         .collect::<TokenStream2>();
@@ -283,7 +286,7 @@ fn prove_method(chips: &[&Field]) -> TokenStream2 {
             use ::valida_machine::__internal::p3_commit::{Pcs, UnivariatePcs, UnivariatePcsWithLde};
             use ::valida_machine::__internal::p3_matrix::{Matrix, MatrixRowSlices, dense::RowMajorMatrix};
             use ::valida_machine::__internal::p3_util::log2_strict_usize;
-            use ::valida_machine::{generate_permutation_trace, MachineProof, ChipProof, Commitments};
+            use ::valida_machine::{generate_permutation_trace, Chip, MachineProof, ChipProof, Commitments};
             use ::valida_machine::OpenedValues;
             use alloc::vec;
             use alloc::vec::Vec;
@@ -360,6 +363,7 @@ fn prove_method(chips: &[&Field]) -> TokenStream2 {
             let alpha: SC::Challenge = challenger.sample_ext_element();
 
             let mut quotients: Vec<RowMajorMatrix<SC::Val>> = vec![];
+            let mut public_inputs: Vec<Vec<(ColumnIndex, ColumnVector<SC::Val>)>> = vec![];
             #compute_quotients
             assert_eq!(quotients.len(), #num_chips);
             assert_eq!(log_quotient_degrees.len(), #num_chips);
@@ -409,7 +413,8 @@ fn prove_method(chips: &[&Field]) -> TokenStream2 {
                 .zip(perm_openings)
                 .zip(quotient_openings)
                 .zip(perm_traces)
-                .map(|((((log_degree,  main), perm), quotient), perm_trace)| {
+                .zip(public_inputs)
+                .map(|(((((log_degree,  main), perm), quotient), perm_trace), public_inputs)| {
                     // TODO: add preprocessed openings
                     let [preprocessed_local, preprocessed_next] =
                         [vec![], vec![]];
@@ -430,6 +435,7 @@ fn prove_method(chips: &[&Field]) -> TokenStream2 {
 
                     let cumulative_sum = perm_trace.row_slice(perm_trace.height() - 1).last().unwrap().clone();
                     ChipProof {
+                        public_inputs,
                         log_degree: *log_degree,
                         opened_values,
                         cumulative_sum,

machine/src/chip.rs

@@ -1,5 +1,5 @@
 use crate::folding_builder::VerifierConstraintFolder;
-use crate::Machine;
+use crate::{Machine, ColumnVector, ColumnIndex};
 use crate::__internal::{DebugConstraintBuilder, ProverConstraintFolder};
 use alloc::vec;
 use alloc::vec::Vec;
@@ -21,6 +21,10 @@ pub trait Chip<M: Machine<SC::Val>, SC: StarkConfig>:
     /// Generate the main trace for the chip given the provided machine.
     fn generate_trace(&self, machine: &M) -> RowMajorMatrix<SC::Val>;
 
+    fn generate_public_inputs(&self) -> Vec<(ColumnIndex, ColumnVector<SC::Val>)> {
+        vec![]
+    }
+
     fn local_sends(&self) -> Vec<Interaction<SC::Val>> {
         vec![]
     }

machine/src/proof.rs

@@ -15,7 +15,7 @@ type PcsProof<SC> = <<SC as StarkConfig>::Pcs as Pcs<Val<SC>, ValMat<SC>>>::Proo
 pub struct MachineProof<SC: StarkConfig> {
     pub commitments: Commitments<Com<SC>>,
     pub opening_proof: PcsProof<SC>,
-    pub chip_proofs: Vec<ChipProof<SC::Challenge>>,
+    pub chip_proofs: Vec<ChipProof<SC>>,
 }
 
 #[derive(Serialize, Deserialize)]
@@ -26,10 +26,17 @@ pub struct Commitments<Com> {
 }
 
 #[derive(Serialize, Deserialize)]
-pub struct ChipProof<Challenge> {
+pub struct ColumnIndex(usize);
+
+#[derive(Serialize, Deserialize)]
+pub struct ColumnVector<A>(Vec<A>);
+
+#[derive(Serialize, Deserialize)]
+pub struct ChipProof<SC: StarkConfig> {
+    pub public_inputs: Vec<(ColumnIndex, ColumnVector<SC::Val>)>,
     pub log_degree: usize,
-    pub opened_values: OpenedValues<Challenge>,
-    pub cumulative_sum: Challenge,
+    pub opened_values: OpenedValues<SC::Challenge>,
+    pub cumulative_sum: SC::Challenge,
 }
 
 #[derive(Serialize, Deserialize)]

program/src/columns.rs

@@ -6,30 +6,16 @@ use valida_util::indices_arr;
 
 #[derive(AlignedBorrow, Default)]
 pub struct ProgramCols<T> {
-    pub multiplicity: T,
-}
-
-#[derive(AlignedBorrow, Default)]
-pub struct ProgramPreprocessedCols<T> {
     pub pc: T,
     pub opcode: T,
     pub operands: Operands<T>,
+    pub multiplicity: T,
 }
 
 pub const NUM_PROGRAM_COLS: usize = size_of::<ProgramCols<u8>>();
 pub const COL_MAP: ProgramCols<usize> = make_col_map();
 
-pub const NUM_PREPROCESSED_COLS: usize = size_of::<ProgramPreprocessedCols<u8>>();
-pub const PREPROCESSED_COL_MAP: ProgramPreprocessedCols<usize> = make_preprocessed_col_map();
-
 const fn make_col_map() -> ProgramCols<usize> {
     let indices_arr = indices_arr::<NUM_PROGRAM_COLS>();
     unsafe { transmute::<[usize; NUM_PROGRAM_COLS], ProgramCols<usize>>(indices_arr) }
 }
-
-const fn make_preprocessed_col_map() -> ProgramPreprocessedCols<usize> {
-    let indices_arr = indices_arr::<NUM_PREPROCESSED_COLS>();
-    unsafe {
-        transmute::<[usize; NUM_PREPROCESSED_COLS], ProgramPreprocessedCols<usize>>(indices_arr)
-    }
-}

program/src/lib.rs

@@ -11,7 +11,7 @@ use valida_util::pad_to_power_of_two;
 
 use p3_field::{AbstractField, Field};
 use p3_matrix::dense::RowMajorMatrix;
-use valida_machine::StarkConfig;
+use valida_machine::{InstructionWord, StarkConfig};
 
 pub mod columns;
 pub mod stark;
@@ -36,15 +36,33 @@ where
     SC: StarkConfig,
 {
     fn generate_trace(&self, _machine: &M) -> RowMajorMatrix<SC::Val> {
-        let mut values = self
+        // Pad the ROM to a power of two.
+        let mut rom = self.program_rom.0.clone();
+        let n = rom.len();
+        rom.resize(n.next_power_of_two(), InstructionWord::default());
+
+        let mut counts = self
             .counts
             .iter()
             .map(|c| SC::Val::from_canonical_u32(*c))
             .collect();
 
-        pad_to_power_of_two::<NUM_PROGRAM_COLS, SC::Val>(&mut values);
+        pad_to_power_of_two::<1, SC::Val>(&mut counts);
+
+        let flattened = rom
+            .into_iter()
+            .enumerate()
+            .zip(counts)
+            .flat_map(|((n, word), c)| {
+                let mut row = vec![SC::Val::zero(); NUM_PROGRAM_COLS];
+                row.push(SC::Val::from_canonical_usize(n));
+                row.append(&mut Vec::from(&word.flatten()));
+                row.push(c);
+                row
+            })
+            .collect();
 
-        RowMajorMatrix::new(values, NUM_PROGRAM_COLS)
+        RowMajorMatrix::new(flattened, NUM_PROGRAM_COLS)
     }
 
     fn global_receives(&self, _machine: &M) -> Vec<Interaction<SC::Val>> {

program/src/stark.rs

@@ -1,4 +1,4 @@
-use crate::columns::{NUM_PREPROCESSED_COLS, NUM_PROGRAM_COLS};
+use crate::columns::{NUM_PROGRAM_COLS};
 use crate::ProgramChip;
 use alloc::vec;
 use valida_machine::InstructionWord;
@@ -18,24 +18,4 @@ impl<F: Field> BaseAir<F> for ProgramChip {
     fn width(&self) -> usize {
         NUM_PROGRAM_COLS
     }
-
-    fn preprocessed_trace(&self) -> Option<RowMajorMatrix<F>> {
-        // Pad the ROM to a power of two.
-        let mut rom = self.program_rom.0.clone();
-        let n = rom.len();
-        rom.resize(n.next_power_of_two(), InstructionWord::default());
-
-        let flattened = rom
-            .into_iter()
-            .enumerate()
-            .flat_map(|(n, word)| {
-                let mut row = vec![F::zero(); NUM_PREPROCESSED_COLS];
-                row[0] = F::from_canonical_usize(n);
-                row[1..].copy_from_slice(&word.flatten());
-                row
-            })
-            .collect();
-        let trace = RowMajorMatrix::new(flattened, NUM_PREPROCESSED_COLS);
-        Some(trace)
-    }
 }