Merge pull request powdr-labs#705 from powdr-labs/custom-witgen

External Witness Generation

compiler/benches/executor_benchmark.rs

@@ -29,7 +29,8 @@ fn get_pil() -> Analyzed<T> {
 fn run_witgen<T: FieldElement>(analyzed: &Analyzed<T>, input: Vec<T>) {
     let query_callback = Some(inputs_to_query_callback(input));
     let (constants, degree) = constant_evaluator::generate(analyzed);
-    executor::witgen::generate(analyzed, degree, &constants, query_callback);
+    executor::witgen::WitnessGenerator::new(analyzed, degree, &constants, query_callback)
+        .generate();
 }
 
 fn criterion_benchmark(c: &mut Criterion) {

compiler/src/lib.rs

@@ -48,6 +48,7 @@ pub fn compile_pil_or_asm<T: FieldElement>(
             output_dir,
             Some(inputs_to_query_callback(inputs)),
             prove_with,
+            vec![],
         )))
     }
 }
@@ -65,13 +66,15 @@ pub fn compile_pil<T: FieldElement, Q: QueryCallback<T>>(
     output_dir: &Path,
     query_callback: Option<Q>,
     prove_with: Option<BackendType>,
+    external_witness_values: Vec<(&str, Vec<T>)>,
 ) -> CompilationResult<T> {
     compile(
         pil_analyzer::analyze(pil_file),
         pil_file.file_name().unwrap(),
         output_dir,
         query_callback,
         prove_with,
+        external_witness_values,
     )
 }
 
@@ -92,6 +95,7 @@ pub fn compile_pil_ast<T: FieldElement, Q: QueryCallback<T>>(
         output_dir,
         query_callback,
         prove_with,
+        vec![],
     )
 }
 
