Give access to column allocator. (#3456)

If we need to create new variables / column in the optimization
procedure, we have to make sure that their IDs do not clash and this is
what the ColumnAllocator is for.

Author: chriseth

autoprecompiles/src/lib.rs

@@ -393,15 +393,22 @@ impl<T, I> Apc<T, I> {
 }
 
 /// Allocates global poly_ids and keeps track of substitutions
-struct ColumnAllocator {
+pub 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 {
+    pub fn from_max_poly_id_of_machine(machine: &SymbolicMachine<impl FieldElement>) -> Self {
+        Self {
+            subs: Vec::new(),
+            next_poly_id: machine.main_columns().map(|c| c.id).max().unwrap_or(0) + 1,
+        }
+    }
+
+    pub fn issue_next_poly_id(&mut self) -> u64 {
         let id = self.next_poly_id;
         self.next_poly_id += 1;
         id
@@ -430,11 +437,12 @@ pub fn build<A: Adapter>(
     metrics::counter!("before_opt_interactions", &labels)
         .absolute(machine.unique_references().count() as u64);
 
-    let machine = optimizer::optimize::<A>(
+    let (machine, column_allocator) = optimizer::optimize::<A>(
         machine,
         vm_config.bus_interaction_handler,
         degree_bound,
         &vm_config.bus_map,
+        column_allocator,
     )?;
 
     // add guards to constraints that are not satisfied by zeroes

autoprecompiles/src/optimizer.rs

@@ -16,6 +16,7 @@ use powdr_number::FieldElement;
 
 use crate::constraint_optimizer::trivial_simplifications;
 use crate::range_constraint_optimizer::optimize_range_constraints;
+use crate::ColumnAllocator;
 use crate::{
     adapter::Adapter,
     constraint_optimizer::optimize_constraints,
@@ -26,12 +27,16 @@ use crate::{
     BusMap, BusType, DegreeBound, SymbolicBusInteraction, SymbolicMachine,
 };
 
+/// Optimizes a given symbolic machine and returns an equivalent, but "simpler" one.
+/// All constraints in the returned machine will respect the given degree bound.
+/// New variables may be introduced in the process.
 pub fn optimize<A: Adapter>(
     mut machine: SymbolicMachine<A::PowdrField>,
     bus_interaction_handler: A::BusInteractionHandler,
     degree_bound: DegreeBound,
     bus_map: &BusMap<A::CustomBusTypes>,
-) -> Result<SymbolicMachine<A::PowdrField>, crate::constraint_optimizer::Error> {
+    column_allocator: ColumnAllocator,
+) -> Result<(SymbolicMachine<A::PowdrField>, ColumnAllocator), crate::constraint_optimizer::Error> {
     let mut stats_logger = StatsLogger::start(&machine);
 
     if let Some(exec_bus_id) = bus_map.get_bus_id(&BusType::ExecutionBridge) {
@@ -112,7 +117,10 @@ pub fn optimize<A: Adapter>(
                 == GroupedExpression::from_number(A::PowdrField::from(pc_lookup_bus_id))),
         "Expected all PC lookups to be removed."
     );
-    Ok(constraint_system_to_symbolic_machine(constraint_system))
+    Ok((
+        constraint_system_to_symbolic_machine(constraint_system),
+        column_allocator,
+    ))
 }
 
 pub fn optimize_exec_bus<T: FieldElement>(

openvm/benches/optimizer_benchmark.rs

@@ -1,5 +1,5 @@
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
-use powdr_autoprecompiles::{optimizer::optimize, SymbolicMachine};
+use powdr_autoprecompiles::{optimizer::optimize, ColumnAllocator, SymbolicMachine};
 use powdr_number::BabyBearField;
 use powdr_openvm::{
     bus_interaction_handler::OpenVmBusInteractionHandler, bus_map::default_openvm_bus_map,
@@ -14,16 +14,21 @@ fn optimize_keccak_benchmark(c: &mut Criterion) {
     let file = std::fs::File::open("tests/keccak_apc_pre_opt.cbor").unwrap();
     let reader = std::io::BufReader::new(file);
     let machine: SymbolicMachine<BabyBearField> = serde_cbor::from_reader(reader).unwrap();
-
     group.bench_function("optimize", |b| {
         b.iter_batched(
-            || machine.clone(),
-            |machine| {
+            || {
+                (
+                    machine.clone(),
+                    ColumnAllocator::from_max_poly_id_of_machine(&machine),
+                )
+            },
+            |(machine, column_allocator)| {
                 optimize::<BabyBearOpenVmApcAdapter>(
                     black_box(machine),
                     OpenVmBusInteractionHandler::default(),
                     DEFAULT_DEGREE_BOUND,
                     &default_openvm_bus_map(),
+                    column_allocator,
                 )
                 .unwrap()
             },

openvm/tests/optimizer.rs

@@ -1,6 +1,6 @@
 use expect_test::expect;
 use powdr_autoprecompiles::optimizer::optimize;
-use powdr_autoprecompiles::SymbolicMachine;
+use powdr_autoprecompiles::{ColumnAllocator, SymbolicMachine};
 use powdr_number::BabyBearField;
 use powdr_openvm::{
     bus_interaction_handler::OpenVmBusInteractionHandler, bus_map::default_openvm_bus_map,
@@ -39,13 +39,16 @@ fn test_optimize() {
     let machine: SymbolicMachine<BabyBearField> = serde_cbor::from_reader(reader).unwrap();
     assert!(machine.derived_columns.is_empty());
 
+    let column_allocator = ColumnAllocator::from_max_poly_id_of_machine(&machine);
     let machine = optimize::<BabyBearOpenVmApcAdapter>(
         machine,
         OpenVmBusInteractionHandler::default(),
         DEFAULT_DEGREE_BOUND,
         &default_openvm_bus_map(),
+        column_allocator,
     )
-    .unwrap();
+    .unwrap()
+    .0;
 
     // This cbor file above has the `is_valid` column removed, this is why the number below
     // might be one less than in other tests.
@@ -69,14 +72,17 @@ fn test_sha256() {
     let reader = flate2::read::GzDecoder::new(file);
     let machine: SymbolicMachine<BabyBearField> = serde_cbor::from_reader(reader).unwrap();
     assert!(machine.derived_columns.is_empty());
+    let column_allocator = ColumnAllocator::from_max_poly_id_of_machine(&machine);
 
     let machine = optimize::<BabyBearOpenVmApcAdapter>(
         machine,
         OpenVmBusInteractionHandler::default(),
         DEFAULT_DEGREE_BOUND,
         &default_openvm_bus_map(),
+        column_allocator,
     )
-    .unwrap();
+    .unwrap()
+    .0;
 
     // This cbor file above has the `is_valid` column removed, this is why the number below
     // might be one less than in other tests.