Fully implement irregular block shapes

executor/src/witgen/data_structures/finalizable_data.rs

@@ -93,6 +93,23 @@ impl<T: FieldElement> CompactData<T> {
         }
     }
 
+    pub fn set(&mut self, row: usize, new_row: Row<T>) {
+        let idx = row * self.column_count;
+        for (i, col_id) in
+            (self.first_column_id..(self.first_column_id + self.column_count as u64)).enumerate()
+        {
+            if let Some(v) = new_row.value(&PolyID {
+                id: col_id,
+                ptype: PolynomialType::Committed,
+            }) {
+                self.data[idx + i] = v;
+                self.known_cells.set(row, i as u64, true);
+            } else {
+                self.known_cells.set(row, i as u64, false);
+            }
+        }
+    }
+
     pub fn append_new_rows(&mut self, count: usize) {
         self.data
             .resize(self.data.len() + count * self.column_count, T::zero());
@@ -394,6 +411,7 @@ impl<T: FieldElement> FinalizableData<T> {
     pub fn append_new_finalized_rows(&mut self, count: usize) -> CompactDataRef<'_, T> {
         assert!(self.post_finalized_data.is_empty());
         let row_zero = self.finalized_data.len();
+        assert!(row_zero > 0);
         self.finalized_data.append_new_rows(count);
         CompactDataRef::new(&mut self.finalized_data, row_zero)
     }
@@ -450,6 +468,38 @@ impl<T: FieldElement> FinalizableData<T> {
             },
         )
     }
+
+    pub fn get_in_progress_row(&self, i: usize, make_fresh_row: impl Fn() -> Row<T>) -> Row<T> {
+        match self.location_of_row(i) {
+            Location::PreFinalized(local) => self.pre_finalized_data[local].clone(),
+            Location::Finalized(local) => {
+                let layout = self.layout();
+                let mut row = make_fresh_row();
+                for column in row.columns() {
+                    if column.id >= layout.first_column_id
+                        && column.id < layout.first_column_id + layout.column_count as u64
+                    {
+                        let (value, known) = self.finalized_data.get(local, column.id);
+                        if known {
+                            row.set_cell_known(&column, value);
+                        }
+                    }
+                }
+                row
+            }
+            Location::PostFinalized(local) => self.post_finalized_data[local].clone(),
+        }
+    }
+
+    pub fn set(&mut self, i: usize, row: Row<T>) {
+        match self.location_of_row(i) {
+            Location::PreFinalized(local) => self.pre_finalized_data[local] = row,
+            Location::Finalized(local) => {
+                self.finalized_data.set(local, row);
+            }
+            Location::PostFinalized(local) => self.post_finalized_data[local] = row,
+        }
+    }
 }
 
 impl<T: FieldElement> Index<usize> for FinalizableData<T> {

executor/src/witgen/data_structures/padded_bitvec.rs

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

executor/src/witgen/jit/block_machine_processor.rs

@@ -64,6 +64,17 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> {
         // In the latch row, set the RHS selector to 1.
         witgen.assign_constant(&connection_rhs.selector, self.latch_row as i32, T::one());
 
+        // Set all other selectors to 0 in the latch row.
+        for other_connection in self.machine_parts.connections.values() {
+            if other_connection.right.selector != connection_rhs.selector {
+                witgen.assign_constant(
+                    &other_connection.right.selector,
+                    self.latch_row as i32,
+                    T::zero(),
+                );
+            }
+        }
+
         // For each argument, connect the expression on the RHS with the formal parameter.
         for (index, expr) in connection_rhs.expressions.iter().enumerate() {
             witgen.assign_variable(expr, self.latch_row as i32, Variable::Param(index));

executor/src/witgen/jit/function_cache.rs

@@ -5,7 +5,6 @@ use powdr_number::{FieldElement, KnownField};
 
 use crate::witgen::{
     data_structures::finalizable_data::{ColumnLayout, CompactDataRef},
-    jit::effect::Effect,
     machines::{LookupCell, MachineParts},
     EvalError, FixedData, MutableState, QueryCallback,
 };
@@ -29,7 +28,6 @@ pub struct FunctionCache<'a, T: FieldElement> {
     /// but failed.
     witgen_functions: HashMap<CacheKey, Option<WitgenFunction<T>>>,
     column_layout: ColumnLayout,
-    block_size: usize,
 }
 
 impl<'a, T: FieldElement> FunctionCache<'a, T> {
@@ -47,7 +45,6 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
             processor,
             column_layout: metadata,
             witgen_functions: HashMap::new(),
-            block_size,
         }
     }
 
@@ -96,25 +93,6 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
         self.processor
             .generate_code(mutable_state, cache_key.identity_id, &cache_key.known_args)
             .ok()
-            .and_then(|code| {
-                // TODO: Remove this once BlockMachine passes the right amount of context for machines with
-                // non-rectangular block shapes.
-                let is_rectangular = code
-                    .iter()
-                    .filter_map(|effect| match effect {
-                        Effect::Assignment(v, _) => Some(v),
-                        _ => None,
-                    })
-                    .filter_map(|assigned_variable| match assigned_variable {
-                        Variable::Cell(cell) => Some(cell.row_offset),
-                        _ => None,
-                    })
-                    .all(|row_offset| row_offset >= 0 && row_offset < self.block_size as i32);
-                if !is_rectangular {
-                    log::debug!("Filtering out code for non-rectangular block shape");
-                }
-                is_rectangular.then_some(code)
-            })
             .map(|code| {
                 log::trace!("Generated code ({} steps)", code.len());
                 let known_inputs = cache_key

executor/src/witgen/machines/block_machine.rs

@@ -67,9 +67,6 @@ pub struct BlockMachine<'a, T: FieldElement> {
     /// The data of the machine.
     data: FinalizableData<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,
@@ -133,7 +130,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 +241,14 @@ 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, || {
+                            Row::fresh(
+                                self.fixed_data,
+                                RowIndex::from_degree(i as u64, self.degree),
+                            )
+                        })
+                    }),
             );
 
             // Instantiate a processor
@@ -279,7 +282,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 +487,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_range(0..self.data.len());
         //TODO can we properly access the last row of the dummy block?
         let data = self.data.append_new_finalized_rows(self.block_size);
 
@@ -560,11 +561,9 @@ 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)
+        // 4. Finalize everything so far
         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();
+            self.data.finalize_range(0..self.data.len());
         }
 
         // 5. Append the new block (including the merged last row of the previous block)

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;
     }