@@ -194,6 +198,7 @@ fn compile<T: FieldElement, Q: QueryCallback<T>>(
     output_dir: &Path,
     query_callback: Option<Q>,
     prove_with: Option<BackendType>,
+    external_witness_values: Vec<(&str, Vec<T>)>,
 ) -> CompilationResult<T> {
     log::info!("Optimizing pil...");
     let analyzed = pilopt::optimize(analyzed);
@@ -210,7 +215,10 @@ fn compile<T: FieldElement, Q: QueryCallback<T>>(
 
     let witness = (analyzed.constant_count() == constants.len()).then(|| {
         log::info!("Deducing witness columns...");
-        let commits = executor::witgen::generate(&analyzed, degree, &constants, query_callback);
+        let commits =
+            executor::witgen::WitnessGenerator::new(&analyzed, degree, &constants, query_callback)
+                .with_external_witness_values(external_witness_values)
+                .generate();
 
         let witness = commits
             .into_iter()

compiler/tests/pil.rs

@@ -4,6 +4,14 @@ use std::path::Path;
 use test_log::test;
 
 pub fn verify_pil(file_name: &str, query_callback: Option<fn(&str) -> Option<GoldilocksField>>) {
+    verify_pil_with_external_witness(file_name, query_callback, vec![]);
+}
+
+pub fn verify_pil_with_external_witness(
+    file_name: &str,
+    query_callback: Option<fn(&str) -> Option<GoldilocksField>>,
+    external_witness_values: Vec<(&str, Vec<GoldilocksField>)>,
+) {
     let input_file = Path::new(&format!(
         "{}/../test_data/pil/{file_name}",
         env!("CARGO_MANIFEST_DIR")
@@ -16,7 +24,8 @@ pub fn verify_pil(file_name: &str, query_callback: Option<fn(&str) -> Option<Gol
         &input_file,
         &temp_dir,
         query_callback,
-        Some(BackendType::PilStarkCli)
+        Some(BackendType::PilStarkCli),
+        external_witness_values
     )
     .witness
     .is_some());
@@ -81,6 +90,49 @@ fn test_fibonacci_macro() {
     gen_estark_proof(f, Default::default());
 }
 
+#[test]
+#[should_panic = "Witness generation failed."]
+fn test_external_witgen_fails_if_none_provided() {
+    let f = "external_witgen.pil";
+    verify_pil(f, None);
+}
+
+#[test]
+fn test_external_witgen_a_provided() {
+    let f = "external_witgen.pil";
+    let external_witness = vec![("main.a", vec![GoldilocksField::from(3); 16])];
+    verify_pil_with_external_witness(f, None, external_witness);
+}
+
+#[test]
+fn test_external_witgen_b_provided() {
+    let f = "external_witgen.pil";
+    let external_witness = vec![("main.b", vec![GoldilocksField::from(4); 16])];
+    verify_pil_with_external_witness(f, None, external_witness);
+}
+
+#[test]
+fn test_external_witgen_both_provided() {
+    let f = "external_witgen.pil";
+    let external_witness = vec![
+        ("main.a", vec![GoldilocksField::from(3); 16]),
+        ("main.b", vec![GoldilocksField::from(4); 16]),
+    ];
+    verify_pil_with_external_witness(f, None, external_witness);
+}
+
+#[test]
+#[should_panic = "called `Result::unwrap()` on an `Err` value: ConstraintUnsatisfiable(\"-1\")"]
+fn test_external_witgen_fails_on_conflicting_external_witness() {
+    let f = "external_witgen.pil";
+    let external_witness = vec![
+        ("main.a", vec![GoldilocksField::from(3); 16]),
+        // Does not satisfy b = a + 1
+        ("main.b", vec![GoldilocksField::from(3); 16]),
+    ];
+    verify_pil_with_external_witness(f, None, external_witness);
+}
+
 #[test]
 fn test_global() {
     verify_pil("global.pil", None);

executor/src/witgen/generator.rs

@@ -91,7 +91,10 @@ impl<'a, T: FieldElement> Generator<'a, T> {
             mutable_state.machines.iter_mut().into(),
         );
         let row_factory = RowFactory::new(self.fixed_data, self.global_range_constraints.clone());
-        let data = vec![row_factory.fresh_row(); 2];
+        let data = vec![
+            row_factory.fresh_row(self.fixed_data.degree - 1),
+            row_factory.fresh_row(0),
+        ];
         let mut processor = Processor::new(
             self.fixed_data.degree - 1,
             data,

executor/src/witgen/machines/block_machine.rs

@@ -69,7 +69,9 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
                 // Start out with a block filled with unknown values so that we do not have to deal with wrap-around
                 // when storing machine witness data.
                 // This will be filled with the default block in `take_witness_col_values`
-                let data = vec![row_factory.fresh_row(); block_size];
+                let data = (0..block_size)
+                    .map(|i| row_factory.fresh_row(i as DegreeType))
+                    .collect();
                 return Some(BlockMachine {
                     block_size,
                     selected_expressions: id.right.clone(),
@@ -274,7 +276,7 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
             let current = &self.data[row as usize];
             // We don't have the next row, because it would be the first row of the next block.
             // We'll use a fresh row instead.
-            let next = self.row_factory.fresh_row();
+            let next = self.row_factory.fresh_row(row + 1);
             let row_pair = RowPair::new(
                 current,
                 &next,
@@ -353,11 +355,13 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         right: &'a SelectedExpressions<Expression<T>>,
         sequence_iterator: &mut ProcessingSequenceIterator,
     ) -> Result<ProcessResult<'a, T>, EvalError<T>> {
+        // We start at the last row of the previous block.
+        let row_offset = self.rows() - 1;
         // Make the block two rows larger than the block size, it includes the last row of the previous block
         // and the first row of the next block.
-        let block = vec![self.row_factory.fresh_row(); self.block_size + 2];
-        // We start at the last row of the previous block.
-        let row_offset = self.data.len() as DegreeType - 1;
+        let block = (0..(self.block_size + 2))
+            .map(|i| self.row_factory.fresh_row(i as DegreeType + row_offset))
+            .collect();
         let mut processor = Processor::new(
             row_offset,
             block,