Simplify FinalizableData (#2304)

A refactoring pulled out of #2281, with the following changes:
- `FinalizableData` no longer supports unfinalized rows before the
section of finalized rows. This simplifies `FinalizableData`.
- To handle wrapping, we might still night the "unfinalized" row. For
that purpose, I added `FinalizableData::get_in_progress_row`. Note that
using this means that any range constraints will be lost, but that's not
a problem in practice.

executor/src/witgen/block_processor.rs

@@ -177,7 +177,8 @@ mod tests {
             .collect();
         let data = FinalizableData::with_initial_rows_in_progress(
             &columns,
-            (0..degree).map(|i| Row::fresh(&fixed_data, RowIndex::from_degree(i, degree))),
+            (0..degree).map(|i| Row::fresh(&fixed_data, i)),
+            &fixed_data,
         );
 
         let mutable_state = MutableState::new(iter::empty(), &query_callback);

executor/src/witgen/data_structures/padded_bitvec.rs

@@ -70,11 +70,28 @@ impl PaddedBitVec {
     }
 
     pub fn get(&self, row: usize, col: u64) -> bool {
+        let (word_index, bit_mask) = self.to_word_and_bit_mask(row, col);
+        let word = &self.data[word_index];
+        (word & bit_mask) != 0
+    }
+
+    pub fn set(&mut self, row: usize, col: u64, value: bool) {
+        let (word_index, bit_mask) = self.to_word_and_bit_mask(row, col);
+        let word = &mut self.data[word_index];
+        if value {
+            *word |= bit_mask;
+        } else {
+            *word &= !bit_mask;
+        }
+    }
+
+    fn to_word_and_bit_mask(&self, row: usize, col: u64) -> (usize, u32) {
         if row >= self.rows || (row + 1 == self.rows && col >= self.bits_in_last_row as u64) {
             panic!("Out of bounds");
         }
-        let word = &self.data[row * self.words_per_row + (col / 32) as usize];
-        (word & (1 << (col % 32))) != 0
+        let word_index = row * self.words_per_row + (col / 32) as usize;
+        let bit_mask = 1 << (col % 32);
+        (word_index, bit_mask)
     }
 
     pub fn as_mut_slice(&mut self) -> &mut [u32] {

executor/src/witgen/machines/block_machine.rs

@@ -65,11 +65,8 @@ pub struct BlockMachine<'a, T: FieldElement> {
     /// The type of constraint used to connect this machine to its caller.
     connection_type: ConnectionKind,
     /// The data of the machine.
-    data: FinalizableData<T>,
+    data: FinalizableData<'a, T>,
     publics: BTreeMap<&'a str, T>,
-    /// The index of the first row that has not been finalized yet.
-    /// At all times, all rows in the range [block_size..first_in_progress_row) are finalized.
-    first_in_progress_row: usize,
     /// Cache that states the order in which to evaluate identities
     /// to make progress most quickly.
     processing_sequence_cache: ProcessingSequenceCache,
@@ -120,6 +117,7 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         let data = FinalizableData::with_initial_rows_in_progress(
             &parts.witnesses,
             (0..block_size).map(|i| Row::fresh(fixed_data, start_index + i)),
+            fixed_data,
         );
         let layout = data.layout();
         Some(BlockMachine {
@@ -133,7 +131,6 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
             connection_type: is_permutation,
             data,
             publics: Default::default(),
-            first_in_progress_row: block_size,
             multiplicity_counter: MultiplicityCounter::new(&parts.connections),
             processing_sequence_cache: ProcessingSequenceCache::new(
                 block_size,
@@ -245,7 +242,8 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
                 iter::once(self.block_size - 1)
                     .chain(0..self.block_size)
                     .chain(iter::once(0))
-                    .map(|i| self.data[i].clone()),
+                    .map(|i| self.data.get_in_progress_row(i)),
+                self.fixed_data,
             );
 
             // Instantiate a processor
@@ -279,7 +277,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
             // Replace the dummy block, discarding first and last row
             dummy_block.pop().unwrap();
             for i in (0..self.block_size).rev() {
-                self.data[i] = dummy_block.pop().unwrap();
+                self.data.set(i, dummy_block.pop().unwrap());
             }
         }
 
@@ -484,9 +482,7 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
             (self.rows() + self.block_size as DegreeType) <= self.degree,
             "Block machine is full (this should have been checked before)"
         );
-        self.data
-            .finalize_range(self.first_in_progress_row..self.data.len());
-        self.first_in_progress_row = self.data.len() + self.block_size;
+        self.data.finalize_all();
         //TODO can we properly access the last row of the dummy block?
         let data = self.data.append_new_finalized_rows(self.block_size);
 
@@ -513,6 +509,7 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         let block = FinalizableData::with_initial_rows_in_progress(
             &self.parts.witnesses,
             (0..(self.block_size + 2)).map(|i| Row::fresh(self.fixed_data, row_offset + i)),
+            self.fixed_data,
         );
         let mut processor = BlockProcessor::new(
             row_offset,
@@ -535,7 +532,7 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
     /// the last row of its previous block is merged with the one we have already.
     /// This is necessary to handle non-rectangular block machines, which already use
     /// unused cells in the previous block.
-    fn append_block(&mut self, mut new_block: FinalizableData<T>) -> Result<(), EvalError<T>> {
+    fn append_block(&mut self, mut new_block: FinalizableData<'a, T>) -> Result<(), EvalError<T>> {
         assert!(
             (self.rows() + self.block_size as DegreeType) <= self.degree,
             "Block machine is full (this should have been checked before)"
@@ -560,12 +557,8 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         // 3. Remove the last row of the previous block from data
         self.data.truncate(self.data.len() - 1);
 
-        // 4. Finalize everything so far (except the dummy block)
-        if self.data.len() > self.block_size {
-            self.data
-                .finalize_range(self.first_in_progress_row..self.data.len());
-            self.first_in_progress_row = self.data.len();
-        }
+        // 4. Finalize everything so far
+        self.data.finalize_all();
 
         // 5. Append the new block (including the merged last row of the previous block)
         self.data.extend(new_block);

executor/src/witgen/machines/dynamic_machine.rs

@@ -26,7 +26,7 @@ struct ProcessResult<'a, T: FieldElement> {
 pub struct DynamicMachine<'a, T: FieldElement> {
     fixed_data: &'a FixedData<'a, T>,
     parts: MachineParts<'a, T>,
-    data: FinalizableData<T>,
+    data: FinalizableData<'a, T>,
     publics: BTreeMap<&'a str, T>,
     latch: Option<Expression<T>>,
     name: String,
@@ -138,7 +138,7 @@ impl<'a, T: FieldElement> DynamicMachine<'a, T> {
         parts: MachineParts<'a, T>,
         latch: Option<Expression<T>>,
     ) -> Self {
-        let data = FinalizableData::new(&parts.witnesses);
+        let data = FinalizableData::new(&parts.witnesses, fixed_data);
         let multiplicity_counter = MultiplicityCounter::new(&parts.connections);
 
         Self {
@@ -194,6 +194,7 @@ impl<'a, T: FieldElement> DynamicMachine<'a, T> {
                 Row::fresh(self.fixed_data, RowIndex::from_i64(0, self.degree)),
             ]
             .into_iter(),
+            self.fixed_data,
         );
 
         // We're only interested in the first row anyway, so identities without a next reference
@@ -236,6 +237,7 @@ impl<'a, T: FieldElement> DynamicMachine<'a, T> {
         let data = FinalizableData::with_initial_rows_in_progress(
             &self.parts.witnesses,
             [first_row].into_iter(),
+            self.fixed_data,
         );
 
         let mut processor = VmProcessor::new(
@@ -268,15 +270,17 @@ impl<'a, T: FieldElement> DynamicMachine<'a, T> {
     fn fix_first_row(&mut self) {
         assert_eq!(self.data.len() as DegreeType, self.degree + 1);
 
+        let mut first_row = self.data.get_in_progress_row(0);
         let last_row = self.data.pop().unwrap();
-        if self.data[0].merge_with(&last_row).is_err() {
-            log::error!("{}", self.data[0].render("First row", false, &self.parts));
+        if first_row.merge_with(&last_row).is_err() {
+            log::error!("{}", first_row.render("First row", false, &self.parts));
             log::error!("{}", last_row.render("Last row", false, &self.parts));
             panic!(
                 "Failed to merge the first and last row of the VM '{}'",
                 self.name()
             );
         }
+        self.data.set(0, first_row);
     }
 
     #[cfg(test)]

executor/src/witgen/machines/second_stage_machine.rs

@@ -20,7 +20,7 @@ use super::LookupCell;
 pub struct SecondStageMachine<'a, T: FieldElement> {
     fixed_data: &'a FixedData<'a, T>,
     parts: MachineParts<'a, T>,
-    data: FinalizableData<T>,
+    data: FinalizableData<'a, T>,
     publics: BTreeMap<&'a str, T>,
     name: String,
     degree: DegreeType,
@@ -78,7 +78,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for SecondStageMachine<'a, T> {
 
 impl<'a, T: FieldElement> SecondStageMachine<'a, T> {
     pub fn new(name: String, fixed_data: &'a FixedData<'a, T>, parts: MachineParts<'a, T>) -> Self {
-        let data = FinalizableData::new(&parts.witnesses);
+        let data = FinalizableData::new(&parts.witnesses, fixed_data);
 
         // Only keep polynomial identities. We assume other constraints to be handled in stage 0.
         let polynomial_identities = parts
@@ -132,6 +132,7 @@ impl<'a, T: FieldElement> SecondStageMachine<'a, T> {
                 Row::fresh(self.fixed_data, RowIndex::from_i64(0, self.degree)),
             ]
             .into_iter(),
+            self.fixed_data,
         );
 
         // We're only interested in the first row anyway, so identities without a next reference
@@ -163,14 +164,15 @@ impl<'a, T: FieldElement> SecondStageMachine<'a, T> {
         &mut self,
         first_row: Row<T>,
         mutable_state: &MutableState<'a, T, Q>,
-    ) -> FinalizableData<T> {
+    ) -> FinalizableData<'a, T> {
         log::trace!(
             "Running Second-Stage Machine with the following initial values in the first row:\n{}",
             first_row.render_values(false, &self.parts)
         );
         let data = FinalizableData::with_initial_rows_in_progress(
             &self.parts.witnesses,
             [first_row].into_iter(),
+            self.fixed_data,
         );
 
         let mut processor = VmProcessor::new(

executor/src/witgen/processor.rs

@@ -29,17 +29,17 @@ pub type Left<'a, T> = Vec<AffineExpression<AlgebraicVariable<'a>, T>>;
 /// The data mutated by the processor
 pub(crate) struct SolverState<'a, T: FieldElement> {
     /// The block of trace cells
-    pub block: FinalizableData<T>,
+    pub block: FinalizableData<'a, T>,
     /// The values of publics
     pub publics: BTreeMap<&'a str, T>,
 }
 
 impl<'a, T: FieldElement> SolverState<'a, T> {
-    pub fn new(block: FinalizableData<T>, publics: BTreeMap<&'a str, T>) -> Self {
+    pub fn new(block: FinalizableData<'a, T>, publics: BTreeMap<&'a str, T>) -> Self {
         Self { block, publics }
     }
 
-    pub fn without_publics(block: FinalizableData<T>) -> Self {
+    pub fn without_publics(block: FinalizableData<'a, T>) -> Self {
         Self {
             block,
             publics: BTreeMap::new(),
@@ -140,7 +140,7 @@ pub struct Processor<'a, 'c, T: FieldElement, Q: QueryCallback<T>> {
     /// The global index of the first row of [Processor::data].
     row_offset: RowIndex,
     /// The rows that are being processed.
-    data: FinalizableData<T>,
+    data: FinalizableData<'a, T>,
     /// The values of the publics
     publics: BTreeMap<&'a str, T>,
     /// The mutable state
@@ -575,12 +575,12 @@ Known values in current row (local: {row_index}, global {global_row_index}):
         self.data.len()
     }
 
-    pub fn finalize_range(&mut self, range: std::ops::Range<usize>) {
+    pub fn finalize_until(&mut self, end: usize) {
         assert!(
             self.copy_constraints.is_empty(),
             "Machines with copy constraints should not be finalized while being processed."
         );
-        self.data.finalize_range(range);
+        self.data.finalize_until(end);
     }
 
     pub fn row(&self, i: usize) -> &Row<T> {

executor/src/witgen/rows.rs

@@ -183,6 +183,10 @@ pub struct Row<T: FieldElement> {
 }
 
 impl<T: FieldElement> Row<T> {
+    pub fn columns(&self) -> impl Iterator<Item = PolyID> {
+        self.values.keys()
+    }
+
     pub fn value_or_zero(&self, poly_id: &PolyID) -> T {
         self.values[poly_id].unwrap_or_zero()
     }
@@ -249,6 +253,10 @@ impl<T: FieldElement> Row<T> {
         self.values[poly_id].is_known()
     }
 
+    pub fn set_cell_known(&mut self, poly_id: &PolyID, value: T) {
+        self.values[poly_id] = CellValue::Known(value);
+    }
+
     pub fn set_cell_unknown(&mut self, poly_id: &PolyID) {
         self.values[poly_id] = CellValue::Unknown;
     }
@@ -260,19 +268,24 @@ impl<T: FieldElement> Row<T> {
 
 impl<T: FieldElement> Row<T> {
     /// Creates a "fresh" row, i.e., one that is empty but initialized with the global range constraints.
-    pub fn fresh(fixed_data: &FixedData<'_, T>, row: RowIndex) -> Row<T> {
+    pub fn fresh<E>(
+        fixed_data: &FixedData<'_, T>,
+        row: impl TryInto<DegreeType, Error = E>,
+    ) -> Row<T>
+    where
+        E: std::fmt::Debug,
+    {
         // TODO this instance could be computed exactly once (per column set) and then cloned.
         // TODO and we could copy in the external witnesses later on
         // TODO we should really only have a subset of the columns.
+        let row = row.try_into().unwrap();
         let values = WitnessColumnMap::from(
             fixed_data
                 .global_range_constraints()
                 .witness_constraints
                 .iter()
                 .map(|(poly_id, rc)| {
-                    if let Some(external_witness) =
-                        fixed_data.external_witness(row.into(), &poly_id)
-                    {
+                    if let Some(external_witness) = fixed_data.external_witness(row, &poly_id) {
                         CellValue::Known(external_witness)
                     } else if let Some(rc) = rc {
                         CellValue::RangeConstraint(rc.clone())

executor/src/witgen/vm_processor.rs

@@ -156,7 +156,6 @@ impl<'a, 'c, T: FieldElement, Q: QueryCallback<T>> VmProcessor<'a, 'c, T, Q> {
         } else {
             log::Level::Debug
         };
-        let mut finalize_start = 1;
 
         for row_index in 0.. {
             // The total number of rows to run for. Note that `self.degree` might change during
@@ -174,8 +173,7 @@ impl<'a, 'c, T: FieldElement, Q: QueryCallback<T>> VmProcessor<'a, 'c, T, Q> {
                 // Periodically make sure most rows are finalized.
                 // Row 0 and the last MAX_PERIOD rows might be needed later, so they are not finalized.
                 let finalize_end = row_index as usize - MAX_PERIOD;
-                self.processor.finalize_range(finalize_start..finalize_end);
-                finalize_start = finalize_end;
+                self.processor.finalize_until(finalize_end);
             }
 
             if row_index >= rows_to_run - 2 {