Update public examples

backend/src/composite/sub_prover.rs

@@ -74,6 +74,9 @@ pub struct SubProver<'s, F: FieldElement> {
 
 impl<'s, F: FieldElement> SubProver<'s, F> {
     pub fn resume(self, response: Vec<(String, Vec<F>)>) -> RunStatus<'s, F> {
+        // `response_sender` has to send `(witness, publics)` to match `response_receiver`.
+        // TODO: currently publics is set to empty, but we should update this once composite backend
+        // uses public reference values instead of witness values in the future.
         self.response_sender
             .send((response, BTreeMap::new()))
             .unwrap();

backend/src/mock/machine.rs

@@ -21,6 +21,7 @@ impl<'a, F: FieldElement> Machine<'a, F> {
     pub fn try_new(
         machine_name: String,
         witness: &'a [(String, Vec<F>)],
+        publics: &'a BTreeMap<String, Option<F>>,
         fixed: &'a [(String, VariablySizedColumn<F>)],
         pil: &'a Analyzed<F>,
         witgen_callback: &WitgenCallback<F>,
@@ -55,9 +56,8 @@ impl<'a, F: FieldElement> Machine<'a, F> {
 
         let intermediate_definitions = pil.intermediate_definitions();
 
-        // TODO: Supports publics.
-        let values =
-            OwnedTerminalValues::new(pil, witness, fixed).with_challenges(challenges.clone());
+        let values = OwnedTerminalValues::new(pil, witness, publics.clone(), fixed)
+            .with_challenges(challenges.clone());
 
         Some(Self {
             machine_name,

backend/src/mock/mod.rs

@@ -76,7 +76,7 @@ impl<F: FieldElement> Backend<F> for MockBackend<F> {
     fn prove(
         &self,
         witness: &[(String, Vec<F>)],
-        _publics: &BTreeMap<String, Option<F>>,
+        publics: &BTreeMap<String, Option<F>>,
         prev_proof: Option<Proof>,
         witgen_callback: WitgenCallback<F>,
     ) -> Result<Proof, Error> {
@@ -111,6 +111,7 @@ impl<F: FieldElement> Backend<F> for MockBackend<F> {
                 Machine::try_new(
                     machine_name.clone(),
                     witness,
+                    publics,
                     &self.fixed,
                     pil,
                     &witgen_callback,

backend/src/plonky3/stark.rs

@@ -375,12 +375,22 @@ mod tests {
             pol fixed FIRST = [1] + [0]*;
             pol witness x;
             pol witness y;
-            y * y = y;
+            y = 1;
             x' = (1 - FIRST') * (x + 1);
             public out0 = x(6);
             public out1 = x(7);
             public out2 = y(3);
             public out3 = y(5);
+
+            pol fixed IS_6 = [0, 0, 0, 0, 0, 0, 1, 0];
+            pol fixed IS_7 = [0, 0, 0, 0, 0, 0, 0, 1];
+            pol fixed IS_3 = [0, 0, 0, 1, 0, 0, 0, 0];
+            pol fixed IS_5 = [0, 0, 0, 0, 0, 1, 0, 0];
+
+            IS_6 * (x - out0) = 0;
+            IS_7 * (x - out1) = 0;
+            IS_3 * (y - out2) = 0;
+            IS_5 * (y - out3) = 0;
         ";
         run_test(content);
     }
@@ -412,6 +422,9 @@ mod tests {
             x + y = z;
 
             public outz = z(7);
+
+            col fixed IS_LAST = [0]* + [1];
+            IS_LAST * (z - outz) = 0;
         "#;
         let malicious_publics = Some(vec![0]);
         run_test_publics(content, &malicious_publics);
@@ -486,6 +499,9 @@ mod tests {
             x = y + beta * alpha;
 
             public out = y(N - 1);
+
+            col fixed IS_LAST = [0]* + [1];
+            IS_LAST * (y - out) = 0;
         "#;
         run_test(content);
     }

executor-utils/src/expression_evaluator.rs

@@ -37,10 +37,11 @@ pub struct RowValues<'a, F> {
     row: usize,
 }
 
-impl<F: std::fmt::Debug> OwnedTerminalValues<F> {
+impl<F: std::fmt::Debug + Clone> OwnedTerminalValues<F> {
     pub fn new(
         pil: &Analyzed<F>,
         witness_columns: Vec<(String, Vec<F>)>,
+        publics: BTreeMap<String, Option<F>>,
         fixed_columns: Vec<(String, Vec<F>)>,
     ) -> Self {
         let mut columns_by_name = witness_columns
@@ -57,9 +58,13 @@ impl<F: std::fmt::Debug> OwnedTerminalValues<F> {
                     .map(|column| (poly_id, column))
             })
             .collect();
+        let publics = publics
+            .iter()
+            .map(|(name, value)| (name.clone(), value.clone().unwrap()))
+            .collect();
         Self {
             trace,
-            public_values: Default::default(),
+            public_values: publics,
             challenge_values: Default::default(),
         }
     }
@@ -123,7 +128,7 @@ impl<F: FieldElement, T: From<F>> TerminalAccess<T> for RowValues<'_, F> {
 
 /// Evaluates an algebraic expression to a value.
 pub struct ExpressionEvaluator<'a, T, Expr, TA> {
-    terminal_access: TA,
+    pub terminal_access: TA,
     intermediate_definitions: &'a BTreeMap<AlgebraicReferenceThin, Expression<T>>,
     /// Maps intermediate reference to their evaluation. Updated throughout the lifetime of the
     /// ExpressionEvaluator.

executor/src/witgen/bus_accumulator/mod.rs

@@ -108,6 +108,7 @@ impl<'a, T: FieldElement, Ext: ExtensionField<T> + Sync> BusAccumulatorGenerator
         let values = OwnedTerminalValues::new(
             pil,
             witness_columns.to_vec(),
+            BTreeMap::new(),
             fixed_columns
                 .map(|(name, values)| (name, values.to_vec()))
                 .collect(),

executor/src/witgen/machines/write_once_memory.rs

@@ -3,7 +3,8 @@ use std::collections::{BTreeMap, HashMap};
 use itertools::{Either, Itertools};
 
 use num_traits::One;
-use powdr_ast::analyzed::{PolyID, PolynomialType};
+use powdr_ast::analyzed::{AlgebraicExpression, PolyID, PolynomialType};
+use powdr_ast::parsed::visitor::{AllChildren, Children};
 use powdr_number::{DegreeType, FieldElement};
 
 use crate::witgen::data_structures::identity::BusReceive;
@@ -42,6 +43,7 @@ pub struct WriteOnceMemory<'a, T: FieldElement> {
     /// The memory content
     data: BTreeMap<DegreeType, Vec<Option<T>>>,
     name: String,
+    public_names: Vec<String>,
 }
 
 impl<'a, T: FieldElement> WriteOnceMemory<'a, T> {
@@ -50,7 +52,13 @@ impl<'a, T: FieldElement> WriteOnceMemory<'a, T> {
         fixed_data: &'a FixedData<'a, T>,
         parts: &MachineParts<'a, T>,
     ) -> Option<Self> {
-        if !parts.identities.is_empty() {
+        // All identities should have a public reference
+        if !parts.identities.iter().all(|id| {
+            (*id).children().any(|c| {
+                c.all_children()
+                    .any(|c| matches!(c, AlgebraicExpression::PublicReference(_)))
+            })
+        }) {
             return None;
         }
 
@@ -135,6 +143,12 @@ impl<'a, T: FieldElement> WriteOnceMemory<'a, T> {
             value_polys,
             key_to_index,
             data: BTreeMap::new(),
+            public_names: parts
+                .fixed_data
+                .analyzed
+                .public_declarations_in_source_order()
+                .map(|(name, _)| name.clone())
+                .collect(),
         })
     }
 
@@ -263,7 +277,8 @@ impl<'a, T: FieldElement> Machine<'a, T> for WriteOnceMemory<'a, T> {
         &mut self,
         _mutable_state: &'b MutableState<'a, T, Q>,
     ) -> HashMap<String, Vec<T>> {
-        self.value_polys
+        let witness = self
+            .value_polys
             .iter()
             .enumerate()
             .map(|(value_index, poly)| {
@@ -285,6 +300,42 @@ impl<'a, T: FieldElement> Machine<'a, T> for WriteOnceMemory<'a, T> {
                 (*poly, column)
             })
             .map(|(poly_id, column)| (self.fixed_data.column_name(&poly_id).to_string(), column))
-            .collect()
+            .collect();
+        witness
+    }
+
+    fn take_public_values(&mut self) -> BTreeMap<String, T> {
+        if self.public_names.is_empty() {
+            BTreeMap::new()
+        } else {
+            let public_values: Vec<_> = self
+                .value_polys
+                .iter()
+                .enumerate()
+                .flat_map(|(value_index, poly)| {
+                    self.fixed_data.witness_cols[poly]
+                        .external_values
+                        .cloned()
+                        .map(|mut external_values| {
+                            // External witness values might only be provided partially.
+                            external_values.resize(self.degree as usize, T::zero());
+                            external_values
+                        })
+                        .unwrap_or_else(|| {
+                            let mut column = vec![T::zero(); 8]; // 8 public values
+                            for (row, values) in self.data.iter() {
+                                column[*row as usize] = values[value_index].unwrap_or_default();
+                            }
+                            column
+                        })
+                })
+                .collect();
+
+            self.public_names
+                .clone()
+                .into_iter()
+                .zip(public_values)
+                .collect()
+        }
     }
 }

executor/src/witgen/mod.rs

@@ -247,10 +247,11 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
         let machines = MachineExtractor::new(&fixed).split_out_machines();
 
         // Run main machine and extract columns from all machines.
-        let (columns, _publics) =
+        let (columns, publics) =
             MutableState::new(machines.into_iter(), &self.query_callback).run();
 
-        let publics = extract_publics(&columns, self.analyzed);
+        let publics = extract_publics(publics, self.analyzed);
+
         if !publics.is_empty() {
             log::debug!("Publics:");
         }
@@ -291,22 +292,16 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
     }
 }
 
-pub fn extract_publics<'a, T, I>(witness: I, pil: &Analyzed<T>) -> BTreeMap<String, Option<T>>
+pub fn extract_publics<T>(
+    public_references: BTreeMap<String, T>,
+    pil: &Analyzed<T>,
+) -> BTreeMap<String, Option<T>>
 where
     T: FieldElement,
-    I: IntoIterator<Item = (&'a String, &'a Vec<T>)>,
 {
-    let witness = witness
-        .into_iter()
-        .map(|(name, col)| (name.clone(), col))
-        .collect::<BTreeMap<_, _>>();
     pil.public_declarations_in_source_order()
-        .map(|(name, public_declaration)| {
-            let poly_name = &public_declaration.referenced_poly_name();
-            let poly_row = public_declaration.row();
-            let value = witness
-                .get(poly_name)
-                .map(|column| column[poly_row as usize]);
+        .map(|(name, _)| {
+            let value = public_references.get(name).cloned();
             ((*name).clone(), value)
         })
         .collect()

parser/src/lib.rs

@@ -306,7 +306,8 @@ namespace Fibonacci(N);
     [x + 2, y'] in [ISLAST, 7];
     y $ [x + 2, y'] is ISLAST $ [ISLAST, 7];
     (x - 2) * y = 8;
-    public out = y(%last_row);"#;
+    public out = y(%last_row);
+    ISLAST * (y - out) = 0;"#;
         let printed = format!("{}", parse(Some("input"), input).unwrap());
         assert_eq!(input.trim(), printed.trim());
     }

pil-analyzer/tests/parse_display.rs

@@ -31,6 +31,8 @@ namespace std::prover(65536);
 namespace std::convert(65536);
     let int = [];
 namespace T(65536);
+    col fixed IS_2 = [0_fe, 0_fe, 1_fe] + [0_fe]*;
+    T::IS_2 * (P - T::pc) = 0;
     col fixed first_step = [1_fe] + [0_fe]*;
     col fixed line(i) { i };
     let ops: int -> bool = |i| i < 7_int && 6_int >= -i;

pipeline/src/pipeline.rs

@@ -541,16 +541,19 @@ impl<T: FieldElement> Pipeline<T> {
     }
 
     /// Sets the witness to the provided value.
-    pub fn set_witness(mut self, witness: Vec<(String, Vec<T>)>) -> Self {
+    pub fn set_witness_and_publics(
+        mut self,
+        witness: Vec<(String, Vec<T>)>,
+        publics: BTreeMap<String, Option<T>>,
+    ) -> Self {
         if self.output_dir.is_some() {
             // Some future steps require the witness to be persisted.
             let fixed_cols = self.compute_fixed_cols().unwrap();
             self.maybe_write_witness(&fixed_cols, &witness).unwrap();
         }
         Pipeline {
             artifact: Artifacts {
-                // need to set publics to Some, or `compute_witness` will run auto witgen
-                witness_and_publics: Some(Arc::new((witness, BTreeMap::new()))),
+                witness_and_publics: Some(Arc::new((witness, publics))),
                 // we're changing the witness, clear the current proof
                 proof: None,
                 ..self.artifact
@@ -708,7 +711,12 @@ impl<T: FieldElement> Pipeline<T> {
     // but give it different witnesses and generate different proofs.
     // The previous alternative to this was cloning the entire pipeline.
     pub fn rollback_from_witness(&mut self) {
-        self.artifact.witness_and_publics = None;
+        if let Some(old_arc) = &self.artifact.witness_and_publics {
+            // Clone the public part
+            let public = old_arc.1.clone();
+            // Replace with an empty witness and the preserved public
+            self.artifact.witness_and_publics = Some(Arc::new((Vec::new(), public)));
+        }
         self.artifact.proof = None;
         self.arguments.external_witness_values.clear();
     }
@@ -1047,7 +1055,10 @@ impl<T: FieldElement> Pipeline<T> {
 
     pub fn compute_witness(&mut self) -> WitgenResult<T> {
         if let Some(arc) = &self.artifact.witness_and_publics {
-            return Ok(arc.clone());
+            // If exists, witness will always be non-empty, but not necessarily publics
+            if !arc.0.is_empty() {
+                return Ok(arc.clone());
+            }
         }
 
         self.host_context.clear();
@@ -1079,8 +1090,16 @@ impl<T: FieldElement> Pipeline<T> {
             .all(|name| external_witness_values.iter().any(|(e, _)| e == name))
         {
             self.log("All witness columns externally provided, skipping witness generation.");
-            self.artifact.witness_and_publics =
-                Some(Arc::new((external_witness_values, BTreeMap::new())));
+            if let Some(old_arc) = &self.artifact.witness_and_publics {
+                // Clone the public part
+                let public = old_arc.1.clone();
+                // Replace with an empty witness and the preserved public
+                self.artifact.witness_and_publics =
+                    Some(Arc::new((external_witness_values, public)));
+            } else {
+                self.artifact.witness_and_publics =
+                    Some(Arc::new((external_witness_values, BTreeMap::new())));
+            }
         } else {
             self.log("Deducing witness columns...");
             let start = Instant::now();