Merge pull request #903 from powdr-labs/input-hash

Support serde data structures in pipeline

pipeline/Cargo.toml

@@ -24,6 +24,8 @@ analysis = { path = "../analysis" }
 linker = { path = "../linker" }
 airgen = { path = "../airgen" }
 importer = { path = "../importer" }
+serde = { version = "1.0", default-features = false, features = ["alloc", "derive", "rc"] }
+serde_cbor = "0.11.2"
 
 [dev-dependencies]
 test-log = "0.2.12"

pipeline/src/lib.rs

@@ -2,6 +2,8 @@
 
 #![deny(clippy::print_stdout)]
 
+use std::marker::{Send, Sync};
+
 pub mod pipeline;
 pub mod test_util;
 pub mod util;
@@ -45,6 +47,38 @@ pub fn access_element<T: FieldElement>(
     }
 }
 
+#[allow(clippy::print_stdout)]
+pub fn serde_data_to_query_callback<T: FieldElement, S: serde::Serialize + Send + Sync>(
+    channel: u32,
+    data: &S,
+) -> impl QueryCallback<T> {
+    let bytes = serde_cbor::to_vec(&data).unwrap();
+    move |query: &str| -> Result<Option<T>, String> {
+        match &parse_query(query)?[..] {
+            ["\"data_identifier\"", index, cb_channel] => {
+                let cb_channel = cb_channel
+                    .parse::<u32>()
+                    .map_err(|e| format!("Error parsing callback data channel: {e})"))?;
+
+                if channel != cb_channel {
+                    return Ok(None);
+                }
+
+                let index = index
+                    .parse::<usize>()
+                    .map_err(|e| format!("Error parsing index: {e})"))?;
+
+                // query index 0 means the length
+                Ok(Some(match index {
+                    0 => (bytes.len() as u64).into(),
+                    index => (bytes[index - 1] as u64).into(),
+                }))
+            }
+            k => Err(format!("Unsupported query: {}", k.iter().format(", "))),
+        }
+    }
+}
+
 #[allow(clippy::print_stdout)]
 pub fn inputs_to_query_callback<T: FieldElement>(inputs: Vec<T>) -> impl QueryCallback<T> {
     move |query: &str| -> Result<Option<T>, String> {
@@ -54,14 +88,6 @@ pub fn inputs_to_query_callback<T: FieldElement>(inputs: Vec<T>) -> impl QueryCa
 
         match &parse_query(query)?[..] {
             ["\"input\"", index] => access_element("prover inputs", &inputs, index),
-            ["\"data\"", index, what] => {
-                let what = what
-                    .parse::<usize>()
-                    .map_err(|e| format!("Error parsing what: {e})"))?;
-                assert_eq!(what, 0);
-
-                access_element("prover inputs", &inputs, index)
-            }
             ["\"print_char\"", ch] => {
                 print!(
                     "{}",

pipeline/src/pipeline.rs

@@ -2,6 +2,7 @@ use std::{
     fmt::Display,
     fs,
     io::{BufWriter, Read, Write},
+    marker::Send,
     path::{Path, PathBuf},
     rc::Rc,
     time::Instant,
@@ -24,7 +25,7 @@ use mktemp::Temp;
 use number::{write_polys_csv_file, write_polys_file, CsvRenderMode, FieldElement};
 
 use crate::{
-    inputs_to_query_callback,
+    inputs_to_query_callback, serde_data_to_query_callback,
     util::{read_poly_set, FixedPolySet, WitnessPolySet},
 };
 
@@ -70,7 +71,7 @@ pub enum Stage {
     Proof,
 }
 
-enum Artifact<T: FieldElement> {
+pub enum Artifact<T: FieldElement> {
     /// The path to a single .asm file.
     AsmFilePath(PathBuf),
     /// The contents of a single .asm file, with an optional Path (for imports).
@@ -106,6 +107,53 @@ enum Artifact<T: FieldElement> {
     Proof(ProofResult<T>),
 }
 
+// These are implementations of specific artifacts we want to retrieve
+// from Pipeline, in a way that allows us to get immutable references
+// to the artifacts.
+impl<T: FieldElement> Artifact<T> {
+    pub fn to_asm_string(&self) -> Option<&String> {
+        match self {
+            Artifact::AsmString(_, asm_string) => Some(asm_string),
+            _ => None,
+        }
+    }
+
+    pub fn to_analyzed_asm(&self) -> Option<&AnalysisASMFile<T>> {
+        match self {
+            Artifact::AnalyzedAsm(analyzed_asm) => Some(analyzed_asm),
+            _ => None,
+        }
+    }
+
+    pub fn to_optimized_pil(&self) -> Option<&Analyzed<T>> {
+        match self {
+            Artifact::OptimzedPil(optimized_pil) => Some(optimized_pil),
+            _ => None,
+        }
+    }
+
+    pub fn to_pil_with_evaluated_fixed_cols(&self) -> Option<&PilWithEvaluatedFixedCols<T>> {
+        match self {
+            Artifact::PilWithEvaluatedFixedCols(pil_with_constants) => Some(pil_with_constants),
+            _ => None,
+        }
+    }
+
+    pub fn to_generated_witness(&self) -> Option<&GeneratedWitness<T>> {
+        match self {
+            Artifact::GeneratedWitness(generated_witness) => Some(generated_witness),
+            _ => None,
+        }
+    }
+
+    pub fn to_proof(&self) -> Option<&ProofResult<T>> {
+        match self {
+            Artifact::Proof(proof) => Some(proof),
+            _ => None,
+        }
+    }
+}
+
 /// Optional Arguments for various stages of the pipeline.
 #[derive(Default)]
 struct Arguments<T: FieldElement> {
@@ -261,6 +309,14 @@ impl<T: FieldElement> Pipeline<T> {
         self
     }
 
+    pub fn add_data<S: serde::Serialize + Send + Sync + 'static>(
+        self,
+        channel: u32,
+        data: &S,
+    ) -> Self {
+        self.add_query_callback(Box::new(serde_data_to_query_callback(channel, data)))
+    }
+
     pub fn with_prover_inputs(self, inputs: Vec<T>) -> Self {
         self.add_query_callback(Box::new(inputs_to_query_callback(inputs)))
     }
@@ -736,4 +792,12 @@ impl<T: FieldElement> Pipeline<T> {
     pub fn name(&self) -> &str {
         self.name.as_ref().unwrap()
     }
+
+    pub fn artifact(&self) -> Option<&Artifact<T>> {
+        self.artifact.as_ref()
+    }
+
+    pub fn data_callback(&self) -> Option<&dyn QueryCallback<T>> {
+        self.arguments.query_callback.as_deref()
+    }
 }

pipeline/tests/asm.rs

@@ -22,17 +22,17 @@ fn simple_sum_asm() {
 #[test]
 fn secondary_block_machine_add2() {
     let f = "asm/secondary_block_machine_add2.asm";
-    verify_asm::<GoldilocksField>(f, vec![]);
-    gen_halo2_proof(f, vec![]);
-    gen_estark_proof(f, vec![]);
+    verify_asm::<GoldilocksField>(f, Default::default());
+    gen_halo2_proof(f, Default::default());
+    gen_estark_proof(f, Default::default());
 }
 
 #[test]
 fn mem_write_once() {
     let f = "asm/mem_write_once.asm";
-    verify_asm::<GoldilocksField>(f, vec![]);
-    gen_halo2_proof(f, vec![]);
-    gen_estark_proof(f, vec![]);
+    verify_asm::<GoldilocksField>(f, Default::default());
+    gen_halo2_proof(f, Default::default());
+    gen_estark_proof(f, Default::default());
 }
 
 #[test]
@@ -42,15 +42,15 @@ fn mem_write_once_external_write() {
     mem[17] = GoldilocksField::from(42);
     mem[62] = GoldilocksField::from(123);
     mem[255] = GoldilocksField::from(-1);
-    verify_test_file::<GoldilocksField>(f, vec![], vec![("main.v".to_string(), mem)]);
+    verify_test_file::<GoldilocksField>(f, Default::default(), vec![("main.v".to_string(), mem)]);
 }
 
 #[test]
 fn block_machine_cache_miss() {
     let f = "asm/block_machine_cache_miss.asm";
-    verify_asm::<GoldilocksField>(f, vec![]);
-    gen_halo2_proof(f, vec![]);
-    gen_estark_proof(f, vec![]);
+    verify_asm::<GoldilocksField>(f, Default::default());
+    gen_halo2_proof(f, Default::default());
+    gen_estark_proof(f, Default::default());
 }
 
 #[test]

pipeline/tests/pil.rs

@@ -10,23 +10,23 @@ pub fn verify_pil(file_name: &str, inputs: Vec<GoldilocksField>) {
 #[test]
 fn test_fibonacci() {
     let f = "pil/fibonacci.pil";
-    verify_pil(f, vec![]);
+    verify_pil(f, Default::default());
     gen_halo2_proof(f, Default::default());
     gen_estark_proof(f, Default::default());
 }
 
 #[test]
 fn test_constant_in_identity() {
     let f = "pil/constant_in_identity.pil";
-    verify_pil(f, vec![]);
+    verify_pil(f, Default::default());
     gen_halo2_proof(f, Default::default());
     gen_estark_proof(f, Default::default());
 }
 
 #[test]
 fn fib_arrays() {
     let f = "pil/fib_arrays.pil";
-    verify_pil(f, vec![]);
+    verify_pil(f, Default::default());
     gen_halo2_proof(f, Default::default());
     gen_estark_proof(f, Default::default());
 }
@@ -35,21 +35,21 @@ fn fib_arrays() {
 #[should_panic = "Witness generation failed."]
 fn test_external_witgen_fails_if_none_provided() {
     let f = "pil/external_witgen.pil";
-    verify_pil(f, vec![]);
+    verify_pil(f, Default::default());
 }
 
 #[test]
 fn test_external_witgen_a_provided() {
     let f = "pil/external_witgen.pil";
     let external_witness = vec![("main.a".to_string(), vec![GoldilocksField::from(3); 16])];
-    verify_test_file(f, vec![], external_witness);
+    verify_test_file(f, Default::default(), external_witness);
 }
 
 #[test]
 fn test_external_witgen_b_provided() {
     let f = "pil/external_witgen.pil";
     let external_witness = vec![("main.b".to_string(), vec![GoldilocksField::from(4); 16])];
-    verify_test_file(f, vec![], external_witness);
+    verify_test_file(f, Default::default(), external_witness);
 }
 
 #[test]
@@ -59,7 +59,7 @@ fn test_external_witgen_both_provided() {
         ("main.a".to_string(), vec![GoldilocksField::from(3); 16]),
         ("main.b".to_string(), vec![GoldilocksField::from(4); 16]),
     ];
-    verify_test_file(f, vec![], external_witness);
+    verify_test_file(f, Default::default(), external_witness);
 }
 
 #[test]
@@ -71,12 +71,12 @@ fn test_external_witgen_fails_on_conflicting_external_witness() {
         // Does not satisfy b = a + 1
         ("main.b".to_string(), vec![GoldilocksField::from(3); 16]),
     ];
-    verify_test_file(f, vec![], external_witness);
+    verify_test_file(f, Default::default(), external_witness);
 }
 
 #[test]
 fn test_global() {
-    verify_pil("pil/global.pil", vec![]);
+    verify_pil("pil/global.pil", Default::default());
     // Halo2 would take too long for this.
     // Starky requires at least one witness column, this test has none.
 }
@@ -107,69 +107,72 @@ fn test_witness_lookup() {
 #[test]
 #[should_panic(expected = "Witness generation failed.")]
 fn test_underdetermined_zero_no_solution() {
-    verify_pil("pil/underdetermined_zero_no_solution.pil", vec![]);
+    verify_pil(
+        "pil/underdetermined_zero_no_solution.pil",
+        Default::default(),
+    );
 }
 
 #[test]
 fn test_pair_lookup() {
     let f = "pil/pair_lookup.pil";
-    verify_pil(f, vec![]);
+    verify_pil(f, Default::default());
     // halo2 would take too long for this
     // starky would take too long for this in debug mode
 }
 
 #[test]
 fn test_block_lookup_or() {
     let f = "pil/block_lookup_or.pil";
-    verify_pil(f, vec![]);
+    verify_pil(f, Default::default());
     // halo2 would take too long for this
     // starky would take too long for this in debug mode
 }
 
 #[test]
 fn test_halo_without_lookup() {
     let f = "pil/halo_without_lookup.pil";
-    verify_pil(f, vec![]);
+    verify_pil(f, Default::default());
     gen_halo2_proof(f, Default::default());
     gen_estark_proof(f, Default::default());
 }
 
 #[test]
 fn test_simple_div() {
     let f = "pil/simple_div.pil";
-    verify_pil(f, vec![]);
+    verify_pil(f, Default::default());
     // starky would take too long for this in debug mode
 }
 
 #[test]
 fn test_single_line_blocks() {
     let f = "pil/single_line_blocks.pil";
-    verify_pil(f, vec![]);
+    verify_pil(f, Default::default());
     gen_estark_proof(f, Default::default());
 }
 
 #[test]
 fn test_two_block_machine_functions() {
     let f = "pil/two_block_machine_functions.pil";
-    verify_pil(f, vec![]);
+    verify_pil(f, Default::default());
     gen_estark_proof(f, Default::default());
 }
 
 #[test]
 fn test_fixed_columns() {
     let f = "pil/fixed_columns.pil";
-    verify_pil(f, vec![]);
+    verify_pil(f, Default::default());
     // Starky requires at least one witness column, this test has none.
 }
 
 #[test]
 fn test_witness_via_let() {
-    verify_pil("pil/witness_via_let.pil", vec![]);
+    verify_pil("pil/witness_via_let.pil", Default::default());
 }
 
 #[test]
 fn conditional_fixed_constraints() {
-    verify_pil("pil/conditional_fixed_constraints.pil", vec![]);
+    verify_pil("pil/conditional_fixed_constraints.pil", Default::default());
 }
 
 #[test]
@@ -187,7 +190,7 @@ mod book {
     use test_log::test;
 
     fn run_book_test(file: &str) {
-        verify_pil(file, vec![]);
+        verify_pil(file, Default::default());
         gen_halo2_proof(file, Default::default());
         gen_estark_proof(file, Default::default());
     }