test(public_table): test public tables behind optional cfg `cfg(test_…

…public_table)` (#1412)

Mainly to test #1411 with a
dummy example

---------

Co-authored-by: Matthias Görgens <matthias.goergens@gmail.com>
Co-authored-by: Matthias Görgens <matthias.goergens@mozak.com>
Co-authored-by: bing <bingcicle@keemail.me>
Co-authored-by: bing <bingcicle@mozak.com>

.github/workflows/ci.yml

@@ -47,7 +47,7 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Compile and archive all the tests
-        run: nixdo nice cargo nextest archive --cargo-profile="${{ matrix.profile }}" --locked --features="enable_poseidon_starks" --all-targets --archive-file mozak-vm-tests.tar.zst
+        run: nixdo nice cargo nextest archive --cargo-profile="${{ matrix.profile }}" --locked --features="enable_poseidon_starks, test_public_table" --all-targets --archive-file mozak-vm-tests.tar.zst
 
       - name: Run all the tests from the archive
         run: nixdo MOZAK_STARK_DEBUG=true nice cargo nextest run --no-fail-fast --archive-file mozak-vm-tests.tar.zst

circuits/Cargo.toml

@@ -42,6 +42,7 @@ proptest = "1.4"
 enable_poseidon_starks = []
 parallel = ["plonky2_maybe_rayon/parallel"]
 test = []
+test_public_table = []
 timing = ["plonky2/timing", "starky/timing"]
 
 [[test]]

circuits/src/bitshift/columns.rs

@@ -1,5 +1,7 @@
 use crate::columns_view::{columns_view_impl, make_col_map};
 use crate::linear_combination::Column;
+#[cfg(feature = "test_public_table")]
+use crate::public_sub_table::PublicSubTable;
 use crate::stark::mozak_stark::{BitshiftTable, TableWithTypedOutput};
 
 columns_view_impl!(Bitshift);
@@ -40,3 +42,21 @@ pub struct BitshiftView<T> {
 pub fn lookup_for_cpu() -> TableWithTypedOutput<Bitshift<Column>> {
     BitshiftTable::new(COL_MAP.executed, COL_MAP.multiplicity)
 }
+
+/// This function makes `amount, multiplier` columns of
+/// Bitshift table, public. It looks something like this:
+/// [0, 1]
+/// [1, 2]
+/// [2, 4]
+/// ....
+/// [31, 2^31]
+#[must_use]
+#[cfg(feature = "test_public_table")]
+pub fn public_sub_table() -> PublicSubTable {
+    use crate::linear_combination_typed::ColumnWithTypedInput;
+
+    PublicSubTable {
+        table: BitshiftTable::new(COL_MAP.executed, ColumnWithTypedInput::constant(1)),
+        num_rows: 32,
+    }
+}

circuits/src/stark/mozak_stark.rs

@@ -140,6 +140,9 @@ pub struct MozakStark<F: RichField + Extendable<D>, const D: usize> {
     )]
     pub poseidon2_output_bytes_stark: Poseidon2OutputBytesStark<F, D>,
     pub cross_table_lookups: [CrossTableLookup; NUM_CROSS_TABLE_LOOKUP],
+    #[cfg(feature = "test_public_table")]
+    pub public_sub_tables: [PublicSubTable; 1],
+    #[cfg(not(feature = "test_public_table"))]
     pub public_sub_tables: [PublicSubTable; 0],
     pub debug: bool,
 }
@@ -411,7 +414,10 @@ impl<F: RichField + Extendable<D>, const D: usize> Default for MozakStark<F, D>
                 #[cfg(feature = "enable_poseidon_starks")]
                 Poseidon2OutputBytesPoseidon2SpongeTable::lookups(),
             ],
+            #[cfg(not(feature = "test_public_table"))]
             public_sub_tables: [],
+            #[cfg(feature = "test_public_table")]
+            public_sub_tables: [crate::bitshift::columns::public_sub_table()],
             debug: false,
         }
     }

circuits/src/stark/recursive_verifier.rs

@@ -49,7 +49,11 @@ pub const VM_RECURSION_THRESHOLD_DEGREE_BITS: usize = 12;
 ///                          hash) = 64
 ///   `ElfMemoryInit trace cap`: 64
 ///   `MozakMemoryInit trace cap`: 64
+///   `bitshift public sub table`: 2 rows * 32 columns = 64
+#[cfg(not(feature = "test_public_table"))]
 pub const VM_PUBLIC_INPUT_SIZE: usize = 1 + 64 + 64 + 64;
+#[cfg(feature = "test_public_table")]
+pub const VM_PUBLIC_INPUT_SIZE: usize = 1 + 64 + 64 + 64 + 64;
 pub const VM_RECURSION_CONFIG: CircuitConfig = CircuitConfig::standard_recursion_config();
 
 /// Represents a circuit which recursively verifies STARK proofs.
@@ -652,15 +656,19 @@ mod tests {
         verify_recursive_vm_proof, VM_PUBLIC_INPUT_SIZE, VM_RECURSION_CONFIG,
         VM_RECURSION_THRESHOLD_DEGREE_BITS,
     };
-    use crate::stark::verifier::verify_proof;
     use crate::test_utils::{C, D, F};
     use crate::utils::from_u32;
 
     type S = MozakStark<F, D>;
 
     #[test]
-    #[ignore]
+    #[cfg(feature = "test_public_table")]
     fn recursive_verify_mozak_starks() -> Result<()> {
+        use itertools::Itertools;
+        use plonky2::field::types::Field;
+
+        use crate::stark::verifier::verify_proof;
+
         let stark = S::default();
         let mut config = StarkConfig::standard_fast_config();
         config.fri_config.cap_height = 1;
@@ -700,6 +708,15 @@ mod tests {
         );
 
         let recursive_proof = mozak_stark_circuit.prove(&mozak_proof)?;
+        let expected_bitshift_subtable = (0..32)
+            .flat_map(|i| [F::from_canonical_u64(i), F::from_canonical_u64(1 << i)])
+            .collect_vec();
+        assert_eq!(
+            recursive_proof.public_inputs[25..].to_vec(),
+            expected_bitshift_subtable,
+            "Could not find bitshift subtable in recursive proof's public inputs"
+        );
+
         mozak_stark_circuit.circuit.verify(recursive_proof)
     }