Plonky3 public inputs (#1465)

Implemented public inputs in plonky3-- public inputs are implemented as
selector witness columns in `precompute_trace` and constrained by
asserting `builder.assert_zero(local[selector_idx] * (public_value -
local[witness_col_idx])`. Implemented a pil test for public inputs, but
doesn't seem to work as witnesses are not filled in during
generate_trace_rows

---------

Co-authored-by: schaeff <[email protected]>
Co-authored-by: Thibaut Schaeffer <[email protected]>
Co-authored-by: Georg Wiese <[email protected]>

ast/src/analyzed/mod.rs

@@ -64,6 +64,10 @@ impl<T> Analyzed<T> {
     pub fn constant_count(&self) -> usize {
         self.declaration_type_count(PolynomialType::Constant)
     }
+    /// @returns the number of public inputs
+    pub fn publics_count(&self) -> usize {
+        self.public_declarations.len()
+    }
 
     pub fn constant_polys_in_source_order(
         &self,

plonky3/src/circuit_builder.rs

@@ -1,8 +1,16 @@
 //! A plonky3 adapter for powdr
+//!
+//! Support for public values invokes fixed selector columns, currently
+//! implemented as extra witness columns in the execution trace.
+//!
+//! Namely, given ith public value pub[i] corresponding to a witness value in
+//! row j of column Ci, a corresponding selector column Pi is constructed to
+//! constrain Pi * (pub[i] - Ci) on every row. Pi is precomputed in the trace
+//! to be 0 everywhere and 1 in row j.
 
-use std::any::TypeId;
+use std::{any::TypeId, collections::BTreeMap};
 
-use p3_air::{Air, AirBuilder, BaseAir};
+use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
 use p3_field::AbstractField;
 use p3_goldilocks::Goldilocks;
 use p3_matrix::{dense::RowMajorMatrix, MatrixRowSlices};
@@ -28,17 +36,21 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
     pub fn generate_trace_rows(&self) -> RowMajorMatrix<Goldilocks> {
         // an iterator over all columns, committed then fixed
         let witness = self.witness().iter();
+        let publics = self.get_publics().into_iter();
         let len = self.analyzed.degree.unwrap();
 
         // for each row, get the value of each column
         let values = (0..len)
             .flat_map(move |i| {
-                witness
-                    .clone()
-                    .map(move |(_, v)| cast_to_goldilocks(v[i as usize]))
+                // witness values
+                witness.clone().map(move |(_, v)| v[i as usize]).chain(
+                    publics
+                        .clone()
+                        .map(move |(_, _, idx)| T::from(i as usize == idx)),
+                )
             })
+            .map(cast_to_goldilocks)
             .collect();
-
         RowMajorMatrix::new(values, self.width())
     }
 }
@@ -53,9 +65,6 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
         if analyzed.constant_count() > 0 {
             unimplemented!("Fixed columns are not supported in Plonky3");
         }
-        if !analyzed.public_declarations.is_empty() {
-            unimplemented!("Public declarations are not supported in Plonky3");
-        }
         if analyzed
             .definitions
             .iter()
@@ -75,6 +84,52 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
         self.witness.as_ref().unwrap()
     }
 
+    /// Retrieves (col_name, col_idx, offset) of each public witness in the trace.
+    pub(crate) fn get_publics(&self) -> Vec<(String, usize, usize)> {
+        let mut publics = self
+            .analyzed
+            .public_declarations
+            .values()
+            .map(|public_declaration| {
+                let witness_name = public_declaration.referenced_poly_name();
+                let witness_column = {
+                    let base = public_declaration.polynomial.poly_id.unwrap().id as usize;
+                    match public_declaration.array_index {
+                        Some(array_idx) => base + array_idx,
+                        None => base,
+                    }
+                };
+                let witness_offset = public_declaration.index as usize;
+                (witness_name, witness_column, witness_offset)
+            })
+            .collect::<Vec<_>>();
+
+        // Sort, so that the order is deterministic
+        publics.sort();
+        publics
+    }
+
+    /// Calculates public values from generated witness values.
+    pub(crate) fn get_public_values(&self) -> Vec<Goldilocks> {
+        let publics = self.get_publics();
+
+        let witness = self
+            .witness
+            .as_ref()
+            .expect("Witness needs to be set")
+            .iter()
+            .map(|(name, values)| (name, values))
+            .collect::<BTreeMap<_, _>>();
+
+        publics
+            .into_iter()
+            .map(|(col_name, _, idx)| {
+                let vals = *witness.get(&col_name).unwrap();
+                cast_to_goldilocks(vals[idx])
+            })
+            .collect()
+    }
+
     pub(crate) fn with_witness(self, witness: &'a [(String, Vec<T>)]) -> Self {
         assert_eq!(witness.len(), self.analyzed.commitment_count());
         Self {
@@ -89,7 +144,6 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
             ..self
         }
     }
-
     /// Conversion to plonky3 expression
     fn to_plonky3_expr<AB: AirBuilder<F = Val>>(
         &self,
@@ -161,18 +215,35 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
 impl<'a, T: FieldElement> BaseAir<Val> for PowdrCircuit<'a, T> {
     fn width(&self) -> usize {
         assert_eq!(self.analyzed.constant_count(), 0);
-        self.analyzed.commitment_count()
+        self.analyzed.commitment_count() + self.analyzed.publics_count()
     }
 
     fn preprocessed_trace(&self) -> Option<RowMajorMatrix<Val>> {
         unimplemented!()
     }
 }
 
-impl<'a, T: FieldElement, AB: AirBuilder<F = Val>> Air<AB> for PowdrCircuit<'a, T> {
+impl<'a, T: FieldElement, AB: AirBuilderWithPublicValues<F = Val>> Air<AB> for PowdrCircuit<'a, T> {
     fn eval(&self, builder: &mut AB) {
         let matrix = builder.main();
+        let pi = builder.public_values();
+        let publics = self.get_publics();
+        assert_eq!(publics.len(), pi.len());
+
+        // public constraints
+        let pi_moved = pi.to_vec();
+        let local = matrix.row_slice(0);
+
+        // constraining Pi * (Ci - pub[i]) = 0
+        publics.iter().zip(pi_moved).enumerate().for_each(
+            |(index, ((_, col_id, _), public_value))| {
+                let selector = local[self.analyzed.commitment_count() + index];
+                let witness_col = local[*col_id];
+                builder.assert_zero(selector * (public_value.into() - witness_col));
+            },
+        );
 
+        // circuit constraints
         for identity in &self
             .analyzed
             .identities_with_inlined_intermediate_polynomials()

plonky3/src/prover/mod.rs

@@ -37,7 +37,7 @@ impl<'a, T: FieldElement> Plonky3Prover<'a, T> {
             .with_witgen_callback(witgen_callback)
             .with_witness(witness);
 
-        let publics = vec![];
+        let publics = circuit.get_public_values();
 
         let trace = circuit.generate_trace_rows();
 
@@ -86,24 +86,50 @@ mod tests {
 
     /// Prove and verify execution
     fn run_test_goldilocks(pil: &str) {
+        run_test_goldilocks_publics(pil, None)
+    }
+
+    fn run_test_goldilocks_publics(pil: &str, malicious_publics: Option<Vec<GoldilocksField>>) {
         let mut pipeline = Pipeline::<GoldilocksField>::default().from_pil_string(pil.to_string());
 
         let pil = pipeline.compute_optimized_pil().unwrap();
         let witness_callback = pipeline.witgen_callback().unwrap();
         let witness = pipeline.compute_witness().unwrap();
 
-        let proof = Plonky3Prover::new(&pil).prove(&witness, witness_callback);
+        let prover = Plonky3Prover::new(&pil);
+        let proof = prover.prove(&witness, witness_callback);
 
         assert!(proof.is_ok());
+
+        if let Some(publics) = malicious_publics {
+            prover.verify(&proof.unwrap(), &[publics]).unwrap()
+        }
     }
 
     #[test]
-    #[should_panic = "not implemented"]
     fn publics() {
         let content = "namespace Global(8); pol witness x; x * (x - 1) = 0; public out = x(7);";
         run_test_goldilocks(content);
     }
 
+    #[test]
+    #[should_panic = r#"called `Result::unwrap()` on an `Err` value: "Failed to verify proof: OodEvaluationMismatch""#]
+    fn public_inputs_malicious() {
+        let content = r#"
+        namespace Add(8);
+            col witness x;
+            col witness y;
+            col witness z;
+            y - 1 = 0;
+            x = 0;
+            x + y = z;
+
+            public outz = z(7);
+        "#;
+        let malicious_publics = Some(vec![GoldilocksField::from(0)]);
+        run_test_goldilocks_publics(content, malicious_publics);
+    }
+
     #[test]
     #[should_panic = "assertion failed: width >= 1"]
     fn empty() {