Stage 0 public reference

backend-utils/src/lib.rs

@@ -153,7 +153,7 @@ pub fn referenced_namespaces_algebraic_expression<F: FieldElement>(
         .all_children()
         .filter_map(|expr| match expr {
             AlgebraicExpression::Reference(reference) => Some(extract_namespace(&reference.name)),
-            AlgebraicExpression::PublicReference(_) => unimplemented!(),
+            AlgebraicExpression::PublicReference(name) => Some(extract_namespace(name)),
             AlgebraicExpression::Challenge(_)
             | AlgebraicExpression::Number(_)
             | AlgebraicExpression::BinaryOperation(_)

backend/src/composite/mod.rs

@@ -304,6 +304,7 @@ impl<F: FieldElement> Backend<F> for CompositeBackend<F> {
     fn prove(
         &self,
         witness: &[(String, Vec<F>)],
+        publics: &BTreeMap<String, Option<F>>,
         prev_proof: Option<Proof>,
         witgen_callback: WitgenCallback<F>,
     ) -> Result<Proof, Error> {
@@ -339,7 +340,12 @@ impl<F: FieldElement> Backend<F> for CompositeBackend<F> {
                         .expect("Machine does not support the given size");
 
                     let status = time_stage(machine, size, 0, || {
-                        sub_prover::run(scope, &inner_machine_data.backend, witness)
+                        sub_prover::run(
+                            scope,
+                            &inner_machine_data.backend,
+                            witness,
+                            publics.clone(),
+                        )
                     });
 
                     Some((status, machine_entry, size))

backend/src/composite/sub_prover.rs

@@ -17,6 +17,7 @@ pub fn run<'s, 'env, F: FieldElement>(
     scope: &'s Scope<'s, 'env>,
     prover: &'env Mutex<Box<dyn Backend<F>>>,
     witness: Vec<(String, Vec<F>)>,
+    publics: BTreeMap<String, Option<F>>,
 ) -> RunStatus<'s, F>
 where
 {
@@ -44,7 +45,10 @@ where
         // proof, even if it's not needed anymore. We should probably change
         // this API so the Vec is moved into the prover, and returned in the
         // callback and result.
-        prover.lock().unwrap().prove(&witness, None, callback)
+        prover
+            .lock()
+            .unwrap()
+            .prove(&witness, &publics, None, callback)
     });
 
     SubProver {

backend/src/estark/mod.rs

@@ -5,6 +5,7 @@ pub mod polygon_wrapper;
 pub mod starky_wrapper;
 
 use std::{
+    collections::BTreeMap,
     fs::File,
     io::{self, BufWriter, Write},
     iter::{once, repeat},
@@ -261,6 +262,7 @@ impl<F: FieldElement> Backend<F> for DumpBackend<F> {
     fn prove(
         &self,
         witness: &[(String, Vec<F>)],
+        _publics: &BTreeMap<String, Option<F>>,
         prev_proof: Option<Proof>,
         // TODO: Implement challenges
         _witgen_callback: WitgenCallback<F>,

backend/src/estark/polygon_wrapper.rs

@@ -1,4 +1,4 @@
-use std::{fs, path::PathBuf, sync::Arc};
+use std::{collections::BTreeMap, fs, path::PathBuf, sync::Arc};
 
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::{
@@ -49,6 +49,7 @@ impl<F: FieldElement> Backend<F> for PolygonBackend<F> {
     fn prove(
         &self,
         witness: &[(String, Vec<F>)],
+        _publics: &BTreeMap<String, Option<F>>,
         prev_proof: Option<Proof>,
         // TODO: Implement challenges
         _witgen_callback: WitgenCallback<F>,

backend/src/estark/starky_wrapper.rs

@@ -1,3 +1,4 @@
+use std::collections::BTreeMap;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::Instant;
@@ -192,6 +193,7 @@ impl Backend<GoldilocksField> for EStark {
     fn prove(
         &self,
         witness: &[(String, Vec<GoldilocksField>)],
+        _publics: &BTreeMap<String, Option<GoldilocksField>>,
         prev_proof: Option<crate::Proof>,
         // TODO: Implement challenges
         _witgen_callback: WitgenCallback<GoldilocksField>,

backend/src/halo2/mod.rs

@@ -1,3 +1,4 @@
+use std::collections::BTreeMap;
 use std::io;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -144,6 +145,7 @@ impl Backend<Bn254Field> for Halo2Prover {
     fn prove(
         &self,
         witness: &[(String, Vec<Bn254Field>)],
+        _publics: &BTreeMap<String, Option<Bn254Field>>,
         prev_proof: Option<Proof>,
         witgen_callback: WitgenCallback<Bn254Field>,
     ) -> Result<Proof, Error> {
@@ -231,6 +233,7 @@ impl<T: FieldElement> Backend<T> for Halo2Mock<T> {
     fn prove(
         &self,
         witness: &[(String, Vec<T>)],
+        _publics: &BTreeMap<String, Option<T>>,
         prev_proof: Option<Proof>,
         witgen_callback: WitgenCallback<T>,
     ) -> Result<Proof, Error> {

backend/src/lib.rs

@@ -14,7 +14,7 @@ mod mock;
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::{constant_evaluator::VariablySizedColumn, witgen::WitgenCallback};
 use powdr_number::{DegreeType, FieldElement};
-use std::{io, path::PathBuf, sync::Arc};
+use std::{collections::BTreeMap, io, path::PathBuf, sync::Arc};
 use strum::{Display, EnumString, EnumVariantNames};
 
 #[derive(Clone, EnumString, EnumVariantNames, Display, Copy)]
@@ -197,6 +197,7 @@ pub trait Backend<F: FieldElement>: Send {
     fn prove(
         &self,
         witness: &[(String, Vec<F>)],
+        publics: &BTreeMap<String, Option<F>>,
         prev_proof: Option<Proof>,
         witgen_callback: WitgenCallback<F>,
     ) -> Result<Proof, Error>;

backend/src/mock/mod.rs

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

backend/src/plonky3/mod.rs

@@ -1,6 +1,6 @@
 mod stark;
 
-use std::{io, path::PathBuf, sync::Arc};
+use std::{collections::BTreeMap, io, path::PathBuf, sync::Arc};
 
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::{constant_evaluator::VariablySizedColumn, witgen::WitgenCallback};
@@ -71,14 +71,15 @@ where
     fn prove(
         &self,
         witness: &[(String, Vec<T>)],
+        publics: &BTreeMap<String, Option<T>>,
         prev_proof: Option<Proof>,
         witgen_callback: WitgenCallback<T>,
     ) -> Result<Proof, Error> {
         if prev_proof.is_some() {
             return Err(Error::NoAggregationAvailable);
         }
 
-        Ok(self.prove(witness, witgen_callback)?)
+        Ok(self.prove(witness, publics, witgen_callback)?)
     }
 
     fn export_verification_key(&self, output: &mut dyn io::Write) -> Result<(), Error> {

backend/src/plonky3/stark.rs

@@ -206,6 +206,7 @@ where
     pub fn prove(
         &self,
         witness: &[(String, Vec<T>)],
+        _publics: &BTreeMap<String, Option<T>>,
         witgen_callback: WitgenCallback<T>,
     ) -> Result<Vec<u8>, String> {
         let mut witness_by_machine = self

backend/src/stwo/mod.rs

@@ -1,5 +1,6 @@
 use serde::de::DeserializeOwned;
 use serde::Serialize;
+use std::collections::BTreeMap;
 use std::io;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -79,6 +80,7 @@ where
     fn prove(
         &self,
         witness: &[(String, Vec<M31>)],
+        _publics: &BTreeMap<String, Option<M31>>,
         prev_proof: Option<Proof>,
         witgen_callback: WitgenCallback<M31>,
     ) -> Result<Proof, Error> {

executor/src/witgen/data_structures/mutable_state.rs

@@ -39,11 +39,11 @@ impl<'a, T: FieldElement, Q: QueryCallback<T>> MutableState<'a, T, Q> {
 
     /// Runs the first machine (unless there are no machines) end returns the generated columns.
     /// The first machine might call other machines, which is handled automatically.
-    pub fn run(self) -> HashMap<String, Vec<T>> {
+    pub fn run(self) -> (HashMap<String, Vec<T>>, BTreeMap<String, T>) {
         if let Some(first_machine) = self.machines.first() {
             first_machine.try_borrow_mut().unwrap().run_timed(&self);
         }
-        self.take_witness_col_values()
+        self.take_witness_col_and_public_values()
     }
 
     pub fn can_process_call_fully(
@@ -96,25 +96,37 @@ impl<'a, T: FieldElement, Q: QueryCallback<T>> MutableState<'a, T, Q> {
     }
 
     /// Extracts the witness column values from the machines.
-    fn take_witness_col_values(self) -> HashMap<String, Vec<T>> {
-        // We keep the already processed machines mutably borrowed so that
-        // "later" machines do not try to create new rows in already processed
-        // machines.
-        let mut processed = vec![];
-        self.machines
+    fn take_witness_col_and_public_values(self) -> (HashMap<String, Vec<T>>, BTreeMap<String, T>) {
+        let witness_columns = {
+            // We keep the already processed machines mutably borrowed so that
+            // "later" machines do not try to create new rows in already processed
+            // machines.
+            let mut processed = vec![];
+            self.machines
+                .iter()
+                .flat_map(|machine| {
+                    let mut machine = machine.try_borrow_mut().unwrap_or_else(|_| {
+                        panic!("Recursive machine dependencies while finishing machines.");
+                    });
+                    let columns = machine.take_witness_col_values(&self).into_iter();
+                    processed.push(machine);
+                    columns
+                })
+                .collect()
+        };
+
+        let public_values = self
+            .machines
             .iter()
             .flat_map(|machine| {
-                let mut machine = machine
-                    .try_borrow_mut()
-                    .map_err(|_| {
-                        panic!("Recursive machine dependencies while finishing machines.");
-                    })
-                    .unwrap();
-                let columns = machine.take_witness_col_values(&self).into_iter();
-                processed.push(machine);
-                columns
+                let mut machine = machine.try_borrow_mut().unwrap_or_else(|_| {
+                    panic!("Recursive machine dependencies while finishing machines.");
+                });
+                machine.take_public_values().into_iter()
             })
-            .collect()
+            .collect();
+
+        (witness_columns, public_values)
     }
 
     pub fn query_callback(&self) -> &Q {

executor/src/witgen/machines/block_machine.rs

@@ -393,6 +393,13 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
             .map(|(id, values)| (self.fixed_data.column_name(&id).to_string(), values))
             .collect()
     }
+
+    fn take_public_values(&mut self) -> BTreeMap<String, T> {
+        std::mem::take(&mut self.publics)
+            .into_iter()
+            .map(|(key, value)| (key.to_string(), value))
+            .collect()
+    }
 }
 
 impl<'a, T: FieldElement> BlockMachine<'a, T> {

executor/src/witgen/machines/dynamic_machine.rs

@@ -55,10 +55,11 @@ impl<'a, T: FieldElement> Machine<'a, T> for DynamicMachine<'a, T> {
     fn run<Q: QueryCallback<T>>(&mut self, mutable_state: &MutableState<'a, T, Q>) {
         assert!(self.data.is_empty());
         let first_row = self.compute_partial_first_row(mutable_state);
-        self.data = self
+        let process_result = self
             .process(first_row, 0, mutable_state, None, true)
-            .updated_data
-            .block;
+            .updated_data;
+        self.data = process_result.block;
+        self.publics.extend(process_result.publics);
     }
 
     fn process_plookup<'b, Q: QueryCallback<T>>(
@@ -124,6 +125,13 @@ impl<'a, T: FieldElement> Machine<'a, T> for DynamicMachine<'a, T> {
             .map(|(id, values)| (self.fixed_data.column_name(&id).to_string(), values))
             .collect()
     }
+
+    fn take_public_values(&mut self) -> BTreeMap<String, T> {
+        std::mem::take(&mut self.publics)
+            .into_iter()
+            .map(|(key, value)| (key.to_string(), value))
+            .collect()
+    }
 }
 
 impl<'a, T: FieldElement> DynamicMachine<'a, T> {

executor/src/witgen/machines/mod.rs

@@ -142,6 +142,10 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync {
 
     /// Returns the identity IDs of the connecting identities that this machine is responsible for.
     fn bus_ids(&self) -> Vec<T>;
+
+    fn take_public_values(&mut self) -> BTreeMap<String, T> {
+        BTreeMap::new()
+    }
 }
 
 #[repr(C)]
@@ -239,6 +243,10 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
         match_variant!(self, m => m.take_witness_col_values(mutable_state))
     }
 
+    fn take_public_values(&mut self) -> BTreeMap<String, T> {
+        match_variant!(self, m => m.take_public_values())
+    }
+
     fn bus_ids(&self) -> Vec<T> {
         match_variant!(self, m => m.bus_ids())
     }

executor/src/witgen/mod.rs

@@ -50,6 +50,9 @@ mod vm_processor;
 pub use affine_expression::{AffineExpression, AffineResult, AlgebraicVariable};
 pub use evaluators::partial_expression_evaluator::{PartialExpressionEvaluator, SymbolicVariables};
 
+pub type Witness<T> = Vec<(String, Vec<T>)>;
+pub type Publics<T> = BTreeMap<String, Option<T>>;
+
 static OUTER_CODE_NAME: &str = "witgen (outer code)";
 
 // TODO change this so that it has functions
@@ -141,6 +144,7 @@ impl<T: FieldElement> WitgenCallbackContext<T> {
                 .with_external_witness_values(current_witness)
                 .with_challenges(stage, challenges)
                 .generate()
+                .0
         }
     }
 }
@@ -202,7 +206,7 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
 
     /// Generates the committed polynomial values
     /// @returns the values (in source order) and the degree of the polynomials.
-    pub fn generate(self) -> Vec<(String, Vec<T>)> {
+    pub fn generate(self) -> (Witness<T>, Publics<T>) {
         record_start(OUTER_CODE_NAME);
         let fixed = FixedData::new(
             self.analyzed,
@@ -241,7 +245,8 @@ 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 = MutableState::new(machines.into_iter(), &self.query_callback).run();
+        let (columns, _publics) =
+            MutableState::new(machines.into_iter(), &self.query_callback).run();
 
         let publics = extract_publics(&columns, self.analyzed);
         if !publics.is_empty() {
@@ -258,7 +263,7 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
 
         let mut columns = if self.stage == 0 {
             // Multiplicities should be computed in the first stage
-            MultiplicityColumnGenerator::new(&fixed).generate(columns, publics)
+            MultiplicityColumnGenerator::new(&fixed).generate(columns, publics.clone())
         } else {
             columns
         };
@@ -279,7 +284,8 @@ impl<'a, 'b, T: FieldElement> WitnessGenerator<'a, 'b, T> {
                 (name, column)
             })
             .collect::<Vec<_>>();
-        witness_cols
+
+        (witness_cols, publics)
     }
 }