Introduce column allocator (#3387)

Alternative to #3385 which solves it in apc generation already.
The original bug is that when introducing a derived column (right now
only `is_valid`), the poly_id is chosen to be the max of the current
poly_ids in the machine plus one. But this poly_id may already be
pointed to by an original column in the substitutions, so in witgen it
looks like an apc column is both computed from an original column, and a
derived column.

This PR introduces a column allocator which issues new poly_ids that are
guaranteed not to be already used.
When returning the substitutions, we only keep the substitutions which
actually point to an apc column.

Author: Schaeff

autoprecompiles/src/lib.rs

@@ -12,6 +12,7 @@ use powdr_expression::{
     visitors::Children, AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicUnaryOperation,
 };
 use serde::{Deserialize, Serialize};
+use std::collections::BTreeSet;
 use std::fmt::Display;
 use std::io::BufWriter;
 use std::iter::once;
@@ -316,15 +317,26 @@ pub trait InstructionHandler {
     fn is_branching(&self, instruction: &Self::Instruction) -> bool;
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Substitution {
+    /// The index of the original column in the original air
+    original_poly_index: usize,
+    /// The `poly_id` of the target column in the APC air
+    apc_poly_id: u64,
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Apc<T, I> {
+    /// The basic block this APC is based on
     pub block: BasicBlock<I>,
+    /// The symbolic machine for this APC
     pub machine: SymbolicMachine<T>,
-    pub subs: Vec<Vec<u64>>,
+    /// For each original air, the substitutions from original columns to APC columns
+    pub subs: Vec<Vec<Substitution>>,
 }
 
 impl<T, I> Apc<T, I> {
-    pub fn subs(&self) -> &[Vec<u64>] {
+    pub fn subs(&self) -> &[Vec<Substitution>] {
         &self.subs
     }
 
@@ -341,6 +353,59 @@ impl<T, I> Apc<T, I> {
     pub fn instructions(&self) -> &[I] {
         &self.block.statements
     }
+
+    /// Create a new APC based on the given basic block, symbolic machine and column allocator
+    /// The column allocator only issues the subs which are actually used in the machine
+    fn new(
+        block: BasicBlock<I>,
+        machine: SymbolicMachine<T>,
+        column_allocator: ColumnAllocator,
+    ) -> Self {
+        // Get all poly_ids in the machine
+        let all_references = machine
+            .unique_references()
+            .map(|r| r.id)
+            .collect::<BTreeSet<_>>();
+        // Only keep substitutions from the column allocator if the target poly_id is used in the machine
+        let subs = column_allocator
+            .subs
+            .into_iter()
+            .map(|subs| {
+                subs.into_iter()
+                    .enumerate()
+                    .filter_map(|(original_poly_index, apc_poly_id)| {
+                        all_references
+                            .contains(&apc_poly_id)
+                            .then_some(Substitution {
+                                original_poly_index,
+                                apc_poly_id,
+                            })
+                    })
+                    .collect_vec()
+            })
+            .collect();
+        Self {
+            block,
+            machine,
+            subs,
+        }
+    }
+}
+
+/// Allocates global poly_ids and keeps track of substitutions
+struct ColumnAllocator {
+    /// For each original air, for each original column index, the associated poly_id in the APC air
+    subs: Vec<Vec<u64>>,
+    /// The next poly_id to issue
+    next_poly_id: u64,
+}
+
+impl ColumnAllocator {
+    fn issue_next_poly_id(&mut self) -> u64 {
+        let id = self.next_poly_id;
+        self.next_poly_id += 1;
+        id
+    }
 }
 
 pub fn build<A: Adapter>(
@@ -351,7 +416,7 @@ pub fn build<A: Adapter>(
 ) -> Result<AdapterApc<A>, crate::constraint_optimizer::Error> {
     let start = std::time::Instant::now();
 
-    let (machine, subs) = statements_to_symbolic_machine::<A>(
+    let (machine, column_allocator) = statements_to_symbolic_machine::<A>(
         &block,
         vm_config.instruction_handler,
         &vm_config.bus_map,
@@ -373,7 +438,7 @@ pub fn build<A: Adapter>(
     )?;
 
     // add guards to constraints that are not satisfied by zeroes
-    let machine = add_guards(machine);
+    let (machine, column_allocator) = add_guards(machine, column_allocator);
 
     metrics::counter!("after_opt_cols", &labels)
         .absolute(machine.unique_references().count() as u64);
@@ -384,11 +449,7 @@ pub fn build<A: Adapter>(
 
     let machine = convert_machine_field_type(machine, &A::into_field);
 
-    let apc = Apc {
-        block,
-        machine,
-        subs,
-    };
+    let apc = Apc::new(block, machine, column_allocator);
 
     if let Some(path) = apc_candidates_dir_path {
         let ser_path = path
@@ -447,12 +508,15 @@ fn add_guards_constraint<T: FieldElement>(
 }
 
 /// Adds an `is_valid` guard to all constraints and bus interactions, if needed.
-fn add_guards<T: FieldElement>(mut machine: SymbolicMachine<T>) -> SymbolicMachine<T> {
+fn add_guards<T: FieldElement>(
+    mut machine: SymbolicMachine<T>,
+    mut column_allocator: ColumnAllocator,
+) -> (SymbolicMachine<T>, ColumnAllocator) {
     let pre_degree = machine.degree();
 
     let is_valid_ref = AlgebraicReference {
         name: Arc::new("is_valid".to_string()),
-        id: machine.unique_references().map(|c| c.id).max().unwrap() + 1,
+        id: column_allocator.issue_next_poly_id(),
     };
     let is_valid = AlgebraicExpression::Reference(is_valid_ref.clone());
 
@@ -493,5 +557,5 @@ fn add_guards<T: FieldElement>(mut machine: SymbolicMachine<T>) -> SymbolicMachi
     // so it may increase the degree of the machine.
     machine.constraints.push(powdr::make_bool(is_valid).into());
 
-    machine
+    (machine, column_allocator)
 }

autoprecompiles/src/symbolic_machine_generator.rs

@@ -7,8 +7,8 @@ use crate::{
     adapter::Adapter,
     blocks::{BasicBlock, Instruction},
     expression::AlgebraicExpression,
-    powdr, BusMap, BusType, InstructionHandler, SymbolicBusInteraction, SymbolicConstraint,
-    SymbolicMachine,
+    powdr, BusMap, BusType, ColumnAllocator, InstructionHandler, SymbolicBusInteraction,
+    SymbolicConstraint, SymbolicMachine,
 };
 
 /// Converts the field type of a symbolic machine.
@@ -102,11 +102,11 @@ fn convert_expression<T, U>(
 /// that contains, for each instruction in the basic block,
 /// a mapping from local column IDs to global column IDs
 /// (in the form of a vector).
-pub fn statements_to_symbolic_machine<A: Adapter>(
+pub(crate) fn statements_to_symbolic_machine<A: Adapter>(
     block: &BasicBlock<A::Instruction>,
     instruction_handler: &A::InstructionHandler,
     bus_map: &BusMap<A::CustomBusTypes>,
-) -> (SymbolicMachine<A::PowdrField>, Vec<Vec<u64>>) {
+) -> (SymbolicMachine<A::PowdrField>, ColumnAllocator) {
     let mut constraints: Vec<SymbolicConstraint<_>> = Vec::new();
     let mut bus_interactions: Vec<SymbolicBusInteraction<_>> = Vec::new();
     let mut col_subs: Vec<Vec<u64>> = Vec::new();
@@ -184,7 +184,10 @@ pub fn statements_to_symbolic_machine<A: Adapter>(
             bus_interactions,
             derived_columns: vec![],
         },
-        col_subs,
+        ColumnAllocator {
+            subs: col_subs,
+            next_poly_id: global_idx,
+        },
     )
 }
 

autoprecompiles/src/trace_handler.rs

@@ -87,15 +87,13 @@ where
         .iter()
         .map(|subs| {
             subs.iter()
-                .enumerate()
-                .filter_map(|(dummy_index, poly_id)| {
-                    // Check if this dummy column is present in the final apc row
-                    apc_poly_id_to_index
-                        .get(poly_id)
-                        // If it is, map the dummy index to the apc index
-                        .map(|apc_index| (dummy_index, *apc_index))
+                .map(|substitution| {
+                    (
+                        substitution.original_poly_index,
+                        apc_poly_id_to_index[&substitution.apc_poly_id],
+                    )
                 })
-                .collect()
+                .collect_vec()
         })
         .collect();