From 668d881057d397b65387ba2bb2088e49db494f7b Mon Sep 17 00:00:00 2001
From: Georg Wiese <georgwiese@gmail.com>
Date: Thu, 15 Feb 2024 12:43:47 +0100
Subject: [PATCH] On (N-1, 0) row pair, only execute identities with next
 reference

---
 executor/src/witgen/generator.rs        | 14 ++++++++++++--
 pipeline/tests/pil.rs                   |  2 +-
 test_data/pil/sum_via_witness_query.pil |  8 +++++++-
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/executor/src/witgen/generator.rs b/executor/src/witgen/generator.rs
index 4077d51aa..f125cf578 100644
--- a/executor/src/witgen/generator.rs
+++ b/executor/src/witgen/generator.rs
@@ -179,16 +179,26 @@ impl<'a, T: FieldElement> Generator<'a, T> {
             ]
             .into_iter(),
         );
+
+        // We're only interested in the first row anyway, so identities without a next reference
+        // are irrelevant.
+        // Also, they can lead to problems in the case where some witness columns are provided
+        // externally, e.g. if the last row happens to call into a stateful machine like memory.
+        let identities_with_next_reference = self
+            .identities
+            .iter()
+            .filter_map(|identity| identity.contains_next_ref().then_some(*identity))
+            .collect::<Vec<_>>();
         let mut processor = BlockProcessor::new(
             self.fixed_data.degree - 1,
             data,
             mutable_state,
-            &self.identities,
+            &identities_with_next_reference,
             self.fixed_data,
             &self.witnesses,
         );
         let mut sequence_iterator = ProcessingSequenceIterator::Default(
-            DefaultSequenceIterator::new(0, self.identities.len(), None),
+            DefaultSequenceIterator::new(0, identities_with_next_reference.len(), None),
         );
         processor.solve(&mut sequence_iterator).unwrap();
         let first_row = processor.finish().remove(1);
diff --git a/pipeline/tests/pil.rs b/pipeline/tests/pil.rs
index 05550395f..c65571a2f 100644
--- a/pipeline/tests/pil.rs
+++ b/pipeline/tests/pil.rs
@@ -123,7 +123,7 @@ fn test_external_witgen_both_provided() {
 }
 
 #[test]
-#[should_panic = "called `Result::unwrap()` on an `Err` value: Generic(\"main.b = (main.a + 1);:\\n    Linear constraint is not satisfiable: 18446744069414584320 != 0\")"]
+#[should_panic = "Witness generation failed."]
 fn test_external_witgen_fails_on_conflicting_external_witness() {
     let f = "pil/external_witgen.pil";
     let external_witness = vec![
diff --git a/test_data/pil/sum_via_witness_query.pil b/test_data/pil/sum_via_witness_query.pil
index 47e4f386d..256a2f2bf 100644
--- a/test_data/pil/sum_via_witness_query.pil
+++ b/test_data/pil/sum_via_witness_query.pil
@@ -7,7 +7,13 @@ namespace Sum(%N);
     pol fixed ISALMOSTLAST(i) { if i == %last_row - 1 { 1 } else { 0 } };
     pol fixed ISFIRST = [ 1, 0 ] + [0]*;
 
-    col witness input(i) query ("input", i);
+    col witness input(i) query match i {
+        // A non-exhaustive match statement is the only way to return "None"
+        0 => ("input", 0),
+        1 => ("input", 1),
+        2 => ("input", 2),
+        // No response in the case of i == 3
+    };
     col witness sum;
 
     ISLAST * sum' = 0;