Stage 1 public reference (#2567)

Depends on #2556. Pass stage 1 publics to the backend but not using them
yet (TODO for future PR).

Author: qwang98

backend/src/composite/mod.rs

@@ -383,7 +383,7 @@ impl<F: FieldElement> Backend<F> for CompositeBackend<F> {
                     .map(|(machine_name, machine_data)| {
                         let (machine_witness, size) = witness_by_machine.get(machine_name).unwrap();
                         let machine_data = machine_data.get(size).unwrap();
-                        let new_witness = witgen_callback.next_stage_witness(
+                        let (new_witness, _) = witgen_callback.next_stage_witness(
                             &machine_data.pil,
                             machine_witness,
                             challenges.clone(),

backend/src/composite/sub_prover.rs

@@ -8,7 +8,7 @@ use std::{
 };
 
 use crate::{Backend, Error};
-use powdr_executor::witgen::WitgenCallback;
+use powdr_executor::witgen::{Publics, WitgenCallback, Witness};
 use powdr_number::FieldElement;
 
 /// Runs a prover until it either completes or challenges the caller, in
@@ -69,12 +69,14 @@ pub enum RunStatus<'s, F: FieldElement> {
 pub struct SubProver<'s, F: FieldElement> {
     pub thread: ScopedJoinHandle<'s, Result<Vec<u8>, Error>>,
     pub challenge_receiver: Receiver<BTreeMap<u64, F>>,
-    pub response_sender: SyncSender<Vec<(String, Vec<F>)>>,
+    pub response_sender: SyncSender<(Witness<F>, Publics<F>)>,
 }
 
 impl<'s, F: FieldElement> SubProver<'s, F> {
     pub fn resume(self, response: Vec<(String, Vec<F>)>) -> RunStatus<'s, F> {
-        self.response_sender.send(response).unwrap();
+        self.response_sender
+            .send((response, BTreeMap::new()))
+            .unwrap();
         self.wait()
     }
 

backend/src/halo2/circuit_builder.rs

@@ -393,7 +393,7 @@ impl<T: FieldElement, F: PrimeField<Repr = [u8; 32]>> Circuit<F> for PowdrCircui
                     "Running witness generation for stage {stage} ({} challenges)!",
                     challenges.len()
                 );
-                new_witness = self
+                (new_witness, _) = self
                     .witgen_callback
                     .as_ref()
                     .expect("Expected witgen callback!")

backend/src/mock/machine.rs

@@ -41,7 +41,7 @@ impl<'a, F: FieldElement> Machine<'a, F> {
 
         for stage in 1..pil.stage_count() {
             log::debug!("Generating stage-{stage} witness for machine {machine_name}");
-            witness =
+            (witness, _) =
                 witgen_callback.next_stage_witness(pil, &witness, challenges.clone(), stage as u8);
         }
 

backend/src/stwo/prover.rs

@@ -411,12 +411,14 @@ where
             let stage1_witness_cols = witness_by_machine
                 .iter()
                 .map(|(machine_name, machine_witness)| {
-                    witgen_callback.next_stage_witness(
-                        &self.split[&machine_name.clone()],
-                        machine_witness,
-                        stage0_challenges.clone(),
-                        1,
-                    )
+                    witgen_callback
+                        .next_stage_witness(
+                            &self.split[&machine_name.clone()],
+                            machine_witness,
+                            stage0_challenges.clone(),
+                            1,
+                        )
+                        .0
                 })
                 .collect_vec();
 

executor-utils/src/lib.rs

@@ -14,11 +14,19 @@ pub mod expression_evaluator;
 /// - The current witness.
 /// - The challenges sampled so far.
 pub type WitgenCallbackFn<T> = Arc<
-    dyn Fn(&Analyzed<T>, &[(String, Vec<T>)], BTreeMap<u64, T>, u8) -> Vec<(String, Vec<T>)>
+    dyn Fn(
+            &Analyzed<T>,
+            &[(String, Vec<T>)],
+            BTreeMap<u64, T>,
+            u8,
+        ) -> (Vec<(String, Vec<T>)>, BTreeMap<String, Option<T>>)
         + Send
         + Sync,
 >;
 
+pub type Witness<T> = Vec<(String, Vec<T>)>;
+pub type Publics<T> = BTreeMap<String, Option<T>>;
+
 #[derive(Clone)]
 pub struct WitgenCallback<T>(WitgenCallbackFn<T>);
 
@@ -34,7 +42,7 @@ impl<T: FieldElement> WitgenCallback<T> {
         current_witness: &[(String, Vec<T>)],
         challenges: BTreeMap<u64, T>,
         stage: u8,
-    ) -> Vec<(String, Vec<T>)> {
+    ) -> (Witness<T>, Publics<T>) {
         (self.0)(pil, current_witness, challenges, stage)
     }
 }

executor/src/witgen/machines/second_stage_machine.rs

@@ -51,7 +51,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for SecondStageMachine<'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.process(first_row, mutable_state);
+        (self.data, self.publics) = self.process(first_row, mutable_state);
     }
 
     fn process_plookup<'b, Q: QueryCallback<T>>(
@@ -78,6 +78,13 @@ impl<'a, T: FieldElement> Machine<'a, T> for SecondStageMachine<'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> SecondStageMachine<'a, T> {
@@ -169,7 +176,7 @@ impl<'a, T: FieldElement> SecondStageMachine<'a, T> {
         &mut self,
         first_row: Row<T>,
         mutable_state: &MutableState<'a, T, Q>,
-    ) -> FinalizableData<'a, T> {
+    ) -> (FinalizableData<'a, T>, BTreeMap<&'a str, T>) {
         log::trace!(
             "Running Second-Stage Machine with the following initial values in the first row:\n{}",
             first_row.render_values(false, &self.parts)
@@ -196,7 +203,7 @@ impl<'a, T: FieldElement> SecondStageMachine<'a, T> {
         // The processor might have detected a loop, in which case the degree has changed
         self.degree = degree;
 
-        updated_data.block
+        (updated_data.block, updated_data.publics)
     }
 
     /// At the end of the solving algorithm, we'll have computed the first row twice

executor/src/witgen/mod.rs

@@ -111,7 +111,7 @@ impl<T: FieldElement> WitgenCallbackContext<T> {
         current_witness: &[(String, Vec<T>)],
         challenges: BTreeMap<u64, T>,
         stage: u8,
-    ) -> Vec<(String, Vec<T>)> {
+    ) -> (Witness<T>, Publics<T>) {
         let has_phantom_bus_sends = pil
             .identities
             .iter()
@@ -135,7 +135,10 @@ impl<T: FieldElement> WitgenCallbackContext<T> {
                 challenges,
             );
 
-            current_witness.iter().cloned().chain(bus_columns).collect()
+            (
+                current_witness.iter().cloned().chain(bus_columns).collect(),
+                BTreeMap::new(),
+            )
         } else {
             log::debug!("Using automatic stage-1 witgen.");
             let size = current_witness.iter().next().unwrap().1.len() as DegreeType;
@@ -144,7 +147,6 @@ impl<T: FieldElement> WitgenCallbackContext<T> {
                 .with_external_witness_values(current_witness)
                 .with_challenges(stage, challenges)
                 .generate()
-                .0
         }
     }
 }

plonky3/src/circuit_builder.rs

@@ -468,12 +468,13 @@ where
             witness_by_machine
                 .par_iter()
                 .map(|(machine_name, machine_witness)| {
-                    let new_witness = self.witgen_callback.as_ref().unwrap().next_stage_witness(
-                        &self.split[machine_name].0,
-                        machine_witness,
-                        challenge_map.clone(),
-                        trace_stage,
-                    );
+                    let (new_witness, _) =
+                        self.witgen_callback.as_ref().unwrap().next_stage_witness(
+                            &self.split[machine_name].0,
+                            machine_witness,
+                            challenge_map.clone(),
+                            trace_stage,
+                        );
                     (machine_name.clone(), new_witness)
                 })
                 .collect()