From 93e38d4867a1f781b4bfce3f5b4d414d2e16e26e Mon Sep 17 00:00:00 2001
From: Georg Wiese <georgwiese@gmail.com>
Date: Thu, 2 Nov 2023 14:51:13 +0000
Subject: [PATCH] Change process_identity() interface

---
 executor/src/witgen/block_processor.rs | 16 ++++++++++------
 executor/src/witgen/processor.rs       |  9 +--------
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/executor/src/witgen/block_processor.rs b/executor/src/witgen/block_processor.rs
index 514fa8d1e..0d2ad830f 100644
--- a/executor/src/witgen/block_processor.rs
+++ b/executor/src/witgen/block_processor.rs
@@ -19,6 +19,8 @@ use super::{
 /// - `'c`: The duration of this Processor's lifetime (e.g. the reference to the identity processor)
 pub struct BlockProcessor<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> {
     processor: Processor<'a, 'b, 'c, T, Q>,
+    /// The list of identities
+    identities: &'c [&'a Identity<Expression<T>>],
 }
 
 impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> BlockProcessor<'a, 'b, 'c, T, Q> {
@@ -35,12 +37,14 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> BlockProcessor<'a, 'b, 'c
             row_offset,
             data,
             mutable_state,
-            identities,
             fixed_data,
             row_factory,
             witness_cols,
         );
-        Self { processor }
+        Self {
+            processor,
+            identities,
+        }
     }
 
     pub fn with_outer_query(
@@ -48,7 +52,7 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> BlockProcessor<'a, 'b, 'c
         outer_query: OuterQuery<'a, T>,
     ) -> BlockProcessor<'a, 'b, 'c, T, Q> {
         let processor = self.processor.with_outer_query(outer_query);
-        Self { processor }
+        Self { processor, ..self }
     }
 
     /// Figures out unknown values.
@@ -62,9 +66,9 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> BlockProcessor<'a, 'b, 'c
         while let Some(SequenceStep { row_delta, action }) = sequence_iterator.next() {
             let row_index = (1 + row_delta) as usize;
             let progress = match action {
-                Action::InternalIdentity(identity_index) => {
-                    self.processor.process_identity(row_index, identity_index)?
-                }
+                Action::InternalIdentity(identity_index) => self
+                    .processor
+                    .process_identity(row_index, self.identities[identity_index])?,
                 Action::OuterQuery => {
                     let (progress, new_outer_assignments) =
                         self.processor.process_outer_query(row_index)?;
diff --git a/executor/src/witgen/processor.rs b/executor/src/witgen/processor.rs
index f1f51283b..c8466e639 100644
--- a/executor/src/witgen/processor.rs
+++ b/executor/src/witgen/processor.rs
@@ -48,8 +48,6 @@ pub struct Processor<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> {
     row_offset: u64,
     /// The rows that are being processed.
     data: FinalizableData<'a, T>,
-    /// The list of identities
-    identities: &'c [&'a Identity<Expression<T>>],
     /// The mutable state
     mutable_state: &'c mut MutableState<'a, 'b, T, Q>,
     /// The fixed data (containing information about all columns)
@@ -69,7 +67,6 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> Processor<'a, 'b, 'c, T,
         row_offset: u64,
         data: FinalizableData<'a, T>,
         mutable_state: &'c mut MutableState<'a, 'b, T, Q>,
-        identities: &'c [&'a Identity<Expression<T>>],
         fixed_data: &'a FixedData<'a, T>,
         row_factory: RowFactory<'a, T>,
         witness_cols: &'c HashSet<PolyID>,
@@ -84,7 +81,6 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> Processor<'a, 'b, 'c, T,
             row_offset,
             data,
             mutable_state,
-            identities,
             fixed_data,
             row_factory,
             witness_cols,
@@ -99,7 +95,6 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> Processor<'a, 'b, 'c, T,
             row_offset: self.row_offset,
             data: self.data,
             mutable_state: self.mutable_state,
-            identities: self.identities,
             fixed_data: self.fixed_data,
             row_factory: self.row_factory,
             witness_cols: self.witness_cols,
@@ -143,10 +138,8 @@ impl<'a, 'b, 'c, T: FieldElement, Q: QueryCallback<T>> Processor<'a, 'b, 'c, T,
     pub fn process_identity(
         &mut self,
         row_index: usize,
-        identity_index: usize,
+        identity: &'a Identity<Expression<T>>,
     ) -> Result<bool, EvalError<T>> {
-        let identity = &self.identities[identity_index];
-
         // Create row pair
         let global_row_index = self.row_offset + row_index as u64;
         let row_pair = RowPair::new(