From 23002ae725ea3d7d784eb48fc925503ed7947244 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 19 May 2025 13:19:17 -0300
Subject: [PATCH 01/39] always add guards

---
 autoprecompiles/src/lib.rs | 28 ++--------------------------
 1 file changed, 2 insertions(+), 26 deletions(-)

diff --git a/autoprecompiles/src/lib.rs b/autoprecompiles/src/lib.rs
index f08c75310c..aa68e6b7e5 100644
--- a/autoprecompiles/src/lib.rs
+++ b/autoprecompiles/src/lib.rs
@@ -410,14 +410,7 @@ pub fn add_guards<T: FieldElement>(mut machine: SymbolicMachine<T>) -> SymbolicM
     });
 
     for c in &mut machine.constraints {
-        let mut zeroed_expr = c.expr.clone();
-        powdr::make_refs_zero(&mut zeroed_expr);
-        let zeroed_expr = simplify_expression(zeroed_expr);
-        if !powdr::is_zero(&zeroed_expr) {
-            c.expr = is_valid.clone() * c.expr.clone();
-            // TODO this would not have to be cloned if we had *=
-            //c.expr *= guard.clone();
-        }
+        c.expr = is_valid.clone() * c.expr.clone();
     }
 
     let [execution_bus_receive, execution_bus_send] = machine
@@ -447,7 +440,6 @@ pub fn add_guards<T: FieldElement>(mut machine: SymbolicMachine<T>) -> SymbolicM
         .unwrap();
     program_bus_send.mult = is_valid.clone();
 
-    let mut is_valid_mults: Vec<SymbolicConstraint<T>> = Vec::new();
     for b in &mut machine.bus_interactions {
         match b.id {
             EXECUTION_BUS_ID => {
@@ -457,27 +449,11 @@ pub fn add_guards<T: FieldElement>(mut machine: SymbolicMachine<T>) -> SymbolicM
                 // already handled
             }
             _ => {
-                // We check the value of the multiplicity when all variables are set to zero
-                let mut zeroed_expr = b.mult.clone();
-                powdr::make_refs_zero(&mut zeroed_expr);
-                let zeroed_expr = simplify_expression(zeroed_expr);
-                if !powdr::is_zero(&zeroed_expr) {
-                    // if it's not zero, then we guard the multiplicity by `is_valid`
-                    b.mult = is_valid.clone() * b.mult.clone();
-                    // TODO this would not have to be cloned if we had *=
-                    //c.expr *= guard.clone();
-                } else {
-                    // if it's zero, then we do not have to change the multiplicity, but we need to force it to be zero on non-valid rows with a constraint
-                    let one = AlgebraicExpression::Number(1u64.into());
-                    let e = ((one - is_valid.clone()) * b.mult.clone()).into();
-                    is_valid_mults.push(e);
-                }
+                b.mult = is_valid.clone() * b.mult.clone();
             }
         }
     }
 
-    machine.constraints.extend(is_valid_mults);
-
     machine
 }
 

From b05703ddc9b3c3a1748d01d1d05f3c4c35b3998f Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 19 May 2025 13:43:45 -0300
Subject: [PATCH 02/39] apc chip stacking

---
 openvm/src/customize_exe.rs        | 204 +++++++++++++-
 openvm/src/powdr_extension/chip.rs | 430 ++++++++++++++++-------------
 openvm/src/powdr_extension/mod.rs  |   2 +-
 openvm/src/powdr_extension/vm.rs   |  21 +-
 4 files changed, 455 insertions(+), 202 deletions(-)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 34249c554b..4ac797e436 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -10,7 +10,8 @@ use openvm_rv32im_transpiler::{Rv32HintStoreOpcode, Rv32LoadStoreOpcode};
 use openvm_sdk::config::SdkVmConfig;
 use openvm_sha256_transpiler::Rv32Sha256Opcode;
 use openvm_stark_backend::{interaction::SymbolicInteraction, p3_field::PrimeField32};
-use powdr::ast::analyzed::AlgebraicExpression;
+use powdr::ast::analyzed::{AlgebraicExpression, AlgebraicReference, PolyID, PolynomialType};
+use powdr::ast::parsed::visitor::ExpressionVisitable;
 use powdr::{FieldElement, LargeInt};
 use powdr_autoprecompiles::powdr::UniqueColumns;
 use powdr_autoprecompiles::{
@@ -21,7 +22,7 @@ use powdr_autoprecompiles::{
 use crate::bus_interaction_handler::OpenVmBusInteractionHandler;
 use crate::instruction_formatter::openvm_instruction_formatter;
 use crate::{
-    powdr_extension::{OriginalInstruction, PowdrExtension, PowdrOpcode, PowdrPrecompile},
+    powdr_extension::{OriginalInstruction, PowdrExtension, PowdrOpcode, PowdrPrecompile, PowdrStackedPrecompile},
     utils::symbolic_to_algebraic,
 };
 
@@ -248,9 +249,208 @@ pub fn customize<F: PrimeField32>(
         ));
     }
 
+    let extensions = air_stacking(extensions);
+
     (exe, PowdrExtension::new(extensions, base_config))
 }
 
+/// make PolyIDs start from zero, sequential and compact
+fn compact_ids<F: PrimeField32>(pcp: &mut PowdrPrecompile<F>) {
+    let mut curr_id = 0;
+    let mut id_map: HashMap<u64, u64> = Default::default();
+    let mut new_poly_id = |old_id: u64| {
+        *id_map.entry(old_id).or_insert_with(|| {
+            let id = curr_id;
+            // println!("remapping poly id {} to {}", old_id, curr_id);
+            curr_id += 1;
+            id
+        })
+    };
+    pcp.machine.pre_visit_expressions_mut(&mut |expr| {
+        if let AlgebraicExpression::Reference(r) = expr {
+            if r.poly_id.ptype == PolynomialType::Committed {
+                r.poly_id.id = new_poly_id(r.poly_id.id);
+            }
+        }
+    });
+    pcp.original_instructions.iter_mut().for_each(|instr| {
+        instr.subs.iter_mut().for_each(|sub| {
+            *sub = new_poly_id(*sub);
+        });
+    });
+
+    pcp.is_valid_column.id.id = new_poly_id(pcp.is_valid_column.id.id);
+}
+
+fn merge_bus_interactions_simple<F: PrimeField32>(interactions: Vec<Vec<SymbolicBusInteraction<F>>>) -> Vec<SymbolicBusInteraction<F>> {
+    interactions.into_iter().flatten().collect_vec()
+}
+
+fn merge_bus_interactions<F: PrimeField32>(interactions_by_machine: Vec<Vec<SymbolicBusInteraction<F>>>) -> Vec<SymbolicBusInteraction<F>> {
+    let mut interactions_by_bus: HashMap<u64, Vec<Vec<SymbolicBusInteraction<F>>>> = Default::default();
+
+    for interactions in interactions_by_machine {
+        let interactions_by_bus_this_machine = interactions.into_iter().into_group_map_by(|interaction| interaction.id);
+        for (k,v) in interactions_by_bus_this_machine {
+            let e = interactions_by_bus.entry(k).or_default();
+            e.push(v);
+        }
+    }
+    let mut to_merge = vec![];
+    for interactions_per_machine in interactions_by_bus.into_values() {
+        // this is doing a "zip longest" among the interactions from each machine
+        // (for a given bus), and saving the elements that were picked
+        // together to be merged later
+        let mut idx = 0;
+        loop {
+            let mut items = vec![];
+            for machine_interactions in interactions_per_machine.iter() {
+                if let Some(item) = machine_interactions.get(idx) {
+                    items.push(item.clone());
+                }
+            }
+            if items.is_empty() {
+                break;
+            }
+            idx += 1;
+            to_merge.push(items);
+        }
+    }
+
+    to_merge.into_iter().map(|interactions| {
+        assert_eq!(interactions.iter().map(|i| &i.kind).unique().count(), 1);
+
+        // add multiplicities together
+        let mut mult = interactions[0].mult.clone();
+        for mult2 in interactions.iter().skip(1).map(|i| i.mult.clone()) {
+            mult = mult + mult2;
+        }
+
+        // add args together
+        let mut args = interactions[0].args.clone();
+        for args2 in interactions.iter().skip(1).map(|i| i.args.clone()) {
+            args = args.into_iter().zip_eq(args2).map(|(a1, a2)| a1 + a2).collect();
+        }
+
+        SymbolicBusInteraction {
+            kind: interactions[0].kind.clone(),
+            id: interactions[0].id,
+            mult,
+            args,
+        }
+    }).collect_vec()
+}
+
+// perform air stacking of multiple precompiles into stacked precompiles.
+fn air_stacking<F: PrimeField32>(
+    mut extensions: Vec<PowdrPrecompile<F>>,
+) -> Vec<PowdrStackedPrecompile<F>> {
+    assert!(!extensions.is_empty());
+    // TODO: for now group every precompile into a single stacked one. Next,
+    // group them by some heuristic (i.e., by powers of number of witness
+    // columns)
+
+    extensions.iter_mut().for_each(compact_ids);
+
+    // take the max id in all pcps and add 1.
+    let is_valid_start = 1 + extensions.iter()
+        .flat_map(|pcp| pcp.original_instructions.iter().flat_map(|instr| instr.subs.iter()))
+        .max()
+        .unwrap();
+
+    let mut stacked_constraints = vec![];
+    let mut interactions_by_machine = vec![];
+
+    println!("stacking {} precompiles", extensions.len());
+
+    let mut is_valid_sum: Option<AlgebraicExpression<F>> = None;
+
+    for (idx, pcp) in extensions.iter_mut().enumerate() {
+        // is_valid columns cannot be shared between precompiles. Here we do
+        // their remapping into exclusive columns.
+        let is_valid_new_id = is_valid_start + idx as u64;
+
+        println!("pcp width: {}", pcp.machine.unique_columns().count());
+
+        println!(
+            "\tIS_VALID_IDX: {} ===> {}",
+            pcp.is_valid_column.id.id, is_valid_new_id
+        );
+
+        // remap is_valid column in constraints and interactions
+        let mut remapped = pcp.machine.clone();
+        remapped.pre_visit_expressions_mut(&mut |expr| {
+            if let AlgebraicExpression::Reference(r) = expr {
+                assert!(r.poly_id.id <= is_valid_start);
+                if r.poly_id == pcp.is_valid_column.id {
+                    // we assume each pcp to have a specific column named "is_valid"
+                    assert!(r.name == "is_valid");
+                    r.poly_id.id = is_valid_new_id;
+                    r.name = format!("is_valid_{}", pcp.opcode.global_opcode().as_usize());
+                } else {
+                    // we need to rename columns here because `unique_columns` relies on both `Column::name` and `Column::id`
+                    r.name = format!("col_{}", r.poly_id.id);
+                }
+            }
+        });
+
+        // set the is valid column in the original precompile
+        pcp.is_valid_column.id.id = is_valid_new_id;
+
+        let is_valid = AlgebraicExpression::Reference(AlgebraicReference {
+            name: format!("is_valid_{}", pcp.opcode.global_opcode().as_usize()),
+            poly_id: PolyID {
+                id: is_valid_new_id,
+                ptype: PolynomialType::Committed,
+            },
+            next: false,
+        });
+
+        // guard interaction payloads so they can be merged later
+        remapped.bus_interactions.iter_mut().for_each(|interaction| {
+            interaction.args.iter_mut().for_each(|arg| {
+                *arg = arg.clone() * is_valid.clone();
+            });
+        });
+
+
+        is_valid_sum = is_valid_sum
+            .map(|sum| sum + is_valid.clone())
+            .or_else(|| Some(is_valid.clone()));
+
+        stacked_constraints.extend(remapped.constraints);
+        interactions_by_machine.push(remapped.bus_interactions);
+    }
+
+    // enforce only one is_valid is active
+    let one = AlgebraicExpression::Number(F::ONE);
+    let one_hot_is_valid = (one - is_valid_sum.clone().unwrap()) * is_valid_sum.unwrap();
+
+    stacked_constraints.push(one_hot_is_valid.into());
+
+    println!("interaction count before: {}", interactions_by_machine.iter().flatten().count());
+    // let stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
+    let stacked_interactions = merge_bus_interactions(interactions_by_machine);
+
+    let machine = SymbolicMachine {
+        constraints: stacked_constraints,
+        bus_interactions: stacked_interactions,
+    };
+
+    println!("interaction count: {}", machine.bus_interactions.len());
+
+    println!("stacked width: {}", machine.unique_columns().count());
+
+    vec![PowdrStackedPrecompile {
+        precompiles: extensions
+            .into_iter()
+            .map(|p| (p.opcode.clone(), p))
+            .collect(),
+        machine,
+    }]
+}
+
+
 // TODO collect properly from opcode enums
 const BRANCH_OPCODES: [usize; 9] = [
     0x220, 0x221, 0x225, 0x226, 0x227, 0x228, 0x230, 0x231, 0x235,
diff --git a/openvm/src/powdr_extension/chip.rs b/openvm/src/powdr_extension/chip.rs
index 8087145330..e70826c9c1 100644
--- a/openvm/src/powdr_extension/chip.rs
+++ b/openvm/src/powdr_extension/chip.rs
@@ -10,7 +10,7 @@ use crate::utils::algebraic_to_symbolic;
 use super::{
     opcode::PowdrOpcode,
     vm::{OriginalInstruction, SdkVmInventory},
-    PowdrPrecompile,
+    PowdrStackedPrecompile,
 };
 use itertools::Itertools;
 use openvm_circuit::{arch::VmConfig, system::memory::MemoryController};
@@ -58,9 +58,8 @@ use serde::{Deserialize, Serialize};
 
 pub struct PowdrChip<F: PrimeField32> {
     pub name: String,
-    pub opcode: PowdrOpcode,
-    /// An "executor" for this chip, based on the original instructions in the basic block
-    pub executor: PowdrExecutor<F>,
+    /// An "executor" for each precompile stacked in this chip, by opcode.
+    pub executors: BTreeMap<usize, PowdrExecutor<F>>,
     pub air: Arc<PowdrAir<F>>,
     pub periphery: SharedChips,
 }
@@ -277,41 +276,55 @@ impl SharedChips {
 
 impl<F: PrimeField32> PowdrChip<F> {
     pub(crate) fn new(
-        precompile: PowdrPrecompile<F>,
+        precompile: PowdrStackedPrecompile<F>,
         memory: Arc<Mutex<OfflineMemory<F>>>,
         base_config: SdkVmConfig,
         periphery: SharedChips,
     ) -> Self {
         let air: PowdrAir<F> = PowdrAir::new(precompile.machine);
-        let original_airs = precompile
-            .original_airs
+        let name = format!(
+            "StackedPrecompile_{}",
+            precompile
+                .precompiles
+                .keys()
+                .map(|o| o.global_opcode())
+                .join("_")
+        );
+
+        let executors = precompile
+            .precompiles
             .into_iter()
-            .map(|(k, v)| (k, v.into()))
+            .map(|(opcode, pcp)| {
+                let original_airs = pcp
+                    .original_airs
+                    .into_iter()
+                    .map(|(k, v)| (k, v.into()))
+                    .collect();
+                let executor = PowdrExecutor::new(
+                    pcp.original_instructions,
+                    original_airs,
+                    pcp.is_valid_column,
+                    memory.clone(),
+                    &periphery.range_checker,
+                    base_config.clone(),
+                );
+                (opcode.global_opcode().as_usize(), executor)
+            })
             .collect();
-        let executor = PowdrExecutor::new(
-            precompile.original_instructions,
-            original_airs,
-            precompile.is_valid_column,
-            memory,
-            &periphery.range_checker,
-            base_config,
-        );
-        let name = precompile.name;
-        let opcode = precompile.opcode;
 
         Self {
+            // TODO: proper name
             name,
-            opcode,
             air: Arc::new(air),
-            executor,
+            executors,
             periphery,
         }
     }
 
-    /// Returns the index of the is_valid of this air.
-    fn get_is_valid_index(&self) -> usize {
-        self.air.column_index_by_poly_id[&self.executor.is_valid_poly_id]
-    }
+    // /// Returns the index of the is_valid of this air.
+    // fn get_is_valid_index(&self) -> usize {
+    //     self.air.column_index_by_poly_id[&self.executor.is_valid_poly_id]
+    // }
 }
 
 impl<F: PrimeField32> InstructionExecutor<F> for PowdrChip<F> {
@@ -322,9 +335,12 @@ impl<F: PrimeField32> InstructionExecutor<F> for PowdrChip<F> {
         from_state: ExecutionState<u32>,
     ) -> ExecutionResult<ExecutionState<u32>> {
         let &Instruction { opcode, .. } = instruction;
-        assert_eq!(opcode.as_usize(), self.opcode.global_opcode().as_usize());
 
-        let execution_state = self.executor.execute(memory, from_state)?;
+        let execution_state = self
+            .executors
+            .get_mut(&opcode.as_usize())
+            .expect("invalid opcode for stacked chip")
+            .execute(memory, from_state)?;
 
         Ok(execution_state)
     }
@@ -336,10 +352,14 @@ impl<F: PrimeField32> InstructionExecutor<F> for PowdrChip<F> {
 
 impl<F: PrimeField32> ChipUsageGetter for PowdrChip<F> {
     fn air_name(&self) -> String {
-        format!("powdr_air_for_opcode_{}", self.opcode.global_opcode()).to_string()
+        format!("powdr_air_for_opcodes_{}", self.executors.keys().join("_"))
     }
+
     fn current_trace_height(&self) -> usize {
-        self.executor.current_trace_height
+        self.executors
+            .values()
+            .map(|e| e.current_trace_height)
+            .sum()
     }
 
     fn trace_width(&self) -> usize {
@@ -358,198 +378,216 @@ where
     fn generate_air_proof_input(self) -> AirProofInput<SC> {
         tracing::trace!("Generating air proof input for PowdrChip {}", self.name);
 
-        let is_valid_index = self.get_is_valid_index();
         let num_records = self.current_trace_height();
         let height = next_power_of_two_or_zero(num_records);
         let width = self.air.width();
         let mut values = Val::<SC>::zero_vec(height * width);
+        // we'll fill values from each APC in turn, this keeps track where we are in the full trace
+        let mut values_curr_record = 0;
+        // this is just for sanity checking later
+        let all_is_valid_ids = self
+            .executors
+            .values()
+            .map(|executor| executor.is_valid_poly_id)
+            .collect::<Vec<_>>();
 
-        // for each original opcode, the name of the dummy air it corresponds to
-        let air_name_by_opcode = self
-            .executor
-            .instructions
-            .iter()
-            .map(|instruction| instruction.opcode())
-            .unique()
-            .map(|opcode| {
-                (
-                    opcode,
-                    self.executor
-                        .inventory
-                        .get_executor(opcode)
-                        .unwrap()
-                        .air_name(),
-                )
-            })
-            .collect::<HashMap<_, _>>();
+        for (_opcode, executor) in self.executors {
+            let is_valid_index = self.air.column_index_by_poly_id[&executor.is_valid_poly_id];
 
-        let dummy_trace_by_air_name: HashMap<_, _> = self
-            .executor
-            .inventory
-            .executors
-            .into_iter()
-            .map(|executor| {
-                (
-                    executor.air_name(),
-                    Chip::<SC>::generate_air_proof_input(executor)
-                        .raw
-                        .common_main
-                        .unwrap(),
-                )
-            })
-            .collect();
+            // for each original opcode, the name of the dummy air it corresponds to
+            let air_name_by_opcode = executor
+                .instructions
+                .iter()
+                .map(|instruction| instruction.opcode())
+                .unique()
+                .map(|opcode| {
+                    (
+                        opcode,
+                        executor
+                            .inventory
+                            .get_executor(opcode)
+                            .unwrap()
+                            .air_name(),
+                    )
+                })
+                .collect::<HashMap<_, _>>();
 
-        let instruction_index_to_table_offset = self
-            .executor
-            .instructions
-            .iter()
-            .enumerate()
-            .scan(
-                HashMap::default(),
-                |counts: &mut HashMap<&str, usize>, (index, instruction)| {
-                    let air_name = air_name_by_opcode.get(&instruction.opcode()).unwrap();
-                    let count = counts.entry(air_name).or_default();
-                    let current_count = *count;
-                    *count += 1;
-                    Some((index, (air_name, current_count)))
-                },
-            )
-            .collect::<HashMap<_, _>>();
+            let dummy_trace_by_air_name: HashMap<_, _> = executor
+                .inventory
+                .executors
+                .into_iter()
+                .map(|executor| {
+                    (
+                        executor.air_name(),
+                        Chip::<SC>::generate_air_proof_input(executor)
+                            .raw
+                            .common_main
+                            .unwrap(),
+                    )
+                })
+                .collect();
 
-        let occurrences_by_table_name: HashMap<&String, usize> = self
-            .executor
-            .instructions
-            .iter()
-            .map(|instruction| air_name_by_opcode.get(&instruction.opcode()).unwrap())
-            .counts();
+            let instruction_index_to_table_offset = executor
+                .instructions
+                .iter()
+                .enumerate()
+                .scan(
+                    HashMap::default(),
+                    |counts: &mut HashMap<&str, usize>, (index, instruction)| {
+                        let air_name = air_name_by_opcode.get(&instruction.opcode()).unwrap();
+                        let count = counts.entry(air_name).or_default();
+                        let current_count = *count;
+                        *count += 1;
+                        Some((index, (air_name, current_count)))
+                    },
+                )
+                .collect::<HashMap<_, _>>();
 
-        // A vector of HashMap<dummy_trace_index, apc_trace_index> by instruction, empty HashMap if none maps to apc
-        let dummy_trace_index_to_apc_index_by_instruction: Vec<HashMap<usize, usize>> = self
-            .executor
-            .instructions
-            .iter()
-            .map(|instruction| {
-                // look up how many dummy‐cells this AIR produces:
-                let air_width = dummy_trace_by_air_name
-                    .get(air_name_by_opcode.get(&instruction.opcode()).unwrap())
-                    .unwrap()
-                    .width();
-
-                // build a map only of the (dummy_index -> apc_index) pairs
-                let mut map = HashMap::with_capacity(air_width);
-                for dummy_trace_index in 0..air_width {
-                    if let Ok(apc_index) = global_index(
-                        dummy_trace_index,
-                        instruction,
-                        &self.air.column_index_by_poly_id,
-                    ) {
-                        if map.insert(dummy_trace_index, apc_index).is_some() {
-                            panic!(
-                                "duplicate dummy_trace_index {} for instruction opcode {:?}",
-                                dummy_trace_index,
-                                instruction.opcode()
-                            );
+            let occurrences_by_table_name: HashMap<&String, usize> = executor
+                .instructions
+                .iter()
+                .map(|instruction| air_name_by_opcode.get(&instruction.opcode()).unwrap())
+                .counts();
+
+            // A vector of HashMap<dummy_trace_index, apc_trace_index> by instruction, empty HashMap if none maps to apc
+            let dummy_trace_index_to_apc_index_by_instruction: Vec<HashMap<usize, usize>> = executor
+                .instructions
+                .iter()
+                .map(|instruction| {
+                    // look up how many dummy‐cells this AIR produces:
+                    let air_width = dummy_trace_by_air_name
+                        .get(air_name_by_opcode.get(&instruction.opcode()).unwrap())
+                        .unwrap()
+                        .width();
+
+                    // build a map only of the (dummy_index -> apc_index) pairs
+                    let mut map = HashMap::with_capacity(air_width);
+                    for dummy_trace_index in 0..air_width {
+                        if let Ok(apc_index) = global_index(
+                            dummy_trace_index,
+                            instruction,
+                            &self.air.column_index_by_poly_id,
+                        ) {
+                            if map.insert(dummy_trace_index, apc_index).is_some() {
+                                panic!(
+                                    "duplicate dummy_trace_index {} for instruction opcode {:?}",
+                                    dummy_trace_index,
+                                    instruction.opcode()
+                                );
+                            }
                         }
                     }
-                }
-                map
-            })
-            .collect();
-
-        assert_eq!(
-            self.executor.instructions.len(),
-            dummy_trace_index_to_apc_index_by_instruction.len()
-        );
-
-        let dummy_values = (0..num_records).into_par_iter().map(|record_index| {
-            (0..self.executor.instructions.len())
-                .map(|index| {
-                    // get the air name and offset for this instruction (by index)
-                    let (air_name, offset) = instruction_index_to_table_offset.get(&index).unwrap();
-                    // get the table
-                    let table = dummy_trace_by_air_name.get(*air_name).unwrap();
-                    // get how many times this table is used per record
-                    let occurrences_per_record = occurrences_by_table_name.get(air_name).unwrap();
-                    // get the width of each occurrence
-                    let width = table.width();
-                    // start after the previous record ended, and offset by the correct offset
-                    let start = (record_index * occurrences_per_record + offset) * width;
-                    // end at the start + width
-                    let end = start + width;
-                    &table.values[start..end]
+                    map
                 })
-                .collect_vec()
-        });
+                .collect();
+
+            assert_eq!(
+                executor.instructions.len(),
+                dummy_trace_index_to_apc_index_by_instruction.len()
+            );
+
+            let dummy_num_records = executor.current_trace_height;
+
+            let dummy_values = (0..dummy_num_records).into_par_iter().map(|record_index| {
+                (0..executor.instructions.len())
+                    .map(|index| {
+                        // get the air name and offset for this instruction (by index)
+                        let (air_name, offset) = instruction_index_to_table_offset.get(&index).unwrap();
+                        // get the table
+                        let table = dummy_trace_by_air_name.get(*air_name).unwrap();
+                        // get how many times this table is used per record
+                        let occurrences_per_record = occurrences_by_table_name.get(air_name).unwrap();
+                        // get the width of each occurrence
+                        let width = table.width();
+                        // start after the previous record ended, and offset by the correct offset
+                        let start = (record_index * occurrences_per_record + offset) * width;
+                        // end at the start + width
+                        let end = start + width;
+                        &table.values[start..end]
+                    })
+                    .collect_vec()
+            });
 
-        // go through the final table and fill in the values
-        values
+            // go through the final table and fill in the values
+            values
             // a record is `width` values
-            .par_chunks_mut(width)
-            .zip(dummy_values)
-            .for_each(|(row_slice, dummy_values)| {
-                // map the dummy rows to the autoprecompile row
-                for (instruction_id, (instruction, dummy_row)) in self
-                    .executor
-                    .instructions
-                    .iter()
-                    .zip_eq(dummy_values)
-                    .enumerate()
-                {
-                    let evaluator = RowEvaluator::new(dummy_row, None);
-
-                    // first remove the side effects of this row on the main periphery
-                    for range_checker_send in self
-                        .executor
-                        .air_by_opcode_id
-                        .get(&instruction.as_ref().opcode.as_usize())
-                        .unwrap()
-                        .bus_interactions
+                .par_chunks_mut(width)
+                .skip(values_curr_record)
+                .zip(dummy_values)
+                .for_each(|(row_slice, dummy_values)| {
+                    // map the dummy rows to the autoprecompile row
+                    for (instruction_id, (instruction, dummy_row)) in executor
+                        .instructions
                         .iter()
-                        .filter(|i| i.id == 3)
+                        .zip_eq(dummy_values)
+                        .enumerate()
                     {
-                        let mult = evaluator
-                            .eval_expr(&range_checker_send.mult)
-                            .as_canonical_u32();
-                        let args = range_checker_send
-                            .args
+                        let evaluator = RowEvaluator::new(dummy_row, None);
+
+                        // first remove the side effects of this row on the main periphery
+                        for range_checker_send in executor
+                            .air_by_opcode_id
+                            .get(&instruction.as_ref().opcode.as_usize())
+                            .unwrap()
+                            .bus_interactions
                             .iter()
-                            .map(|arg| evaluator.eval_expr(arg).as_canonical_u32())
-                            .collect_vec();
-                        let [value, max_bits] = args.try_into().unwrap();
-                        for _ in 0..mult {
-                            self.periphery
-                                .range_checker
-                                .remove_count(value, max_bits as usize);
+                            .filter(|i| i.id == 3)
+                        {
+                            let mult = evaluator
+                                .eval_expr(&range_checker_send.mult)
+                                .as_canonical_u32();
+                            let args = range_checker_send
+                                .args
+                                .iter()
+                                .map(|arg| evaluator.eval_expr(arg).as_canonical_u32())
+                                .collect_vec();
+                            let [value, max_bits] = args.try_into().unwrap();
+                            for _ in 0..mult {
+                                self.periphery
+                                    .range_checker
+                                    .remove_count(value, max_bits as usize);
+                            }
                         }
+
+                        write_dummy_to_autoprecompile_row(
+                            row_slice,
+                            dummy_row,
+                            &dummy_trace_index_to_apc_index_by_instruction[instruction_id],
+                        );
                     }
 
-                    write_dummy_to_autoprecompile_row(
-                        row_slice,
-                        dummy_row,
-                        &dummy_trace_index_to_apc_index_by_instruction[instruction_id],
-                    );
-                }
+                    // double-check is_valid of every precompile is ZERO
+                    for id in all_is_valid_ids.iter() {
+                        assert_eq!(
+                            row_slice[self.air.column_index_by_poly_id[id]],
+                            <Val<SC>>::ZERO
+                        );
+                    }
 
-                // Set the is_valid column to 1
-                row_slice[is_valid_index] = <Val<SC>>::ONE;
+                    // Set the is_valid of the active chip to 1
+                    row_slice[is_valid_index] = <Val<SC>>::ONE;
 
-                let evaluator =
-                    RowEvaluator::new(row_slice, Some(&self.air.column_index_by_poly_id));
+                    let evaluator =
+                        RowEvaluator::new(row_slice, Some(&self.air.column_index_by_poly_id));
 
-                // replay the side effects of this row on the main periphery
-                for bus_interaction in self.air.machine.bus_interactions.iter() {
-                    let mult = evaluator
-                        .eval_expr(&bus_interaction.mult)
-                        .as_canonical_u32();
-                    let args = bus_interaction
-                        .args
-                        .iter()
-                        .map(|arg| evaluator.eval_expr(arg).as_canonical_u32())
-                        .collect_vec();
+                    // replay the side effects of this row on the main periphery
+                    for bus_interaction in self.air.machine.bus_interactions.iter() {
+                        let mult = evaluator
+                            .eval_expr(&bus_interaction.mult)
+                            .as_canonical_u32();
+                        let args = bus_interaction
+                            .args
+                            .iter()
+                            .map(|arg| evaluator.eval_expr(arg).as_canonical_u32())
+                            .collect_vec();
 
-                    self.periphery.apply(bus_interaction.id, mult, &args);
-                }
-            });
+                        self.periphery.apply(bus_interaction.id, mult, &args);
+                    }
+                });
+
+            // next executor will fill in the next part of the trace
+            values_curr_record += dummy_num_records;
+        }
 
         let trace = RowMajorMatrix::new(values, width);
 
diff --git a/openvm/src/powdr_extension/mod.rs b/openvm/src/powdr_extension/mod.rs
index 354e3b3934..4d75bec1e4 100644
--- a/openvm/src/powdr_extension/mod.rs
+++ b/openvm/src/powdr_extension/mod.rs
@@ -6,4 +6,4 @@ pub mod opcode;
 mod vm;
 
 pub use opcode::PowdrOpcode;
-pub use vm::{OriginalInstruction, PowdrExecutor, PowdrExtension, PowdrPeriphery, PowdrPrecompile};
+pub use vm::{OriginalInstruction, PowdrExecutor, PowdrExtension, PowdrPeriphery, PowdrPrecompile, PowdrStackedPrecompile};
diff --git a/openvm/src/powdr_extension/vm.rs b/openvm/src/powdr_extension/vm.rs
index df4871a243..b8a86bd535 100644
--- a/openvm/src/powdr_extension/vm.rs
+++ b/openvm/src/powdr_extension/vm.rs
@@ -31,7 +31,7 @@ pub type SdkVmInventory<F> = VmInventory<SdkVmConfigExecutor<F>, SdkVmConfigPeri
 #[derive(Clone, Deserialize, Serialize)]
 #[serde(bound = "F: Field")]
 pub struct PowdrExtension<F: PrimeField32> {
-    pub precompiles: Vec<PowdrPrecompile<F>>,
+    pub precompiles: Vec<PowdrStackedPrecompile<F>>,
     pub base_config: SdkVmConfig,
 }
 
@@ -89,8 +89,17 @@ impl<F> PowdrPrecompile<F> {
     }
 }
 
+#[derive(Clone, Serialize, Deserialize)]
+#[serde(bound = "F: Field")]
+pub struct PowdrStackedPrecompile<F> {
+    /// stacked precompiles by opcode
+    pub precompiles: BTreeMap<PowdrOpcode, PowdrPrecompile<F>>,
+    /// constraints
+    pub machine: SymbolicMachine<F>,
+}
+
 impl<F: PrimeField32> PowdrExtension<F> {
-    pub fn new(precompiles: Vec<PowdrPrecompile<F>>, base_config: SdkVmConfig) -> Self {
+    pub fn new(precompiles: Vec<PowdrStackedPrecompile<F>>, base_config: SdkVmConfig) -> Self {
         Self {
             precompiles,
             base_config,
@@ -148,7 +157,13 @@ impl<F: PrimeField32> VmExtension<F> for PowdrExtension<F> {
                 ),
             );
 
-            inventory.add_executor(powdr_chip, once(precompile.opcode.global_opcode()))?;
+            inventory.add_executor(
+                powdr_chip,
+                precompile
+                    .precompiles
+                    .keys()
+                    .map(|opcode| opcode.global_opcode()),
+            )?;
         }
 
         Ok(inventory)

From 7a35e47d3354b68cc473ec0e99d3fe597e37b809 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 19 May 2025 16:17:42 -0300
Subject: [PATCH 03/39] lower degree bound of optimize step

---
 openvm/src/customize_exe.rs        | 16 ++++++----------
 openvm/src/powdr_extension/chip.rs |  1 -
 openvm/src/powdr_extension/vm.rs   |  1 -
 3 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 4ac797e436..ba3bb2f224 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -370,12 +370,7 @@ fn air_stacking<F: PrimeField32>(
         // their remapping into exclusive columns.
         let is_valid_new_id = is_valid_start + idx as u64;
 
-        println!("pcp width: {}", pcp.machine.unique_columns().count());
-
-        println!(
-            "\tIS_VALID_IDX: {} ===> {}",
-            pcp.is_valid_column.id.id, is_valid_new_id
-        );
+        println!("\tpcp width: {}", pcp.machine.unique_columns().count());
 
         // remap is_valid column in constraints and interactions
         let mut remapped = pcp.machine.clone();
@@ -576,9 +571,9 @@ fn generate_autoprecompile<F: PrimeField32, P: FieldElement>(
                 args: [
                     instr.a, instr.b, instr.c, instr.d, instr.e, instr.f, instr.g,
                 ]
-                .iter()
-                .map(|f| to_powdr_field::<F, P>(*f))
-                .collect(),
+                    .iter()
+                    .map(|f| to_powdr_field::<F, P>(*f))
+                    .collect(),
             };
 
             if is_jump(&instr.opcode) {
@@ -601,7 +596,8 @@ fn generate_autoprecompile<F: PrimeField32, P: FieldElement>(
 
     let (precompile, subs) = autoprecompiles.build(
         OpenVmBusInteractionHandler::default(),
-        OPENVM_DEGREE_BOUND,
+        // chip stacking needs to add guards to bus arguments also, so we restrict the optimizer by 1 degree here.
+        OPENVM_DEGREE_BOUND - 1,
         apc_opcode as u32,
     );
 
diff --git a/openvm/src/powdr_extension/chip.rs b/openvm/src/powdr_extension/chip.rs
index e70826c9c1..54b85a105b 100644
--- a/openvm/src/powdr_extension/chip.rs
+++ b/openvm/src/powdr_extension/chip.rs
@@ -8,7 +8,6 @@ use std::{
 use crate::utils::algebraic_to_symbolic;
 
 use super::{
-    opcode::PowdrOpcode,
     vm::{OriginalInstruction, SdkVmInventory},
     PowdrStackedPrecompile,
 };
diff --git a/openvm/src/powdr_extension/vm.rs b/openvm/src/powdr_extension/vm.rs
index b8a86bd535..ea0ca10199 100644
--- a/openvm/src/powdr_extension/vm.rs
+++ b/openvm/src/powdr_extension/vm.rs
@@ -1,7 +1,6 @@
 // Mostly taken from [this openvm extension](https://github.com/openvm-org/openvm/blob/1b76fd5a900a7d69850ee9173969f70ef79c4c76/extensions/rv32im/circuit/src/extension.rs#L185) and simplified to only handle a single opcode with its necessary dependencies
 
 use std::collections::BTreeMap;
-use std::iter::once;
 
 use derive_more::From;
 

From 323e706d5632ab6630cb58e7231d032d9d24c564 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 19 May 2025 17:26:25 -0300
Subject: [PATCH 04/39] group by powers of two

---
 openvm/src/customize_exe.rs | 166 +++++++++++++++++++-----------------
 1 file changed, 90 insertions(+), 76 deletions(-)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index ba3bb2f224..4762aa44f1 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -352,97 +352,111 @@ fn air_stacking<F: PrimeField32>(
 
     extensions.iter_mut().for_each(compact_ids);
 
-    // take the max id in all pcps and add 1.
-    let is_valid_start = 1 + extensions.iter()
-        .flat_map(|pcp| pcp.original_instructions.iter().flat_map(|instr| instr.subs.iter()))
-        .max()
-        .unwrap();
-
-    let mut stacked_constraints = vec![];
-    let mut interactions_by_machine = vec![];
-
-    println!("stacking {} precompiles", extensions.len());
-
-    let mut is_valid_sum: Option<AlgebraicExpression<F>> = None;
-
-    for (idx, pcp) in extensions.iter_mut().enumerate() {
-        // is_valid columns cannot be shared between precompiles. Here we do
-        // their remapping into exclusive columns.
-        let is_valid_new_id = is_valid_start + idx as u64;
-
-        println!("\tpcp width: {}", pcp.machine.unique_columns().count());
-
-        // remap is_valid column in constraints and interactions
-        let mut remapped = pcp.machine.clone();
-        remapped.pre_visit_expressions_mut(&mut |expr| {
-            if let AlgebraicExpression::Reference(r) = expr {
-                assert!(r.poly_id.id <= is_valid_start);
-                if r.poly_id == pcp.is_valid_column.id {
-                    // we assume each pcp to have a specific column named "is_valid"
-                    assert!(r.name == "is_valid");
-                    r.poly_id.id = is_valid_new_id;
-                    r.name = format!("is_valid_{}", pcp.opcode.global_opcode().as_usize());
-                } else {
-                    // we need to rename columns here because `unique_columns` relies on both `Column::name` and `Column::id`
-                    r.name = format!("col_{}", r.poly_id.id);
+    // create apc groups by number of columns
+    let mut groups: HashMap<usize, Vec<PowdrPrecompile<F>>> = Default::default();
+    for pcp in extensions {
+        let n_cols = pcp.machine.unique_columns().count().next_power_of_two();
+        groups.entry(n_cols).or_default().push(pcp);
+    }
+
+    let mut result = vec![];
+
+    for mut extensions in groups.into_values() {
+        let mut stacked_constraints = vec![];
+        let mut interactions_by_machine = vec![];
+
+        println!("stacking {} precompiles", extensions.len());
+
+        // take the max id in all pcps and add 1.
+        let is_valid_start = 1 + extensions.iter()
+            .flat_map(|pcp| pcp.original_instructions.iter().flat_map(|instr| instr.subs.iter()))
+            .max()
+            .unwrap();
+
+
+        let mut is_valid_sum: Option<AlgebraicExpression<F>> = None;
+
+        for (idx, pcp) in extensions.iter_mut().enumerate() {
+            // is_valid columns cannot be shared between precompiles. Here we do
+            // their remapping into exclusive columns.
+            let is_valid_new_id = is_valid_start + idx as u64;
+
+            println!("\tpcp width: {}", pcp.machine.unique_columns().count());
+
+            // remap is_valid column in constraints and interactions
+            let mut remapped = pcp.machine.clone();
+            remapped.pre_visit_expressions_mut(&mut |expr| {
+                if let AlgebraicExpression::Reference(r) = expr {
+                    assert!(r.poly_id.id <= is_valid_start);
+                    if r.poly_id == pcp.is_valid_column.id {
+                        // we assume each pcp to have a specific column named "is_valid"
+                        assert!(r.name == "is_valid");
+                        r.poly_id.id = is_valid_new_id;
+                        r.name = format!("is_valid_{}", pcp.opcode.global_opcode().as_usize());
+                    } else {
+                        // we need to rename columns here because `unique_columns` relies on both `Column::name` and `Column::id`
+                        r.name = format!("col_{}", r.poly_id.id);
+                    }
                 }
-            }
-        });
+            });
 
-        // set the is valid column in the original precompile
-        pcp.is_valid_column.id.id = is_valid_new_id;
+            // set the is valid column in the original precompile
+            pcp.is_valid_column.id.id = is_valid_new_id;
 
-        let is_valid = AlgebraicExpression::Reference(AlgebraicReference {
-            name: format!("is_valid_{}", pcp.opcode.global_opcode().as_usize()),
-            poly_id: PolyID {
-                id: is_valid_new_id,
-                ptype: PolynomialType::Committed,
-            },
-            next: false,
-        });
+            let is_valid = AlgebraicExpression::Reference(AlgebraicReference {
+                name: format!("is_valid_{}", pcp.opcode.global_opcode().as_usize()),
+                poly_id: PolyID {
+                    id: is_valid_new_id,
+                    ptype: PolynomialType::Committed,
+                },
+                next: false,
+            });
 
-        // guard interaction payloads so they can be merged later
-        remapped.bus_interactions.iter_mut().for_each(|interaction| {
-            interaction.args.iter_mut().for_each(|arg| {
-                *arg = arg.clone() * is_valid.clone();
+            // guard interaction payloads so they can be merged later
+            remapped.bus_interactions.iter_mut().for_each(|interaction| {
+                interaction.args.iter_mut().for_each(|arg| {
+                    *arg = arg.clone() * is_valid.clone();
+                });
             });
-        });
 
 
-        is_valid_sum = is_valid_sum
-            .map(|sum| sum + is_valid.clone())
-            .or_else(|| Some(is_valid.clone()));
+            is_valid_sum = is_valid_sum
+                .map(|sum| sum + is_valid.clone())
+                .or_else(|| Some(is_valid.clone()));
 
-        stacked_constraints.extend(remapped.constraints);
-        interactions_by_machine.push(remapped.bus_interactions);
-    }
+            stacked_constraints.extend(remapped.constraints);
+            interactions_by_machine.push(remapped.bus_interactions);
+        }
 
-    // enforce only one is_valid is active
-    let one = AlgebraicExpression::Number(F::ONE);
-    let one_hot_is_valid = (one - is_valid_sum.clone().unwrap()) * is_valid_sum.unwrap();
+        // enforce only one is_valid is active
+        let one = AlgebraicExpression::Number(F::ONE);
+        let one_hot_is_valid = (one - is_valid_sum.clone().unwrap()) * is_valid_sum.unwrap();
 
-    stacked_constraints.push(one_hot_is_valid.into());
+        stacked_constraints.push(one_hot_is_valid.into());
 
-    println!("interaction count before: {}", interactions_by_machine.iter().flatten().count());
-    // let stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
-    let stacked_interactions = merge_bus_interactions(interactions_by_machine);
+        println!("interaction count before: {}", interactions_by_machine.iter().flatten().count());
+        // let stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
+        let stacked_interactions = merge_bus_interactions(interactions_by_machine);
 
-    let machine = SymbolicMachine {
-        constraints: stacked_constraints,
-        bus_interactions: stacked_interactions,
-    };
+        let machine = SymbolicMachine {
+            constraints: stacked_constraints,
+            bus_interactions: stacked_interactions,
+        };
+
+        println!("interaction count: {}", machine.bus_interactions.len());
 
-    println!("interaction count: {}", machine.bus_interactions.len());
+        println!("stacked width: {}", machine.unique_columns().count());
 
-    println!("stacked width: {}", machine.unique_columns().count());
+        result.push(PowdrStackedPrecompile {
+            precompiles: extensions
+                .into_iter()
+                .map(|p| (p.opcode.clone(), p))
+                .collect(),
+            machine,
+        });
+    }
 
-    vec![PowdrStackedPrecompile {
-        precompiles: extensions
-            .into_iter()
-            .map(|p| (p.opcode.clone(), p))
-            .collect(),
-        machine,
-    }]
+    result
 }
 
 

From f09fc3481a4128d751108ceedca80727df3b86d0 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Tue, 20 May 2025 16:34:51 -0300
Subject: [PATCH 05/39] use f64::log

---
 openvm/src/customize_exe.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 4762aa44f1..c161254ac9 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -355,8 +355,8 @@ fn air_stacking<F: PrimeField32>(
     // create apc groups by number of columns
     let mut groups: HashMap<usize, Vec<PowdrPrecompile<F>>> = Default::default();
     for pcp in extensions {
-        let n_cols = pcp.machine.unique_columns().count().next_power_of_two();
-        groups.entry(n_cols).or_default().push(pcp);
+        let idx = f64::log(pcp.machine.unique_columns().count() as f64, 2.0).floor() as usize;
+        groups.entry(idx).or_default().push(pcp);
     }
 
     let mut result = vec![];

From 5d060c527507b456406193e955da659f2ff4f0ab Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Fri, 23 May 2025 12:56:16 -0300
Subject: [PATCH 06/39] fix for scroll

- another unknown instruction
- handle interactions to the same bus id with different number of args
---
 openvm/src/customize_exe.rs | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 7fc8c0bd64..40ded8e940 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -94,6 +94,7 @@ pub fn customize<F: PrimeField32>(
         SystemOpcode::TERMINATE.global_opcode().as_usize(),
         0x510, // not sure yet what this is
         0x513, // not sure yet what this is
+        0x516, // not sure yet what this is
         0x51c, // not sure yet what this is
         0x523, // not sure yet what this is
     ];
@@ -287,10 +288,12 @@ fn merge_bus_interactions_simple<F: PrimeField32>(interactions: Vec<Vec<Symbolic
 }
 
 fn merge_bus_interactions<F: PrimeField32>(interactions_by_machine: Vec<Vec<SymbolicBusInteraction<F>>>) -> Vec<SymbolicBusInteraction<F>> {
-    let mut interactions_by_bus: HashMap<u64, Vec<Vec<SymbolicBusInteraction<F>>>> = Default::default();
+    let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<F>>>> = Default::default();
 
     for interactions in interactions_by_machine {
-        let interactions_by_bus_this_machine = interactions.into_iter().into_group_map_by(|interaction| interaction.id);
+        let interactions_by_bus_this_machine = interactions.into_iter()
+            // we group by bus id and number of args
+            .into_group_map_by(|interaction| (interaction.id, interaction.args.len()));
         for (k,v) in interactions_by_bus_this_machine {
             let e = interactions_by_bus.entry(k).or_default();
             e.push(v);

From c6443f708f4f07898c08d70d9af322cb5316e69d Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Fri, 23 May 2025 16:05:53 -0300
Subject: [PATCH 07/39] another unsupported instruction

---
 openvm/src/customize_exe.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 40ded8e940..f77f6cd0ff 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -97,6 +97,7 @@ pub fn customize<F: PrimeField32>(
         0x516, // not sure yet what this is
         0x51c, // not sure yet what this is
         0x523, // not sure yet what this is
+        0x526, // not sure yet what this is
     ];
 
     let mut blocks = collect_basic_blocks(&exe.program, labels, &opcodes_no_apc);

From c950d3fa7c4b3b0ceba451bd5c0c8ce7d468fd0e Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 26 May 2025 16:19:18 -0300
Subject: [PATCH 08/39] handle group of 1 precompile

---
 openvm/src/customize_exe.rs | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index f77f6cd0ff..24ec35ab10 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -350,9 +350,6 @@ fn air_stacking<F: PrimeField32>(
     mut extensions: Vec<PowdrPrecompile<F>>,
 ) -> Vec<PowdrStackedPrecompile<F>> {
     assert!(!extensions.is_empty());
-    // TODO: for now group every precompile into a single stacked one. Next,
-    // group them by some heuristic (i.e., by powers of number of witness
-    // columns)
 
     extensions.iter_mut().for_each(compact_ids);
 
@@ -366,6 +363,22 @@ fn air_stacking<F: PrimeField32>(
     let mut result = vec![];
 
     for mut extensions in groups.into_values() {
+        // if there is only one precompile in the group, we can just return it as a stacked precompile
+        if extensions.len() == 1 {
+            println!("not stacking precompile");
+            result.push(PowdrStackedPrecompile {
+                machine: SymbolicMachine {
+                    constraints: extensions[0].machine.constraints.clone(),
+                    bus_interactions: extensions[0].machine.bus_interactions.clone(),
+                },
+                precompiles: extensions
+                    .into_iter()
+                    .map(|p| (p.opcode.clone(), p))
+                    .collect(),
+            });
+            continue;
+        }
+
         let mut stacked_constraints = vec![];
         let mut interactions_by_machine = vec![];
 

From 4b3c5d4b4fe6650b3954e2e290b421d03e4627bc Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 2 Jun 2025 11:12:54 -0300
Subject: [PATCH 09/39] test joining constraints

---
 autoprecompiles/src/lib.rs  |  4 ++--
 openvm/src/customize_exe.rs | 48 ++++++++++++++++++++++++++++++++++---
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/autoprecompiles/src/lib.rs b/autoprecompiles/src/lib.rs
index 3cae7a129a..aabd603188 100644
--- a/autoprecompiles/src/lib.rs
+++ b/autoprecompiles/src/lib.rs
@@ -37,7 +37,7 @@ pub struct SymbolicInstructionStatement<T> {
     pub args: Vec<T>,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, Ord, PartialOrd, Eq, PartialEq)]
 pub struct SymbolicConstraint<T> {
     pub expr: AlgebraicExpression<T>,
 }
@@ -277,7 +277,7 @@ fn add_guards_constraint<T: FieldElement>(
 /// - There are exactly one execution bus receive and one execution bus send, in this order.
 /// - There is exactly one program bus send.
 fn add_guards<T: FieldElement>(mut machine: SymbolicMachine<T>) -> SymbolicMachine<T> {
-    let pre_degree = machine.degree();
+    // let pre_degree = machine.degree();
 
     let max_id = machine
         .unique_columns()
diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 718cf1348a..02fe974f6d 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -16,6 +16,9 @@ use openvm_stark_backend::{
     interaction::SymbolicInteraction,
     p3_field::{FieldAlgebra, PrimeField32},
 };
+use powdr_autoprecompiles::SymbolicConstraint;
+use powdr_expression::AlgebraicBinaryOperation;
+use powdr_expression::AlgebraicBinaryOperator;
 use powdr_expression::{AlgebraicExpression, visitors::ExpressionVisitable};
 use powdr_autoprecompiles::legacy_expression::{AlgebraicReference, PolyID, PolynomialType};
 use powdr_autoprecompiles::powdr::UniqueColumns;
@@ -386,13 +389,48 @@ fn merge_bus_interactions<P: IntoOpenVm>(interactions_by_machine: Vec<Vec<Symbol
     }).collect_vec()
 }
 
+fn canonicalize_expression<P: IntoOpenVm>(expr: &mut AlgebraicExpression<P, AlgebraicReference>) {
+    expr.pre_visit_expressions_mut(&mut |expr| {
+        if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) = expr {
+            if *op == AlgebraicBinaryOperator::Mul || *op == AlgebraicBinaryOperator::Add && left > right {
+                std::mem::swap(left, right);
+            }
+        }
+    });
+}
+
+/// assumes ALL constraints are of the form `Y * is_valid_expr`
+fn join_constraints<P: IntoOpenVm>(mut constraints: Vec<SymbolicConstraint<P>>) -> Vec<SymbolicConstraint<P>> {
+    constraints.sort();
+    let mut result: Vec<SymbolicConstraint<P>> = vec![];
+    'outer: for c1 in constraints {
+        for c2 in result.iter_mut() {
+            if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {left: l1, op: AlgebraicBinaryOperator::Mul, right: is_valid1}) = &c1.expr {
+                if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {left: l2, op: AlgebraicBinaryOperator::Mul, right: is_valid2}) = &mut c2.expr {
+                    if l1 == l2 {
+                        *is_valid2 = Box::new(*is_valid1.clone() + *is_valid2.clone());
+                        continue 'outer;
+                    }
+                }
+            }
+        }
+        result.push(c1);
+    }
+    result
+}
+
+
 // perform air stacking of multiple precompiles into stacked precompiles.
 fn air_stacking<P: IntoOpenVm>(
     mut extensions: Vec<PowdrPrecompile<P>>,
 ) -> Vec<PowdrStackedPrecompile<P>> {
     assert!(!extensions.is_empty());
 
-    extensions.iter_mut().for_each(compact_ids);
+    extensions.iter_mut().for_each(|ext| {
+        ext.machine.constraints.sort();
+        ext.machine.constraints.iter_mut().for_each(|c| canonicalize_expression(&mut c.expr));
+        compact_ids(ext)
+    });
 
     // create apc groups by number of columns
     let mut groups: HashMap<usize, Vec<PowdrPrecompile<P>>> = Default::default();
@@ -486,6 +524,10 @@ fn air_stacking<P: IntoOpenVm>(
             interactions_by_machine.push(remapped.bus_interactions);
         }
 
+        println!("before joining constraints: {}", stacked_constraints.len());
+        let mut stacked_constraints = join_constraints(stacked_constraints);
+        println!("after joining constraints: {}", stacked_constraints.len());
+
         // enforce only one is_valid is active
         let one = AlgebraicExpression::Number(P::ONE);
         let one_hot_is_valid = (one - is_valid_sum.clone().unwrap()) * is_valid_sum.unwrap();
@@ -493,8 +535,8 @@ fn air_stacking<P: IntoOpenVm>(
         stacked_constraints.push(one_hot_is_valid.into());
 
         println!("interaction count before: {}", interactions_by_machine.iter().flatten().count());
-        // let stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
-        let stacked_interactions = merge_bus_interactions(interactions_by_machine);
+        let stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
+        // let stacked_interactions = merge_bus_interactions(interactions_by_machine);
 
         let machine = SymbolicMachine {
             constraints: stacked_constraints,

From a7f5ff1e8a176ca55280e5c650fb5282fc6d7cb6 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 2 Jun 2025 15:20:28 -0300
Subject: [PATCH 10/39] sort

---
 autoprecompiles/src/lib.rs  |  4 ++--
 openvm/src/customize_exe.rs | 21 ++++++++++++++++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/autoprecompiles/src/lib.rs b/autoprecompiles/src/lib.rs
index aabd603188..1efea099e7 100644
--- a/autoprecompiles/src/lib.rs
+++ b/autoprecompiles/src/lib.rs
@@ -66,11 +66,11 @@ impl<T: Clone + Ord + std::fmt::Display> Children<AlgebraicExpression<T>>
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
+#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash, Ord, PartialOrd)]
 pub struct SymbolicBusInteraction<T> {
     pub id: u64,
-    pub mult: AlgebraicExpression<T>,
     pub args: Vec<AlgebraicExpression<T>>,
+    pub mult: AlgebraicExpression<T>,
 }
 
 impl<T: Display> Display for SymbolicBusInteraction<T> {
diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 02fe974f6d..94eceb37aa 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -19,6 +19,7 @@ use openvm_stark_backend::{
 use powdr_autoprecompiles::SymbolicConstraint;
 use powdr_expression::AlgebraicBinaryOperation;
 use powdr_expression::AlgebraicBinaryOperator;
+use powdr_expression::AlgebraicUnaryOperation;
 use powdr_expression::{AlgebraicExpression, visitors::ExpressionVisitable};
 use powdr_autoprecompiles::legacy_expression::{AlgebraicReference, PolyID, PolynomialType};
 use powdr_autoprecompiles::powdr::UniqueColumns;
@@ -419,6 +420,22 @@ fn join_constraints<P: IntoOpenVm>(mut constraints: Vec<SymbolicConstraint<P>>)
     result
 }
 
+fn has_same_structure<T: PrimeField32>(expr1: &AlgebraicExpression<T, AlgebraicReference>, expr2: &AlgebraicExpression<T, AlgebraicReference>) -> bool {
+    match (expr1, expr2) {
+        (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation{ left: left1, op: op1, right: right1 }),
+         AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation{ left: left2, op: op2, right: right2 })) => {
+            op1 == op2 && has_same_structure(left1, left2) && has_same_structure(right1, right2)
+        }
+        (AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op1, expr: expr1 }),
+         AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op2, expr: expr2 })) => {
+            op1 == op2 && has_same_structure(expr1, expr2)
+        }
+        (AlgebraicExpression::Reference(_), AlgebraicExpression::Reference(_)) => true,
+        (AlgebraicExpression::Number(n1), AlgebraicExpression::Number(n2)) => n1 == n2,
+        _ => false
+    }
+}
+
 
 // perform air stacking of multiple precompiles into stacked precompiles.
 fn air_stacking<P: IntoOpenVm>(
@@ -526,6 +543,7 @@ fn air_stacking<P: IntoOpenVm>(
 
         println!("before joining constraints: {}", stacked_constraints.len());
         let mut stacked_constraints = join_constraints(stacked_constraints);
+        stacked_constraints.sort();
         println!("after joining constraints: {}", stacked_constraints.len());
 
         // enforce only one is_valid is active
@@ -535,7 +553,8 @@ fn air_stacking<P: IntoOpenVm>(
         stacked_constraints.push(one_hot_is_valid.into());
 
         println!("interaction count before: {}", interactions_by_machine.iter().flatten().count());
-        let stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
+        let mut stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
+        stacked_interactions.sort();
         // let stacked_interactions = merge_bus_interactions(interactions_by_machine);
 
         let machine = SymbolicMachine {

From e48bf1b108e26fe6b8da9aa34c9ec0675374c7da Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Tue, 3 Jun 2025 19:03:25 -0300
Subject: [PATCH 11/39] seems to work!

---
 openvm/src/customize_exe.rs | 238 +++++++++++++++++++++++++++++++++++-
 1 file changed, 232 insertions(+), 6 deletions(-)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 94eceb37aa..379d7d3d3e 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -36,6 +36,8 @@ use crate::{
     utils::symbolic_to_algebraic,
 };
 
+use bimap::BiMap;
+
 pub const OPENVM_DEGREE_BOUND: usize = 5;
 
 // TODO: read this from program
@@ -430,12 +432,223 @@ fn has_same_structure<T: PrimeField32>(expr1: &AlgebraicExpression<T, AlgebraicR
          AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op2, expr: expr2 })) => {
             op1 == op2 && has_same_structure(expr1, expr2)
         }
-        (AlgebraicExpression::Reference(_), AlgebraicExpression::Reference(_)) => true,
+        (AlgebraicExpression::Reference(r1), AlgebraicExpression::Reference(r2)) => {
+            r1.poly_id.id == r2.poly_id.id
+        },
         (AlgebraicExpression::Number(n1), AlgebraicExpression::Number(n2)) => n1 == n2,
         _ => false
     }
 }
 
+fn create_mapping<T: PrimeField32>(from: &AlgebraicExpression<T>, to: &AlgebraicExpression<T>) -> BTreeMap<u64, u64> {
+    fn create_mapping_inner<T: PrimeField32>(from: &AlgebraicExpression<T>, to: &AlgebraicExpression<T>) -> BTreeMap<u64, u64> {
+        match (from, to) {
+            (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left1, op: op1, right: right1 }),
+             AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left2, op: op2, right: right2 })) => {
+                assert_eq!(op1, op2);
+                let mut left = create_mapping_inner(left1, left2);
+                let right = create_mapping_inner(right1, right2);
+                // assert both sides don't conflict
+                assert!(extend_if_no_conflicts(&mut left, right), "conflict in mapping: {from} => {to}");
+                left
+            }
+            (AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op1, expr: expr1 }),
+             AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op2, expr: expr2 })) => {
+                assert_eq!(op1, op2);
+                create_mapping_inner(&expr1, &expr2)
+            }
+            (AlgebraicExpression::Reference(from), AlgebraicExpression::Reference(to)) => {
+                let mappings = BTreeMap::from([(from.poly_id.id, to.poly_id.id)]);
+                // ref2.poly_id.id = ref1.poly_id.id;
+                mappings
+            }
+            _ => BTreeMap::new(),
+        }
+    }
+
+    create_mapping_inner(from, to)
+}
+
+/// assumes all expressions of the form `Y * is_valid_expr`
+// fn create_mapping<T: PrimeField32>(expr1: &AlgebraicExpression<T>, expr2: &AlgebraicExpression<T>) -> BTreeMap<u64, u64> {
+//     fn create_mapping_inner<T: PrimeField32>(expr1: &Box<AlgebraicExpression<T>>, expr2: &Box<AlgebraicExpression<T>>) -> BTreeMap<u64, u64> {
+//         match (expr1.as_ref(), expr2.as_ref()) {
+//             (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left1, op: op1, right: right1 }),
+//              AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left2, op: op2, right: right2 })) => {
+//                 assert_eq!(op1, op2);
+//                 let mut left = create_mapping_inner(left1, left2);
+//                 let right = create_mapping_inner(right1, right2);
+//                 // assert both sides don't conflict
+//                 assert!(extend_if_no_conflicts(&mut left, right), "conflict in mapping: {expr1} => {expr2}");
+//                 left
+//             }
+//             (AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op1, expr: expr1 }),
+//              AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op2, expr: expr2 })) => {
+//                 assert_eq!(op1, op2);
+//                 create_mapping_inner(expr1, expr2)
+//             }
+//             (AlgebraicExpression::Reference(ref1), AlgebraicExpression::Reference(ref2)) => {
+//                 let mappings = BTreeMap::from([(ref2.poly_id.id, ref1.poly_id.id)]); // CHECK
+//                 mappings
+//             }
+//             (AlgebraicExpression::Number(n1), AlgebraicExpression::Number(n2)) => {
+//                 assert_eq!(n1, n2, "numbers do not match: {n1} != {n2}");
+//                 BTreeMap::new() // no mapping needed for numbers
+//             }
+//             _ => unreachable!("invalid expression mapping: {expr1} => {expr2}"),
+//         }
+//     }
+//
+//     let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: expr1, op: _, right: _is_valid }) = expr1 else {
+//         panic!("invalid expression mapping: {expr1} => {expr2}");
+//     };
+//     let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: expr2, op: _, right: _is_valid }) = expr2 else {
+//         panic!("invalid expression: {expr1} => {expr2}");
+//     };
+//     create_mapping_inner(expr1, expr2)
+// }
+
+fn extend_if_no_conflicts(
+    mappings: &mut BTreeMap<u64, u64>,
+    new_mappings: BTreeMap<u64, u64>,
+) -> bool {
+    for (k, v) in &new_mappings {
+        if let Some(existing_v) = mappings.get(k) {
+            if existing_v != v {
+                return false; // conflict found
+            }
+        }
+    }
+    mappings.extend(new_mappings);
+    true
+}
+
+/// changes poly_id to match the order in which they are encountered in the expression.
+/// This allows us to compare two expressions for the same "structure".
+/// e.g., a+a == b+b and a+a != a+b
+fn expr_poly_id_by_order<F: PrimeField32>(mut expr: AlgebraicExpression<F>) -> AlgebraicExpression<F> {
+    let mut curr_id = 0;
+    let mut id_map: HashMap<u64, u64> = Default::default();
+    let mut new_poly_id = |old_id: u64| {
+        *id_map.entry(old_id).or_insert_with(|| {
+            let id = curr_id;
+            curr_id += 1;
+            id
+        })
+    };
+    expr.pre_visit_expressions_mut(&mut |e| {
+        if let AlgebraicExpression::Reference(r) = e {
+            r.poly_id.id = new_poly_id(r.poly_id.id);
+            r.name = format!("col_{}", r.poly_id.id); // this one not needed, just for printing
+        }
+    });
+    expr
+}
+
+fn expr_orig_poly_ids<F: PrimeField32>(mut expr: AlgebraicExpression<F>) -> AlgebraicExpression<F> {
+    expr.pre_visit_expressions_mut(&mut |e| {
+        if let AlgebraicExpression::Reference(r) = e {
+            r.name = format!("col_{}", r.poly_id.id); // this one not needed, just for printing
+        }
+    });
+    expr
+}
+
+#[derive(Default)]
+struct ColumnAssigner<F: PrimeField32> {
+    // original PCPs
+    pcps: Vec<PowdrPrecompile<F>>,
+    // id mapping for each PCP (by idx)
+    mappings: BTreeMap<usize, BTreeMap<u64, u64>>,
+}
+
+impl<F: PrimeField32> ColumnAssigner<F> {
+    fn assign_pcp(
+        &mut self,
+        pcp: &mut PowdrPrecompile<F>,
+    ) {
+        let idx = self.pcps.len();
+        println!("Assinging PCP {}", idx);
+        let mut mappings = BTreeMap::new();
+        // for each constraint, we want to find if there is a constraint in
+        // another machine with the same structure that we can assign the same
+        // ids.
+        // There may be multiple such constraints, so we try them all (as some
+        // of the ids in the current constraint may already be assigned)
+
+        'outer: for c in &mut pcp.machine.constraints {
+            for pcp2 in self.pcps.iter() {
+                for c2 in &pcp2.machine.constraints {
+                    if has_same_structure(&expr_poly_id_by_order(c.expr.clone()),
+                                          &expr_poly_id_by_order(c2.expr.clone())) {
+                        println!("Found same structure: \n\t{}\n\n\t{}",
+                                 &expr_poly_id_by_order(c.expr.clone()),
+                                 &expr_poly_id_by_order(c2.expr.clone()));
+
+                        println!("Found same structure (ids): \n\t{}\n\n\t{}",
+                                 &expr_orig_poly_ids(c.expr.clone()),
+                                 &expr_orig_poly_ids(c2.expr.clone()));
+
+                        // assign the same ids
+                        let new_mappings = create_mapping(&c.expr, &c2.expr);
+                        if extend_if_no_conflicts(
+                            &mut mappings,
+                            new_mappings,
+                        ) {
+                            println!("\tMappings: {:#?}", mappings);
+                            continue 'outer;
+                        }
+                        println!("\tCould not extend due to conflicts!");
+                    }
+                }
+            }
+        }
+
+        // TODO: do the same for bus interactions
+
+        println!("Mappings for PCP {}: {:#?}", idx, mappings);
+
+        // now assign new ids for all references in the PCP
+        let mut curr_id = 0;
+        let mut new_poly_id = |old_id: u64| {
+            if let Some(id) = mappings.get(&old_id) {
+                return *id;
+            }
+            // if the new id is a target in the mappings, find a new id not yet in the mapping
+            while mappings.values().any(|id| curr_id == *id) {
+                curr_id += 1;
+            }
+            assert!(mappings.values().all(|id| curr_id != *id), "New id {} already exists in mappings: {:?}", curr_id, mappings);
+            assert!(!mappings.contains_key(&old_id));
+            println!("Assigning new id {} for old id {}", curr_id, old_id);
+            let id = curr_id;
+            mappings.insert(old_id, id);
+            curr_id += 1;
+            id
+        };
+
+        pcp.machine.pre_visit_expressions_mut(&mut |expr| {
+            if let AlgebraicExpression::Reference(r) = expr {
+                if r.poly_id.ptype == PolynomialType::Committed {
+                    r.poly_id.id = new_poly_id(r.poly_id.id);
+                }
+            }
+        });
+        println!("----- assigning instr subs ------");
+        pcp.original_instructions.iter_mut().for_each(|instr| {
+            instr.subs.iter_mut().for_each(|sub| {
+                *sub = new_poly_id(*sub);
+            });
+        });
+
+        println!("----- assigning is valid ------");
+        pcp.is_valid_column.id.id = new_poly_id(pcp.is_valid_column.id.id);
+
+        self.pcps.push(pcp.clone());
+    }
+}
+
+
 
 // perform air stacking of multiple precompiles into stacked precompiles.
 fn air_stacking<P: IntoOpenVm>(
@@ -444,9 +657,9 @@ fn air_stacking<P: IntoOpenVm>(
     assert!(!extensions.is_empty());
 
     extensions.iter_mut().for_each(|ext| {
-        ext.machine.constraints.sort();
         ext.machine.constraints.iter_mut().for_each(|c| canonicalize_expression(&mut c.expr));
-        compact_ids(ext)
+        // ext.machine.constraints.sort();
+        // compact_ids(ext)
     });
 
     // create apc groups by number of columns
@@ -456,6 +669,17 @@ fn air_stacking<P: IntoOpenVm>(
         groups.entry(idx).or_default().push(pcp);
     }
 
+    // sort each group by number of columns
+    groups.values_mut().for_each(|g| {
+        // assign largest pcp first
+        g.sort_by(|pcp1, pcp2| pcp2.machine.unique_columns().count().cmp(&pcp1.machine.unique_columns().count()));
+        let mut column_assigner = ColumnAssigner::default();
+        g.iter_mut().for_each(|ext| {
+            column_assigner.assign_pcp(ext);
+            compact_ids(ext); // this should undo the above or sth is very wrong?
+        });
+    });
+
     let mut result = vec![];
 
     for mut extensions in groups.into_values() {
@@ -490,6 +714,7 @@ fn air_stacking<P: IntoOpenVm>(
         let mut is_valid_sum: Option<AlgebraicExpression<P, AlgebraicReference>> = None;
 
         for (idx, pcp) in extensions.iter_mut().enumerate() {
+
             // is_valid columns cannot be shared between precompiles. Here we do
             // their remapping into exclusive columns.
             let is_valid_new_id = is_valid_start + idx as u64;
@@ -498,6 +723,7 @@ fn air_stacking<P: IntoOpenVm>(
 
             // remap is_valid column in constraints and interactions
             let mut remapped = pcp.machine.clone();
+
             remapped.pre_visit_expressions_mut(&mut |expr| {
                 if let AlgebraicExpression::Reference(r) = expr {
                     assert!(r.poly_id.id <= is_valid_start);
@@ -542,7 +768,7 @@ fn air_stacking<P: IntoOpenVm>(
         }
 
         println!("before joining constraints: {}", stacked_constraints.len());
-        let mut stacked_constraints = join_constraints(stacked_constraints);
+        // let mut stacked_constraints = join_constraints(stacked_constraints);
         stacked_constraints.sort();
         println!("after joining constraints: {}", stacked_constraints.len());
 
@@ -553,9 +779,9 @@ fn air_stacking<P: IntoOpenVm>(
         stacked_constraints.push(one_hot_is_valid.into());
 
         println!("interaction count before: {}", interactions_by_machine.iter().flatten().count());
-        let mut stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
+        // let mut stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
+        let mut stacked_interactions = merge_bus_interactions(interactions_by_machine);
         stacked_interactions.sort();
-        // let stacked_interactions = merge_bus_interactions(interactions_by_machine);
 
         let machine = SymbolicMachine {
             constraints: stacked_constraints,

From 13850f669bcaebc49a34d2566059bdd96b8e7c5f Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Tue, 3 Jun 2025 19:16:37 -0300
Subject: [PATCH 12/39] seems to work

---
 openvm/Cargo.toml           |  2 ++
 openvm/src/customize_exe.rs | 50 +++++++++++++++++++++----------------
 2 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/openvm/Cargo.toml b/openvm/Cargo.toml
index 25e6b5a9a8..0fde7559a5 100644
--- a/openvm/Cargo.toml
+++ b/openvm/Cargo.toml
@@ -68,6 +68,8 @@ log = "0.4.17"
 serde_cbor = "0.11.2"
 struct-reflection = { git = "https://github.com/gzanitti/struct-reflection-rs.git" }
 
+bimap = "0.6.3"
+
 [dev-dependencies]
 test-log = { version = "0.2.17", features = ["trace"] }
 pretty_assertions = "1.4.0"
diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 379d7d3d3e..3e18d8189f 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -440,8 +440,8 @@ fn has_same_structure<T: PrimeField32>(expr1: &AlgebraicExpression<T, AlgebraicR
     }
 }
 
-fn create_mapping<T: PrimeField32>(from: &AlgebraicExpression<T>, to: &AlgebraicExpression<T>) -> BTreeMap<u64, u64> {
-    fn create_mapping_inner<T: PrimeField32>(from: &AlgebraicExpression<T>, to: &AlgebraicExpression<T>) -> BTreeMap<u64, u64> {
+fn create_mapping<T: PrimeField32>(from: &AlgebraicExpression<T>, to: &AlgebraicExpression<T>) -> BiMap<u64, u64> {
+    fn create_mapping_inner<T: PrimeField32>(from: &AlgebraicExpression<T>, to: &AlgebraicExpression<T>) -> BiMap<u64, u64> {
         match (from, to) {
             (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left1, op: op1, right: right1 }),
              AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left2, op: op2, right: right2 })) => {
@@ -458,11 +458,11 @@ fn create_mapping<T: PrimeField32>(from: &AlgebraicExpression<T>, to: &Algebraic
                 create_mapping_inner(&expr1, &expr2)
             }
             (AlgebraicExpression::Reference(from), AlgebraicExpression::Reference(to)) => {
-                let mappings = BTreeMap::from([(from.poly_id.id, to.poly_id.id)]);
-                // ref2.poly_id.id = ref1.poly_id.id;
+                let mut mappings = BiMap::new();
+                mappings.insert(from.poly_id.id, to.poly_id.id);
                 mappings
             }
-            _ => BTreeMap::new(),
+            _ => BiMap::new(),
         }
     }
 
@@ -509,15 +509,20 @@ fn create_mapping<T: PrimeField32>(from: &AlgebraicExpression<T>, to: &Algebraic
 // }
 
 fn extend_if_no_conflicts(
-    mappings: &mut BTreeMap<u64, u64>,
-    new_mappings: BTreeMap<u64, u64>,
+    mappings: &mut BiMap<u64, u64>,
+    new_mappings: BiMap<u64, u64>,
 ) -> bool {
     for (k, v) in &new_mappings {
-        if let Some(existing_v) = mappings.get(k) {
+        if let Some(existing_v) = mappings.get_by_left(k) {
             if existing_v != v {
                 return false; // conflict found
             }
         }
+        if let Some(existing_k) = mappings.get_by_right(v) {
+            if existing_k != k {
+                return false; // conflict found
+            }
+        }
     }
     mappings.extend(new_mappings);
     true
@@ -569,7 +574,7 @@ impl<F: PrimeField32> ColumnAssigner<F> {
     ) {
         let idx = self.pcps.len();
         println!("Assinging PCP {}", idx);
-        let mut mappings = BTreeMap::new();
+        let mut mappings = BiMap::new();
         // for each constraint, we want to find if there is a constraint in
         // another machine with the same structure that we can assign the same
         // ids.
@@ -585,9 +590,9 @@ impl<F: PrimeField32> ColumnAssigner<F> {
                                  &expr_poly_id_by_order(c.expr.clone()),
                                  &expr_poly_id_by_order(c2.expr.clone()));
 
-                        println!("Found same structure (ids): \n\t{}\n\n\t{}",
-                                 &expr_orig_poly_ids(c.expr.clone()),
-                                 &expr_orig_poly_ids(c2.expr.clone()));
+                        // println!("Found same structure (ids): \n\t{}\n\n\t{}",
+                        //          &expr_orig_poly_ids(c.expr.clone()),
+                        //          &expr_orig_poly_ids(c2.expr.clone()));
 
                         // assign the same ids
                         let new_mappings = create_mapping(&c.expr, &c2.expr);
@@ -595,7 +600,7 @@ impl<F: PrimeField32> ColumnAssigner<F> {
                             &mut mappings,
                             new_mappings,
                         ) {
-                            println!("\tMappings: {:#?}", mappings);
+                            // println!("\tMappings: {:#?}", mappings);
                             continue 'outer;
                         }
                         println!("\tCould not extend due to conflicts!");
@@ -611,16 +616,17 @@ impl<F: PrimeField32> ColumnAssigner<F> {
         // now assign new ids for all references in the PCP
         let mut curr_id = 0;
         let mut new_poly_id = |old_id: u64| {
-            if let Some(id) = mappings.get(&old_id) {
+            if let Some(id) = mappings.get_by_left(&old_id) {
                 return *id;
             }
             // if the new id is a target in the mappings, find a new id not yet in the mapping
-            while mappings.values().any(|id| curr_id == *id) {
+            while mappings.get_by_right(&curr_id).is_some() {
                 curr_id += 1;
             }
-            assert!(mappings.values().all(|id| curr_id != *id), "New id {} already exists in mappings: {:?}", curr_id, mappings);
-            assert!(!mappings.contains_key(&old_id));
-            println!("Assigning new id {} for old id {}", curr_id, old_id);
+            assert!(mappings.get_by_right(&curr_id).is_none(), "New id {} already exists in mappings: {:?}", curr_id, mappings);
+            // assert!(mappings.values().all(|id| curr_id != *id), "New id {} already exists in mappings: {:?}", curr_id, mappings);
+            assert!(!mappings.contains_left(&old_id));
+            // println!("Assigning new id {} for old id {}", curr_id, old_id);
             let id = curr_id;
             mappings.insert(old_id, id);
             curr_id += 1;
@@ -634,14 +640,14 @@ impl<F: PrimeField32> ColumnAssigner<F> {
                 }
             }
         });
-        println!("----- assigning instr subs ------");
+        // println!("----- assigning instr subs ------");
         pcp.original_instructions.iter_mut().for_each(|instr| {
             instr.subs.iter_mut().for_each(|sub| {
                 *sub = new_poly_id(*sub);
             });
         });
 
-        println!("----- assigning is valid ------");
+        // println!("----- assigning is valid ------");
         pcp.is_valid_column.id.id = new_poly_id(pcp.is_valid_column.id.id);
 
         self.pcps.push(pcp.clone());
@@ -676,7 +682,7 @@ fn air_stacking<P: IntoOpenVm>(
         let mut column_assigner = ColumnAssigner::default();
         g.iter_mut().for_each(|ext| {
             column_assigner.assign_pcp(ext);
-            compact_ids(ext); // this should undo the above or sth is very wrong?
+            // compact_ids(ext); // this should undo the above or sth is very wrong?
         });
     });
 
@@ -768,7 +774,7 @@ fn air_stacking<P: IntoOpenVm>(
         }
 
         println!("before joining constraints: {}", stacked_constraints.len());
-        // let mut stacked_constraints = join_constraints(stacked_constraints);
+        let mut stacked_constraints = join_constraints(stacked_constraints);
         stacked_constraints.sort();
         println!("after joining constraints: {}", stacked_constraints.len());
 

From 669b9747243ef0a9590d6e776762fb9419ff098f Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 4 Jun 2025 18:29:37 -0300
Subject: [PATCH 13/39] try handle bus interactions

---
 openvm/src/customize_exe.rs | 236 +++++++++++++++++++++++++++---------
 1 file changed, 182 insertions(+), 54 deletions(-)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 3e18d8189f..541f705e11 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -336,6 +336,140 @@ fn merge_bus_interactions_simple<P: IntoOpenVm>(interactions: Vec<Vec<SymbolicBu
     interactions.into_iter().flatten().collect_vec()
 }
 
+fn merge_bus_interactions2<P: IntoOpenVm>(
+    interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>
+) -> Vec<SymbolicBusInteraction<P>> {
+    // split interactions by bus/args len
+    let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> = Default::default();
+
+    for interactions in interactions_by_machine {
+        let interactions_by_bus_this_machine = interactions.into_iter()
+            // we group by bus id and number of args
+            .into_group_map_by(|interaction| (interaction.id, interaction.args.len()));
+        for (k,v) in interactions_by_bus_this_machine {
+            let e = interactions_by_bus.entry(k).or_default();
+            e.push(v);
+        }
+    }
+    interactions_by_bus.values_mut().for_each(|interactions_by_machine| {
+        interactions_by_machine.sort_by_key(|interactions| interactions.len());
+    });
+
+    let mut result = vec![];
+    for mut interactions_by_machine in interactions_by_bus.into_values() {
+        // to_merge is a vec of vecs, each inner vec is a set of interactions to be merged
+        let mut to_merge = interactions_by_machine.pop().unwrap().into_iter().map(|i| vec![i]).collect_vec();
+        for machine_interactions in interactions_by_machine {
+            let mut used = BTreeSet::new();
+            'outer: for i in machine_interactions {
+                // try to find a call with an exact match
+                for (idx, to_merge_set) in to_merge.iter_mut().enumerate() {
+                    if used.contains(&idx) {
+                        continue;
+                    }
+                    let i2 = to_merge_set.get(0).unwrap();
+                    let all_args_same_structure = i.args.iter()
+                        .zip_eq(i2.args.iter())
+                        .all(|(a1, a2)| has_same_structure(strip_is_valid(&a1),
+                                                           strip_is_valid(&a2)));
+                    if all_args_same_structure {
+                        println!("EXACT MATCH");
+                        to_merge_set.push(i);
+                        used.insert(idx);
+                        continue 'outer;
+                    }
+                }
+                // didn't find exact match, try one where some args match
+                for (idx, to_merge_set) in to_merge.iter_mut().enumerate() {
+                    if used.contains(&idx) {
+                        continue;
+                    }
+                    // check all args have the same structure
+                    let i2 = to_merge_set.get(0).unwrap();
+                    let some_args_same_structure = i.args.iter()
+                        .zip_eq(i2.args.iter())
+                        .any(|(a1, a2)| has_same_structure(strip_is_valid(&a1),
+                                                           strip_is_valid(&a2)));
+                    if some_args_same_structure {
+                        println!("PARTIAL MATCH");
+                        to_merge_set.push(i);
+                        used.insert(idx);
+                        continue 'outer;
+                    }
+                }
+                // just pick the first unused one
+                for (idx, to_merge_set) in to_merge.iter_mut().enumerate() {
+                    if !used.contains(&idx) {
+                        println!("NO MATCH");
+                        to_merge_set.push(i);
+                        used.insert(idx);
+                        continue 'outer;
+                    }
+                }
+                unreachable!("could not find any bus interactions to merge with, but should");
+            }
+        }
+
+        fn merge_args<P: IntoOpenVm>(arg1: AlgebraicExpression<P, AlgebraicReference>, arg2: AlgebraicExpression<P, AlgebraicReference>) -> AlgebraicExpression<P, AlgebraicReference> {
+            match arg1 {
+                // expr * is_valid
+                AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left1, op: AlgebraicBinaryOperator::Mul, right: is_valid1 }) => {
+                    let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left2, op: AlgebraicBinaryOperator::Mul, right: is_valid2 }) = arg2 else {
+                        panic!("Expected binary operation for arg2, got: {arg2}");
+                    };
+                    if has_same_structure(&left1, &left2) {
+                        *left1 * (*is_valid1 + *is_valid2)
+                    } else {
+                        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                            left: left1,
+                            op: AlgebraicBinaryOperator::Mul,
+                            right: is_valid1,
+                        })
+                            +
+                        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                            left: left2,
+                            op: AlgebraicBinaryOperator::Mul,
+                            right: is_valid2,
+                        })
+                    }
+                }
+                // exprA * is_validA + exprB * is_validB + ...
+                AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left:_, op: AlgebraicBinaryOperator::Add, right:_ }) => {
+                    arg1 + arg2
+                }
+                _ => unreachable!()
+            }
+        }
+
+        fn simplify_arg<P: IntoOpenVm>(arg: AlgebraicExpression<P, AlgebraicReference>) -> AlgebraicExpression<P, AlgebraicReference> {
+            // expr * is_valid_expr
+            if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op: AlgebraicBinaryOperator::Mul, right: _is_valid_expr }) = arg {
+                *left
+            } else {
+                // sum of expr * is_validN, can't simplify
+                arg
+            }
+        }
+
+        // merge each set of interactions
+        for set in to_merge {
+            let id = set[0].id;
+            let mult = set.iter().map(|i| i.mult.clone()).reduce(|a, b| a + b).unwrap();
+            let args = set.into_iter().map(|i| i.args)
+                .reduce(|a, b| a.into_iter().zip_eq(b).map(|(a1, a2)| merge_args(a1, a2)).collect())
+                .unwrap();
+            let args = args.into_iter().map(simplify_arg).collect_vec();
+            result.push(SymbolicBusInteraction {
+                id,
+                args,
+                mult,
+            });
+        }
+    }
+
+    result
+}
+
 fn merge_bus_interactions<P: IntoOpenVm>(interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>) -> Vec<SymbolicBusInteraction<P>> {
     let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> = Default::default();
 
@@ -422,7 +556,7 @@ fn join_constraints<P: IntoOpenVm>(mut constraints: Vec<SymbolicConstraint<P>>)
     result
 }
 
-fn has_same_structure<T: PrimeField32>(expr1: &AlgebraicExpression<T, AlgebraicReference>, expr2: &AlgebraicExpression<T, AlgebraicReference>) -> bool {
+fn has_same_structure<P: IntoOpenVm>(expr1: &AlgebraicExpression<P, AlgebraicReference>, expr2: &AlgebraicExpression<P, AlgebraicReference>) -> bool {
     match (expr1, expr2) {
         (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation{ left: left1, op: op1, right: right1 }),
          AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation{ left: left2, op: op2, right: right2 })) => {
@@ -440,8 +574,17 @@ fn has_same_structure<T: PrimeField32>(expr1: &AlgebraicExpression<T, AlgebraicR
     }
 }
 
-fn create_mapping<T: PrimeField32>(from: &AlgebraicExpression<T>, to: &AlgebraicExpression<T>) -> BiMap<u64, u64> {
-    fn create_mapping_inner<T: PrimeField32>(from: &AlgebraicExpression<T>, to: &AlgebraicExpression<T>) -> BiMap<u64, u64> {
+// assumes expression is of the form `some_expr * is_valid_expr`.
+fn strip_is_valid<P: IntoOpenVm>(expr: &AlgebraicExpression<P, AlgebraicReference>) -> &AlgebraicExpression<P, AlgebraicReference> {
+    let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op: AlgebraicBinaryOperator::Mul, right: _ }) = expr else {
+        panic!("Expression is not a binary operation with Mul operator: {expr}");
+    };
+    left.as_ref()
+}
+
+
+fn create_mapping<P: IntoOpenVm>(from: &AlgebraicExpression<P, AlgebraicReference>, to: &AlgebraicExpression<P, AlgebraicReference>) -> BiMap<u64, u64> {
+    fn create_mapping_inner<P: IntoOpenVm>(from: &AlgebraicExpression<P, AlgebraicReference>, to: &AlgebraicExpression<P, AlgebraicReference>) -> BiMap<u64, u64> {
         match (from, to) {
             (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left1, op: op1, right: right1 }),
              AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left2, op: op2, right: right2 })) => {
@@ -469,45 +612,6 @@ fn create_mapping<T: PrimeField32>(from: &AlgebraicExpression<T>, to: &Algebraic
     create_mapping_inner(from, to)
 }
 
-/// assumes all expressions of the form `Y * is_valid_expr`
-// fn create_mapping<T: PrimeField32>(expr1: &AlgebraicExpression<T>, expr2: &AlgebraicExpression<T>) -> BTreeMap<u64, u64> {
-//     fn create_mapping_inner<T: PrimeField32>(expr1: &Box<AlgebraicExpression<T>>, expr2: &Box<AlgebraicExpression<T>>) -> BTreeMap<u64, u64> {
-//         match (expr1.as_ref(), expr2.as_ref()) {
-//             (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left1, op: op1, right: right1 }),
-//              AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left2, op: op2, right: right2 })) => {
-//                 assert_eq!(op1, op2);
-//                 let mut left = create_mapping_inner(left1, left2);
-//                 let right = create_mapping_inner(right1, right2);
-//                 // assert both sides don't conflict
-//                 assert!(extend_if_no_conflicts(&mut left, right), "conflict in mapping: {expr1} => {expr2}");
-//                 left
-//             }
-//             (AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op1, expr: expr1 }),
-//              AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op2, expr: expr2 })) => {
-//                 assert_eq!(op1, op2);
-//                 create_mapping_inner(expr1, expr2)
-//             }
-//             (AlgebraicExpression::Reference(ref1), AlgebraicExpression::Reference(ref2)) => {
-//                 let mappings = BTreeMap::from([(ref2.poly_id.id, ref1.poly_id.id)]); // CHECK
-//                 mappings
-//             }
-//             (AlgebraicExpression::Number(n1), AlgebraicExpression::Number(n2)) => {
-//                 assert_eq!(n1, n2, "numbers do not match: {n1} != {n2}");
-//                 BTreeMap::new() // no mapping needed for numbers
-//             }
-//             _ => unreachable!("invalid expression mapping: {expr1} => {expr2}"),
-//         }
-//     }
-//
-//     let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: expr1, op: _, right: _is_valid }) = expr1 else {
-//         panic!("invalid expression mapping: {expr1} => {expr2}");
-//     };
-//     let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: expr2, op: _, right: _is_valid }) = expr2 else {
-//         panic!("invalid expression: {expr1} => {expr2}");
-//     };
-//     create_mapping_inner(expr1, expr2)
-// }
-
 fn extend_if_no_conflicts(
     mappings: &mut BiMap<u64, u64>,
     new_mappings: BiMap<u64, u64>,
@@ -531,7 +635,7 @@ fn extend_if_no_conflicts(
 /// changes poly_id to match the order in which they are encountered in the expression.
 /// This allows us to compare two expressions for the same "structure".
 /// e.g., a+a == b+b and a+a != a+b
-fn expr_poly_id_by_order<F: PrimeField32>(mut expr: AlgebraicExpression<F>) -> AlgebraicExpression<F> {
+fn expr_poly_id_by_order<P: IntoOpenVm>(mut expr: AlgebraicExpression<P, AlgebraicReference>) -> AlgebraicExpression<P, AlgebraicReference> {
     let mut curr_id = 0;
     let mut id_map: HashMap<u64, u64> = Default::default();
     let mut new_poly_id = |old_id: u64| {
@@ -550,7 +654,7 @@ fn expr_poly_id_by_order<F: PrimeField32>(mut expr: AlgebraicExpression<F>) -> A
     expr
 }
 
-fn expr_orig_poly_ids<F: PrimeField32>(mut expr: AlgebraicExpression<F>) -> AlgebraicExpression<F> {
+fn expr_orig_poly_ids<P: IntoOpenVm>(mut expr: AlgebraicExpression<P, AlgebraicReference>) -> AlgebraicExpression<P, AlgebraicReference> {
     expr.pre_visit_expressions_mut(&mut |e| {
         if let AlgebraicExpression::Reference(r) = e {
             r.name = format!("col_{}", r.poly_id.id); // this one not needed, just for printing
@@ -560,17 +664,17 @@ fn expr_orig_poly_ids<F: PrimeField32>(mut expr: AlgebraicExpression<F>) -> Alge
 }
 
 #[derive(Default)]
-struct ColumnAssigner<F: PrimeField32> {
+struct ColumnAssigner<P: IntoOpenVm> {
     // original PCPs
-    pcps: Vec<PowdrPrecompile<F>>,
+    pcps: Vec<PowdrPrecompile<P>>,
     // id mapping for each PCP (by idx)
     mappings: BTreeMap<usize, BTreeMap<u64, u64>>,
 }
 
-impl<F: PrimeField32> ColumnAssigner<F> {
+impl<P: IntoOpenVm> ColumnAssigner<P> {
     fn assign_pcp(
         &mut self,
-        pcp: &mut PowdrPrecompile<F>,
+        pcp: &mut PowdrPrecompile<P>,
     ) {
         let idx = self.pcps.len();
         println!("Assinging PCP {}", idx);
@@ -581,7 +685,34 @@ impl<F: PrimeField32> ColumnAssigner<F> {
         // There may be multiple such constraints, so we try them all (as some
         // of the ids in the current constraint may already be assigned)
 
-        'outer: for c in &mut pcp.machine.constraints {
+        // do the same for bus interactions
+        'outer: for b in &pcp.machine.bus_interactions {
+            for pcp2 in self.pcps.iter() {
+                for b2 in &pcp2.machine.bus_interactions {
+                    if b.id == b2.id && b.args.len() == b2.args.len() {
+                        // let all_args_same_structure = b.args.iter()
+                        //     .zip_eq(b2.args.iter())
+                        //     .all(|(a1, a2)| has_same_structure(&expr_poly_id_by_order(a1.clone()), &expr_poly_id_by_order(a2.clone())));
+                        // if all_args_same_structure {
+                        // println!("Found same bus interaction: {b:?} == {b2:?}");
+                        for (arg1, arg2) in b.args.iter().zip_eq(b2.args.iter()) {
+                            let new_mappings = create_mapping(&arg1, &arg2);
+                            if extend_if_no_conflicts(
+                                &mut mappings,
+                                new_mappings,
+                            ) {
+                                println!("\tFound mapping for bus interaction");
+                                continue 'outer;
+                            }
+                            println!("\tCould not extend due to conflicts!");
+                        }
+                        // }
+                    }
+                }
+            }
+        }
+
+        'outer: for c in &pcp.machine.constraints {
             for pcp2 in self.pcps.iter() {
                 for c2 in &pcp2.machine.constraints {
                     if has_same_structure(&expr_poly_id_by_order(c.expr.clone()),
@@ -609,8 +740,6 @@ impl<F: PrimeField32> ColumnAssigner<F> {
             }
         }
 
-        // TODO: do the same for bus interactions
-
         println!("Mappings for PCP {}: {:#?}", idx, mappings);
 
         // now assign new ids for all references in the PCP
@@ -624,7 +753,6 @@ impl<F: PrimeField32> ColumnAssigner<F> {
                 curr_id += 1;
             }
             assert!(mappings.get_by_right(&curr_id).is_none(), "New id {} already exists in mappings: {:?}", curr_id, mappings);
-            // assert!(mappings.values().all(|id| curr_id != *id), "New id {} already exists in mappings: {:?}", curr_id, mappings);
             assert!(!mappings.contains_left(&old_id));
             // println!("Assigning new id {} for old id {}", curr_id, old_id);
             let id = curr_id;
@@ -682,7 +810,6 @@ fn air_stacking<P: IntoOpenVm>(
         let mut column_assigner = ColumnAssigner::default();
         g.iter_mut().for_each(|ext| {
             column_assigner.assign_pcp(ext);
-            // compact_ids(ext); // this should undo the above or sth is very wrong?
         });
     });
 
@@ -786,7 +913,8 @@ fn air_stacking<P: IntoOpenVm>(
 
         println!("interaction count before: {}", interactions_by_machine.iter().flatten().count());
         // let mut stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
-        let mut stacked_interactions = merge_bus_interactions(interactions_by_machine);
+        // let mut stacked_interactions = merge_bus_interactions(interactions_by_machine);
+        let mut stacked_interactions = merge_bus_interactions2(interactions_by_machine);
         stacked_interactions.sort();
 
         let machine = SymbolicMachine {

From 9dcff42ce5b2ad1a92737f002a9921ddb8c0990a Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Thu, 5 Jun 2025 11:57:29 -0300
Subject: [PATCH 14/39] stuff

---
 openvm/Cargo.toml           |  1 +
 openvm/src/customize_exe.rs | 65 ++++++++++++++++++++-----------------
 2 files changed, 37 insertions(+), 29 deletions(-)

diff --git a/openvm/Cargo.toml b/openvm/Cargo.toml
index 0fde7559a5..3930838e72 100644
--- a/openvm/Cargo.toml
+++ b/openvm/Cargo.toml
@@ -50,6 +50,7 @@ powdr-number.workspace = true
 powdr-riscv-elf.workspace = true
 powdr-autoprecompiles.workspace = true
 powdr-constraint-solver.workspace = true
+powdr-pilopt.workspace = true
 
 eyre = "0.6.12"
 serde = "1.0.217"
diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 541f705e11..8be026880d 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -28,6 +28,7 @@ use powdr_autoprecompiles::{
     SymbolicBusInteraction, SymbolicInstructionStatement, SymbolicMachine,
 };
 use powdr_number::FieldElement;
+use powdr_pilopt::simplify_expression;
 
 use crate::bus_interaction_handler::{BusMap, OpenVmBusInteractionHandler};
 use crate::instruction_formatter::openvm_instruction_formatter;
@@ -361,6 +362,7 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
         let mut to_merge = interactions_by_machine.pop().unwrap().into_iter().map(|i| vec![i]).collect_vec();
         for machine_interactions in interactions_by_machine {
             let mut used = BTreeSet::new();
+            let mut try_partial_match = vec![];
             'outer: for i in machine_interactions {
                 // try to find a call with an exact match
                 for (idx, to_merge_set) in to_merge.iter_mut().enumerate() {
@@ -379,6 +381,11 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
                         continue 'outer;
                     }
                 }
+                try_partial_match.push(i);
+            }
+
+            let mut no_match = vec![];
+            'outer: for i in try_partial_match {
                 // didn't find exact match, try one where some args match
                 for (idx, to_merge_set) in to_merge.iter_mut().enumerate() {
                     if used.contains(&idx) {
@@ -397,6 +404,10 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
                         continue 'outer;
                     }
                 }
+                no_match.push(i);
+            }
+
+            'outer: for i in no_match {
                 // just pick the first unused one
                 for (idx, to_merge_set) in to_merge.iter_mut().enumerate() {
                     if !used.contains(&idx) {
@@ -454,11 +465,11 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
         // merge each set of interactions
         for set in to_merge {
             let id = set[0].id;
-            let mult = set.iter().map(|i| i.mult.clone()).reduce(|a, b| a + b).unwrap();
+            let mult = simplify_expression(set.iter().map(|i| i.mult.clone()).reduce(|a, b| a + b).unwrap());
             let args = set.into_iter().map(|i| i.args)
                 .reduce(|a, b| a.into_iter().zip_eq(b).map(|(a1, a2)| merge_args(a1, a2)).collect())
                 .unwrap();
-            let args = args.into_iter().map(simplify_arg).collect_vec();
+            let args = args.into_iter().map(|arg| simplify_expression(simplify_arg(arg))).collect_vec();
             result.push(SymbolicBusInteraction {
                 id,
                 args,
@@ -685,33 +696,6 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
         // There may be multiple such constraints, so we try them all (as some
         // of the ids in the current constraint may already be assigned)
 
-        // do the same for bus interactions
-        'outer: for b in &pcp.machine.bus_interactions {
-            for pcp2 in self.pcps.iter() {
-                for b2 in &pcp2.machine.bus_interactions {
-                    if b.id == b2.id && b.args.len() == b2.args.len() {
-                        // let all_args_same_structure = b.args.iter()
-                        //     .zip_eq(b2.args.iter())
-                        //     .all(|(a1, a2)| has_same_structure(&expr_poly_id_by_order(a1.clone()), &expr_poly_id_by_order(a2.clone())));
-                        // if all_args_same_structure {
-                        // println!("Found same bus interaction: {b:?} == {b2:?}");
-                        for (arg1, arg2) in b.args.iter().zip_eq(b2.args.iter()) {
-                            let new_mappings = create_mapping(&arg1, &arg2);
-                            if extend_if_no_conflicts(
-                                &mut mappings,
-                                new_mappings,
-                            ) {
-                                println!("\tFound mapping for bus interaction");
-                                continue 'outer;
-                            }
-                            println!("\tCould not extend due to conflicts!");
-                        }
-                        // }
-                    }
-                }
-            }
-        }
-
         'outer: for c in &pcp.machine.constraints {
             for pcp2 in self.pcps.iter() {
                 for c2 in &pcp2.machine.constraints {
@@ -740,6 +724,29 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
             }
         }
 
+        // do the same for bus interactions
+        'outer: for b in &pcp.machine.bus_interactions {
+            for pcp2 in self.pcps.iter() {
+                for b2 in &pcp2.machine.bus_interactions {
+                    if b.id == b2.id && b.args.len() == b2.args.len() {
+                        let all_args_same_structure = b.args.iter()
+                            .zip_eq(b2.args.iter())
+                            .all(|(a1, a2)| has_same_structure(&expr_poly_id_by_order(a1.clone()), &expr_poly_id_by_order(a2.clone())));
+                        if all_args_same_structure {
+                            for (arg1, arg2) in b.args.iter().zip_eq(b2.args.iter()) {
+                                let new_mappings = create_mapping(&arg1, &arg2);
+                                extend_if_no_conflicts(
+                                    &mut mappings,
+                                    new_mappings,
+                                );
+                            }
+                            continue 'outer;
+                        }
+                    }
+                }
+            }
+        }
+
         println!("Mappings for PCP {}: {:#?}", idx, mappings);
 
         // now assign new ids for all references in the PCP

From c8918f509eb9f9ad3b1d0b4a782556ebe65a1b62 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Tue, 10 Jun 2025 09:22:31 -0300
Subject: [PATCH 15/39] move air_stacking to own file

---
 openvm/src/customize_exe.rs              | 222 +--------------------
 openvm/src/customize_exe/air_stacking.rs | 233 +++++++++++++++++++++++
 2 files changed, 236 insertions(+), 219 deletions(-)
 create mode 100644 openvm/src/customize_exe/air_stacking.rs

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 718cf1348a..db439ba8ce 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -32,6 +32,9 @@ use crate::{
     utils::symbolic_to_algebraic,
 };
 
+mod air_stacking;
+use air_stacking::air_stacking;
+
 pub const OPENVM_DEGREE_BOUND: usize = 5;
 
 // TODO: read this from program
@@ -298,225 +301,6 @@ pub fn customize<P: IntoOpenVm>(
     )
 }
 
-/// make PolyIDs start from zero, sequential and compact
-fn compact_ids<P: IntoOpenVm>(pcp: &mut PowdrPrecompile<P>) {
-    let mut curr_id = 0;
-    let mut id_map: HashMap<u64, u64> = Default::default();
-    let mut new_poly_id = |old_id: u64| {
-        *id_map.entry(old_id).or_insert_with(|| {
-            let id = curr_id;
-            // println!("remapping poly id {} to {}", old_id, curr_id);
-            curr_id += 1;
-            id
-        })
-    };
-    pcp.machine.pre_visit_expressions_mut(&mut |expr| {
-        if let AlgebraicExpression::Reference(r) = expr {
-            if r.poly_id.ptype == PolynomialType::Committed {
-                r.poly_id.id = new_poly_id(r.poly_id.id);
-            }
-        }
-    });
-    pcp.original_instructions.iter_mut().for_each(|instr| {
-        instr.subs.iter_mut().for_each(|sub| {
-            *sub = new_poly_id(*sub);
-        });
-    });
-
-    pcp.is_valid_column.id.id = new_poly_id(pcp.is_valid_column.id.id);
-}
-
-fn merge_bus_interactions_simple<P: IntoOpenVm>(interactions: Vec<Vec<SymbolicBusInteraction<P>>>) -> Vec<SymbolicBusInteraction<P>> {
-    interactions.into_iter().flatten().collect_vec()
-}
-
-fn merge_bus_interactions<P: IntoOpenVm>(interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>) -> Vec<SymbolicBusInteraction<P>> {
-    let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> = Default::default();
-
-    for interactions in interactions_by_machine {
-        let interactions_by_bus_this_machine = interactions.into_iter()
-            // we group by bus id and number of args
-            .into_group_map_by(|interaction| (interaction.id, interaction.args.len()));
-        for (k,v) in interactions_by_bus_this_machine {
-            let e = interactions_by_bus.entry(k).or_default();
-            e.push(v);
-        }
-    }
-    let mut to_merge = vec![];
-    for interactions_per_machine in interactions_by_bus.into_values() {
-        // this is doing a "zip longest" among the interactions from each machine
-        // (for a given bus), and saving the elements that were picked
-        // together to be merged later
-        let mut idx = 0;
-        loop {
-            let mut items = vec![];
-            for machine_interactions in interactions_per_machine.iter() {
-                if let Some(item) = machine_interactions.get(idx) {
-                    items.push(item.clone());
-                }
-            }
-            if items.is_empty() {
-                break;
-            }
-            idx += 1;
-            to_merge.push(items);
-        }
-    }
-
-    to_merge.into_iter().map(|interactions| {
-        // assert_eq!(interactions.iter().map(|i| &i.kind).unique().count(), 1);
-
-        // add multiplicities together
-        let mut mult = interactions[0].mult.clone();
-        for mult2 in interactions.iter().skip(1).map(|i| i.mult.clone()) {
-            mult = mult + mult2;
-        }
-
-        // add args together
-        let mut args = interactions[0].args.clone();
-        for args2 in interactions.iter().skip(1).map(|i| i.args.clone()) {
-            args = args.into_iter().zip_eq(args2).map(|(a1, a2)| a1 + a2).collect();
-        }
-
-        SymbolicBusInteraction {
-            id: interactions[0].id,
-            mult,
-            args,
-        }
-    }).collect_vec()
-}
-
-// perform air stacking of multiple precompiles into stacked precompiles.
-fn air_stacking<P: IntoOpenVm>(
-    mut extensions: Vec<PowdrPrecompile<P>>,
-) -> Vec<PowdrStackedPrecompile<P>> {
-    assert!(!extensions.is_empty());
-
-    extensions.iter_mut().for_each(compact_ids);
-
-    // create apc groups by number of columns
-    let mut groups: HashMap<usize, Vec<PowdrPrecompile<P>>> = Default::default();
-    for pcp in extensions {
-        let idx = f64::log(pcp.machine.unique_columns().count() as f64, 2.0).floor() as usize;
-        groups.entry(idx).or_default().push(pcp);
-    }
-
-    let mut result = vec![];
-
-    for mut extensions in groups.into_values() {
-        // if there is only one precompile in the group, we can just return it as a stacked precompile
-        if extensions.len() == 1 {
-            println!("not stacking precompile");
-            result.push(PowdrStackedPrecompile {
-                machine: SymbolicMachine {
-                    constraints: extensions[0].machine.constraints.clone(),
-                    bus_interactions: extensions[0].machine.bus_interactions.clone(),
-                },
-                precompiles: extensions
-                    .into_iter()
-                    .map(|p| (p.opcode.clone(), p))
-                    .collect(),
-            });
-            continue;
-        }
-
-        let mut stacked_constraints = vec![];
-        let mut interactions_by_machine = vec![];
-
-        println!("stacking {} precompiles", extensions.len());
-
-        // take the max id in all pcps and add 1.
-        let is_valid_start = 1 + extensions.iter()
-            .flat_map(|pcp| pcp.original_instructions.iter().flat_map(|instr| instr.subs.iter()))
-            .max()
-            .unwrap();
-
-
-        let mut is_valid_sum: Option<AlgebraicExpression<P, AlgebraicReference>> = None;
-
-        for (idx, pcp) in extensions.iter_mut().enumerate() {
-            // is_valid columns cannot be shared between precompiles. Here we do
-            // their remapping into exclusive columns.
-            let is_valid_new_id = is_valid_start + idx as u64;
-
-            println!("\tpcp width: {}", pcp.machine.unique_columns().count());
-
-            // remap is_valid column in constraints and interactions
-            let mut remapped = pcp.machine.clone();
-            remapped.pre_visit_expressions_mut(&mut |expr| {
-                if let AlgebraicExpression::Reference(r) = expr {
-                    assert!(r.poly_id.id <= is_valid_start);
-                    if r.poly_id == pcp.is_valid_column.id {
-                        // we assume each pcp to have a specific column named "is_valid"
-                        assert!(r.name == "is_valid");
-                        r.poly_id.id = is_valid_new_id;
-                        r.name = format!("is_valid_{}", pcp.opcode.global_opcode().as_usize());
-                    } else {
-                        // we need to rename columns here because `unique_columns` relies on both `Column::name` and `Column::id`
-                        r.name = format!("col_{}", r.poly_id.id);
-                    }
-                }
-            });
-
-            // set the is valid column in the original precompile
-            pcp.is_valid_column.id.id = is_valid_new_id;
-
-            let is_valid = AlgebraicExpression::Reference(AlgebraicReference {
-                name: format!("is_valid_{}", pcp.opcode.global_opcode().as_usize()),
-                poly_id: PolyID {
-                    id: is_valid_new_id,
-                    ptype: PolynomialType::Committed,
-                },
-                next: false,
-            });
-
-            // guard interaction payloads so they can be merged later
-            remapped.bus_interactions.iter_mut().for_each(|interaction| {
-                interaction.args.iter_mut().for_each(|arg| {
-                    *arg = arg.clone() * is_valid.clone();
-                });
-            });
-
-
-            is_valid_sum = is_valid_sum
-                .map(|sum| sum + is_valid.clone())
-                .or_else(|| Some(is_valid.clone()));
-
-            stacked_constraints.extend(remapped.constraints);
-            interactions_by_machine.push(remapped.bus_interactions);
-        }
-
-        // enforce only one is_valid is active
-        let one = AlgebraicExpression::Number(P::ONE);
-        let one_hot_is_valid = (one - is_valid_sum.clone().unwrap()) * is_valid_sum.unwrap();
-
-        stacked_constraints.push(one_hot_is_valid.into());
-
-        println!("interaction count before: {}", interactions_by_machine.iter().flatten().count());
-        // let stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
-        let stacked_interactions = merge_bus_interactions(interactions_by_machine);
-
-        let machine = SymbolicMachine {
-            constraints: stacked_constraints,
-            bus_interactions: stacked_interactions,
-        };
-
-        println!("interaction count: {}", machine.bus_interactions.len());
-
-        println!("stacked width: {}", machine.unique_columns().count());
-
-        result.push(PowdrStackedPrecompile {
-            precompiles: extensions
-                .into_iter()
-                .map(|p| (p.opcode.clone(), p))
-                .collect(),
-            machine,
-        });
-    }
-
-    result
-}
-
 
 fn branch_opcodes() -> Vec<usize> {
     let mut opcodes = vec![
diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
new file mode 100644
index 0000000000..2cbf2410fb
--- /dev/null
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -0,0 +1,233 @@
+use std::collections::HashMap;
+
+use powdr_autoprecompiles::{legacy_expression::{AlgebraicExpression, AlgebraicReference, PolyID, PolynomialType}, SymbolicBusInteraction, SymbolicMachine};
+
+use powdr_autoprecompiles::powdr::UniqueColumns;
+
+use powdr_expression::visitors::ExpressionVisitable;
+
+use crate::{powdr_extension::{PowdrPrecompile, PowdrStackedPrecompile}, IntoOpenVm};
+
+use openvm_instructions::LocalOpcode;
+
+use itertools::Itertools;
+
+// perform air stacking of multiple precompiles into stacked precompiles.
+pub fn air_stacking<P: IntoOpenVm>(
+    mut extensions: Vec<PowdrPrecompile<P>>,
+) -> Vec<PowdrStackedPrecompile<P>> {
+    assert!(!extensions.is_empty());
+
+    extensions.iter_mut().for_each(compact_ids);
+
+    // create apc groups by number of columns
+    let mut groups: HashMap<usize, Vec<PowdrPrecompile<P>>> = Default::default();
+    for pcp in extensions {
+        let idx = f64::log(pcp.machine.unique_columns().count() as f64, 2.0).floor() as usize;
+        groups.entry(idx).or_default().push(pcp);
+    }
+
+    let mut result = vec![];
+
+    for mut extensions in groups.into_values() {
+        // if there is only one precompile in the group, we can just return it as a stacked precompile
+        if extensions.len() == 1 {
+            println!("not stacking precompile");
+            result.push(PowdrStackedPrecompile {
+                machine: SymbolicMachine {
+                    constraints: extensions[0].machine.constraints.clone(),
+                    bus_interactions: extensions[0].machine.bus_interactions.clone(),
+                },
+                precompiles: extensions
+                    .into_iter()
+                    .map(|p| (p.opcode.clone(), p))
+                    .collect(),
+            });
+            continue;
+        }
+
+        let mut stacked_constraints = vec![];
+        let mut interactions_by_machine = vec![];
+
+        println!("stacking {} precompiles", extensions.len());
+
+        // take the max id in all pcps and add 1.
+        let is_valid_start = 1 + extensions.iter()
+            .flat_map(|pcp| pcp.original_instructions.iter().flat_map(|instr| instr.subs.iter()))
+            .max()
+            .unwrap();
+
+
+        let mut is_valid_sum: Option<AlgebraicExpression<P>> = None;
+
+        for (idx, pcp) in extensions.iter_mut().enumerate() {
+            // is_valid columns cannot be shared between precompiles. Here we do
+            // their remapping into exclusive columns.
+            let is_valid_new_id = is_valid_start + idx as u64;
+
+            println!("\tpcp width: {}", pcp.machine.unique_columns().count());
+
+            // remap is_valid column in constraints and interactions
+            let mut remapped = pcp.machine.clone();
+            remapped.pre_visit_expressions_mut(&mut |expr| {
+                if let AlgebraicExpression::Reference(r) = expr {
+                    assert!(r.poly_id.id <= is_valid_start);
+                    if r.poly_id == pcp.is_valid_column.id {
+                        // we assume each pcp to have a specific column named "is_valid"
+                        assert!(r.name == "is_valid");
+                        r.poly_id.id = is_valid_new_id;
+                        r.name = format!("is_valid_{}", pcp.opcode.global_opcode().as_usize());
+                    } else {
+                        // we need to rename columns here because `unique_columns` relies on both `Column::name` and `Column::id`
+                        r.name = format!("col_{}", r.poly_id.id);
+                    }
+                }
+            });
+
+            // set the is valid column in the original precompile
+            pcp.is_valid_column.id.id = is_valid_new_id;
+
+            let is_valid = AlgebraicExpression::Reference(AlgebraicReference {
+                name: format!("is_valid_{}", pcp.opcode.global_opcode().as_usize()),
+                poly_id: PolyID {
+                    id: is_valid_new_id,
+                    ptype: PolynomialType::Committed,
+                },
+                next: false,
+            });
+
+            // guard interaction payloads so they can be merged later
+            remapped.bus_interactions.iter_mut().for_each(|interaction| {
+                interaction.args.iter_mut().for_each(|arg| {
+                    *arg = arg.clone() * is_valid.clone();
+                });
+            });
+
+
+            is_valid_sum = is_valid_sum
+                .map(|sum| sum + is_valid.clone())
+                .or_else(|| Some(is_valid.clone()));
+
+            stacked_constraints.extend(remapped.constraints);
+            interactions_by_machine.push(remapped.bus_interactions);
+        }
+
+        // enforce only one is_valid is active
+        let one = AlgebraicExpression::Number(P::ONE);
+        let one_hot_is_valid = (one - is_valid_sum.clone().unwrap()) * is_valid_sum.unwrap();
+
+        stacked_constraints.push(one_hot_is_valid.into());
+
+        println!("interaction count before: {}", interactions_by_machine.iter().flatten().count());
+        // let stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
+        let stacked_interactions = merge_bus_interactions(interactions_by_machine);
+
+        let machine = SymbolicMachine {
+            constraints: stacked_constraints,
+            bus_interactions: stacked_interactions,
+        };
+
+        println!("interaction count: {}", machine.bus_interactions.len());
+
+        println!("stacked width: {}", machine.unique_columns().count());
+
+        result.push(PowdrStackedPrecompile {
+            precompiles: extensions
+                .into_iter()
+                .map(|p| (p.opcode.clone(), p))
+                .collect(),
+            machine,
+        });
+    }
+
+    result
+}
+
+
+/// make PolyIDs start from zero, sequential and compact
+fn compact_ids<P: IntoOpenVm>(pcp: &mut PowdrPrecompile<P>) {
+    let mut curr_id = 0;
+    let mut id_map: HashMap<u64, u64> = Default::default();
+    let mut new_poly_id = |old_id: u64| {
+        *id_map.entry(old_id).or_insert_with(|| {
+            let id = curr_id;
+            // println!("remapping poly id {} to {}", old_id, curr_id);
+            curr_id += 1;
+            id
+        })
+    };
+    pcp.machine.pre_visit_expressions_mut(&mut |expr| {
+        if let AlgebraicExpression::Reference(r) = expr {
+            if r.poly_id.ptype == PolynomialType::Committed {
+                r.poly_id.id = new_poly_id(r.poly_id.id);
+            }
+        }
+    });
+    pcp.original_instructions.iter_mut().for_each(|instr| {
+        instr.subs.iter_mut().for_each(|sub| {
+            *sub = new_poly_id(*sub);
+        });
+    });
+
+    pcp.is_valid_column.id.id = new_poly_id(pcp.is_valid_column.id.id);
+}
+
+fn merge_bus_interactions_simple<P: IntoOpenVm>(interactions: Vec<Vec<SymbolicBusInteraction<P>>>) -> Vec<SymbolicBusInteraction<P>> {
+    interactions.into_iter().flatten().collect_vec()
+}
+
+fn merge_bus_interactions<P: IntoOpenVm>(interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>) -> Vec<SymbolicBusInteraction<P>> {
+    let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> = Default::default();
+
+    for interactions in interactions_by_machine {
+        let interactions_by_bus_this_machine = interactions.into_iter()
+            // we group by bus id and number of args
+            .into_group_map_by(|interaction| (interaction.id, interaction.args.len()));
+        for (k,v) in interactions_by_bus_this_machine {
+            let e = interactions_by_bus.entry(k).or_default();
+            e.push(v);
+        }
+    }
+    let mut to_merge = vec![];
+    for interactions_per_machine in interactions_by_bus.into_values() {
+        // this is doing a "zip longest" among the interactions from each machine
+        // (for a given bus), and saving the elements that were picked
+        // together to be merged later
+        let mut idx = 0;
+        loop {
+            let mut items = vec![];
+            for machine_interactions in interactions_per_machine.iter() {
+                if let Some(item) = machine_interactions.get(idx) {
+                    items.push(item.clone());
+                }
+            }
+            if items.is_empty() {
+                break;
+            }
+            idx += 1;
+            to_merge.push(items);
+        }
+    }
+
+    to_merge.into_iter().map(|interactions| {
+        // assert_eq!(interactions.iter().map(|i| &i.kind).unique().count(), 1);
+
+        // add multiplicities together
+        let mut mult = interactions[0].mult.clone();
+        for mult2 in interactions.iter().skip(1).map(|i| i.mult.clone()) {
+            mult = mult + mult2;
+        }
+
+        // add args together
+        let mut args = interactions[0].args.clone();
+        for args2 in interactions.iter().skip(1).map(|i| i.args.clone()) {
+            args = args.into_iter().zip_eq(args2).map(|(a1, a2)| a1 + a2).collect();
+        }
+
+        SymbolicBusInteraction {
+            id: interactions[0].id,
+            mult,
+            args,
+        }
+    }).collect_vec()
+}

From 0f5448e787b32fe10262b85c295d83d7291f0f59 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Tue, 10 Jun 2025 10:45:35 -0300
Subject: [PATCH 16/39] default stacking 1.1

---
 openvm/src/customize_exe/air_stacking.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index d71987d404..34439559d8 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -29,7 +29,7 @@ pub fn air_stacking<P: IntoOpenVm>(
     // create apc groups by number of columns
     let mut groups: HashMap<usize, Vec<PowdrPrecompile<P>>> = Default::default();
     for pcp in extensions {
-        let idx = f64::log(pcp.machine.unique_columns().count() as f64, 2.0).floor() as usize;
+        let idx = f64::log(pcp.machine.unique_columns().count() as f64, 1.1).floor() as usize;
         groups.entry(idx).or_default().push(pcp);
     }
 

From 4dd19ef58e67d7b8de5d0715bf9f70cd51246260 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Fri, 13 Jun 2025 09:18:10 -0300
Subject: [PATCH 17/39] print max degree after handling
 constraints/interactions

---
 openvm/src/customize_exe/air_stacking.rs | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index 34439559d8..3aec5b8174 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -134,6 +134,7 @@ pub fn air_stacking<P: IntoOpenVm>(
         let mut stacked_constraints = join_constraints(stacked_constraints);
         stacked_constraints.sort();
         println!("after joining constraints: {}", stacked_constraints.len());
+        println!("max degree constraints: {}", stacked_constraints.iter().map(|c| c.expr.degree()).max().unwrap_or(0));
 
         // enforce only one is_valid is active
         let one = AlgebraicExpression::Number(P::ONE);
@@ -153,6 +154,14 @@ pub fn air_stacking<P: IntoOpenVm>(
         };
 
         println!("interaction count: {}", machine.bus_interactions.len());
+        println!("max degree interaction args: {}", machine.bus_interactions.iter()
+            .map(|i| i.args.iter().map(|a| a.degree()).max().unwrap_or(0))
+            .max()
+                 .unwrap_or(0));
+        println!("max degree interaction multiplicity:{}", machine.bus_interactions.iter()
+            .map(|i| i.mult.degree())
+            .max()
+                 .unwrap_or(0));
 
         println!("stacked width: {}", machine.unique_columns().count());
 

From a7537c21a0e17c3be58f13d30d286f63d0ae4d05 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Fri, 13 Jun 2025 09:54:30 -0300
Subject: [PATCH 18/39] print degrees before/after optimizer

---
 autoprecompiles/src/lib.rs | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/autoprecompiles/src/lib.rs b/autoprecompiles/src/lib.rs
index afd839840c..c85b4118b8 100644
--- a/autoprecompiles/src/lib.rs
+++ b/autoprecompiles/src/lib.rs
@@ -220,6 +220,9 @@ pub fn build<T: FieldElement, B: BusInteractionHandler<T> + IsBusStateful<T> + C
 ) -> Result<(SymbolicMachine<T>, Vec<Vec<u64>>), crate::constraint_optimizer::Error> {
     let (machine, subs) = statements_to_symbolic_machine(&program, vm_config.instruction_machines);
 
+    println!("before optimizer - machine constraints degree: {}", machine.constraints.iter().map(|c| c.expr.degree()).max().unwrap_or(0));
+    println!("before optimizer - machine bus args degree: {}", machine.bus_interactions.iter().map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0)).max().unwrap_or(0));
+
     let machine = optimizer::optimize(
         machine,
         vm_config.bus_interaction_handler,
@@ -227,6 +230,10 @@ pub fn build<T: FieldElement, B: BusInteractionHandler<T> + IsBusStateful<T> + C
         degree_bound,
     )?;
 
+    println!("after optimizer - machine constraints degree: {}", machine.constraints.iter().map(|c| c.expr.degree()).max().unwrap_or(0));
+    println!("after optimizer - machine bus args degree: {}", machine.bus_interactions.iter().map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0)).max().unwrap_or(0));
+    
+
     // add guards to constraints that are not satisfied by zeroes
     let machine = add_guards(machine);
 

From dfb387fba8dab74e1b5e4f86148ea107e02107bc Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 16 Jun 2025 14:35:17 -0300
Subject: [PATCH 19/39] plonk

---
 autoprecompiles/src/lib.rs               |  41 ++-
 openvm/src/customize_exe.rs              |   1 -
 openvm/src/customize_exe/air_stacking.rs | 428 ++++++++++++++++-------
 openvm/src/powdr_extension/chip.rs       |  18 +-
 openvm/src/powdr_extension/mod.rs        |   5 +-
 openvm/src/powdr_extension/vm.rs         |  37 +-
 6 files changed, 362 insertions(+), 168 deletions(-)

diff --git a/autoprecompiles/src/lib.rs b/autoprecompiles/src/lib.rs
index c85b4118b8..75498c60b1 100644
--- a/autoprecompiles/src/lib.rs
+++ b/autoprecompiles/src/lib.rs
@@ -220,8 +220,24 @@ pub fn build<T: FieldElement, B: BusInteractionHandler<T> + IsBusStateful<T> + C
 ) -> Result<(SymbolicMachine<T>, Vec<Vec<u64>>), crate::constraint_optimizer::Error> {
     let (machine, subs) = statements_to_symbolic_machine(&program, vm_config.instruction_machines);
 
-    println!("before optimizer - machine constraints degree: {}", machine.constraints.iter().map(|c| c.expr.degree()).max().unwrap_or(0));
-    println!("before optimizer - machine bus args degree: {}", machine.bus_interactions.iter().map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0)).max().unwrap_or(0));
+    println!(
+        "before optimizer - machine constraints degree: {}",
+        machine
+            .constraints
+            .iter()
+            .map(|c| c.expr.degree())
+            .max()
+            .unwrap_or(0)
+    );
+    println!(
+        "before optimizer - machine bus args degree: {}",
+        machine
+            .bus_interactions
+            .iter()
+            .map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0))
+            .max()
+            .unwrap_or(0)
+    );
 
     let machine = optimizer::optimize(
         machine,
@@ -230,9 +246,24 @@ pub fn build<T: FieldElement, B: BusInteractionHandler<T> + IsBusStateful<T> + C
         degree_bound,
     )?;
 
-    println!("after optimizer - machine constraints degree: {}", machine.constraints.iter().map(|c| c.expr.degree()).max().unwrap_or(0));
-    println!("after optimizer - machine bus args degree: {}", machine.bus_interactions.iter().map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0)).max().unwrap_or(0));
-    
+    println!(
+        "after optimizer - machine constraints degree: {}",
+        machine
+            .constraints
+            .iter()
+            .map(|c| c.expr.degree())
+            .max()
+            .unwrap_or(0)
+    );
+    println!(
+        "after optimizer - machine bus args degree: {}",
+        machine
+            .bus_interactions
+            .iter()
+            .map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0))
+            .max()
+            .unwrap_or(0)
+    );
 
     // add guards to constraints that are not satisfied by zeroes
     let machine = add_guards(machine);
diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 6895a1547a..918ac7e242 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -304,7 +304,6 @@ pub fn customize<P: IntoOpenVm>(
     )
 }
 
-
 fn branch_opcodes() -> Vec<usize> {
     let mut opcodes = vec![
         openvm_rv32im_transpiler::BranchEqualOpcode::BEQ
diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index 3aec5b8174..265dfd70ab 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -1,12 +1,21 @@
 use std::collections::{BTreeSet, HashMap};
 
-use powdr_autoprecompiles::{legacy_expression::{AlgebraicExpression, AlgebraicReference, PolyID, PolynomialType}, simplify_expression, SymbolicBusInteraction, SymbolicConstraint, SymbolicMachine};
+use powdr_autoprecompiles::{
+    legacy_expression::{AlgebraicExpression, AlgebraicReference, PolyID, PolynomialType},
+    simplify_expression, SymbolicBusInteraction, SymbolicConstraint, SymbolicMachine,
+};
 
 use powdr_autoprecompiles::powdr::UniqueColumns;
 
-use powdr_expression::{visitors::ExpressionVisitable, AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicUnaryOperation};
+use powdr_expression::{
+    visitors::ExpressionVisitable, AlgebraicBinaryOperation, AlgebraicBinaryOperator,
+    AlgebraicUnaryOperation,
+};
 
-use crate::{powdr_extension::{PowdrPrecompile, PowdrStackedPrecompile}, IntoOpenVm};
+use crate::{
+    powdr_extension::{PowdrPrecompile, PowdrStackedPrecompile},
+    IntoOpenVm,
+};
 
 use openvm_instructions::LocalOpcode;
 
@@ -21,7 +30,10 @@ pub fn air_stacking<P: IntoOpenVm>(
     assert!(!extensions.is_empty());
 
     extensions.iter_mut().for_each(|ext| {
-        ext.machine.constraints.iter_mut().for_each(|c| canonicalize_expression(&mut c.expr));
+        ext.machine
+            .constraints
+            .iter_mut()
+            .for_each(|c| canonicalize_expression(&mut c.expr));
         // ext.machine.constraints.sort();
         // compact_ids(ext)
     });
@@ -36,7 +48,12 @@ pub fn air_stacking<P: IntoOpenVm>(
     // sort each group by number of columns
     groups.values_mut().for_each(|g| {
         // assign largest pcp first
-        g.sort_by(|pcp1, pcp2| pcp2.machine.unique_columns().count().cmp(&pcp1.machine.unique_columns().count()));
+        g.sort_by(|pcp1, pcp2| {
+            pcp2.machine
+                .unique_columns()
+                .count()
+                .cmp(&pcp1.machine.unique_columns().count())
+        });
         let mut column_assigner = ColumnAssigner::default();
         g.iter_mut().for_each(|ext| {
             column_assigner.assign_pcp(ext);
@@ -68,16 +85,19 @@ pub fn air_stacking<P: IntoOpenVm>(
         println!("stacking {} precompiles", extensions.len());
 
         // take the max id in all pcps and add 1.
-        let is_valid_start = 1 + extensions.iter()
-            .flat_map(|pcp| pcp.original_instructions.iter().flat_map(|instr| instr.subs.iter()))
+        let is_valid_start = 1 + extensions
+            .iter()
+            .flat_map(|pcp| {
+                pcp.original_instructions
+                    .iter()
+                    .flat_map(|instr| instr.subs.iter())
+            })
             .max()
             .unwrap();
 
-
         let mut is_valid_sum: Option<AlgebraicExpression<P>> = None;
 
         for (idx, pcp) in extensions.iter_mut().enumerate() {
-
             // is_valid columns cannot be shared between precompiles. Here we do
             // their remapping into exclusive columns.
             let is_valid_new_id = is_valid_start + idx as u64;
@@ -115,12 +135,14 @@ pub fn air_stacking<P: IntoOpenVm>(
             });
 
             // guard interaction payloads so they can be merged later
-            remapped.bus_interactions.iter_mut().for_each(|interaction| {
-                interaction.args.iter_mut().for_each(|arg| {
-                    *arg = arg.clone() * is_valid.clone();
+            remapped
+                .bus_interactions
+                .iter_mut()
+                .for_each(|interaction| {
+                    interaction.args.iter_mut().for_each(|arg| {
+                        *arg = arg.clone() * is_valid.clone();
+                    });
                 });
-            });
-
 
             is_valid_sum = is_valid_sum
                 .map(|sum| sum + is_valid.clone())
@@ -134,7 +156,14 @@ pub fn air_stacking<P: IntoOpenVm>(
         let mut stacked_constraints = join_constraints(stacked_constraints);
         stacked_constraints.sort();
         println!("after joining constraints: {}", stacked_constraints.len());
-        println!("max degree constraints: {}", stacked_constraints.iter().map(|c| c.expr.degree()).max().unwrap_or(0));
+        println!(
+            "max degree constraints: {}",
+            stacked_constraints
+                .iter()
+                .map(|c| c.expr.degree())
+                .max()
+                .unwrap_or(0)
+        );
 
         // enforce only one is_valid is active
         let one = AlgebraicExpression::Number(P::ONE);
@@ -142,7 +171,10 @@ pub fn air_stacking<P: IntoOpenVm>(
 
         stacked_constraints.push(one_hot_is_valid.into());
 
-        println!("interaction count before: {}", interactions_by_machine.iter().flatten().count());
+        println!(
+            "interaction count before: {}",
+            interactions_by_machine.iter().flatten().count()
+        );
         // let mut stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
         // let mut stacked_interactions = merge_bus_interactions(interactions_by_machine);
         let mut stacked_interactions = merge_bus_interactions2(interactions_by_machine);
@@ -154,14 +186,24 @@ pub fn air_stacking<P: IntoOpenVm>(
         };
 
         println!("interaction count: {}", machine.bus_interactions.len());
-        println!("max degree interaction args: {}", machine.bus_interactions.iter()
-            .map(|i| i.args.iter().map(|a| a.degree()).max().unwrap_or(0))
-            .max()
-                 .unwrap_or(0));
-        println!("max degree interaction multiplicity:{}", machine.bus_interactions.iter()
-            .map(|i| i.mult.degree())
-            .max()
-                 .unwrap_or(0));
+        println!(
+            "max degree interaction args: {}",
+            machine
+                .bus_interactions
+                .iter()
+                .map(|i| i.args.iter().map(|a| a.degree()).max().unwrap_or(0))
+                .max()
+                .unwrap_or(0)
+        );
+        println!(
+            "max degree interaction multiplicity:{}",
+            machine
+                .bus_interactions
+                .iter()
+                .map(|i| i.mult.degree())
+                .max()
+                .unwrap_or(0)
+        );
 
         println!("stacked width: {}", machine.unique_columns().count());
 
@@ -177,7 +219,6 @@ pub fn air_stacking<P: IntoOpenVm>(
     result
 }
 
-
 /// make PolyIDs start from zero, sequential and compact
 fn compact_ids<P: IntoOpenVm>(pcp: &mut PowdrPrecompile<P>) {
     let mut curr_id = 0;
@@ -206,33 +247,44 @@ fn compact_ids<P: IntoOpenVm>(pcp: &mut PowdrPrecompile<P>) {
     pcp.is_valid_column.id.id = new_poly_id(pcp.is_valid_column.id.id);
 }
 
-fn merge_bus_interactions_simple<P: IntoOpenVm>(interactions: Vec<Vec<SymbolicBusInteraction<P>>>) -> Vec<SymbolicBusInteraction<P>> {
+fn merge_bus_interactions_simple<P: IntoOpenVm>(
+    interactions: Vec<Vec<SymbolicBusInteraction<P>>>,
+) -> Vec<SymbolicBusInteraction<P>> {
     interactions.into_iter().flatten().collect_vec()
 }
 
 fn merge_bus_interactions2<P: IntoOpenVm>(
-    interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>
+    interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>,
 ) -> Vec<SymbolicBusInteraction<P>> {
     // split interactions by bus/args len
-    let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> = Default::default();
+    let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> =
+        Default::default();
 
     for interactions in interactions_by_machine {
-        let interactions_by_bus_this_machine = interactions.into_iter()
+        let interactions_by_bus_this_machine = interactions
+            .into_iter()
             // we group by bus id and number of args
             .into_group_map_by(|interaction| (interaction.id, interaction.args.len()));
-        for (k,v) in interactions_by_bus_this_machine {
+        for (k, v) in interactions_by_bus_this_machine {
             let e = interactions_by_bus.entry(k).or_default();
             e.push(v);
         }
     }
-    interactions_by_bus.values_mut().for_each(|interactions_by_machine| {
-        interactions_by_machine.sort_by_key(|interactions| interactions.len());
-    });
+    interactions_by_bus
+        .values_mut()
+        .for_each(|interactions_by_machine| {
+            interactions_by_machine.sort_by_key(|interactions| interactions.len());
+        });
 
     let mut result = vec![];
     for mut interactions_by_machine in interactions_by_bus.into_values() {
         // to_merge is a vec of vecs, each inner vec is a set of interactions to be merged
-        let mut to_merge = interactions_by_machine.pop().unwrap().into_iter().map(|i| vec![i]).collect_vec();
+        let mut to_merge = interactions_by_machine
+            .pop()
+            .unwrap()
+            .into_iter()
+            .map(|i| vec![i])
+            .collect_vec();
         for machine_interactions in interactions_by_machine {
             let mut used = BTreeSet::new();
             let mut try_partial_match = vec![];
@@ -243,10 +295,10 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
                         continue;
                     }
                     let i2 = to_merge_set.get(0).unwrap();
-                    let all_args_same_structure = i.args.iter()
-                        .zip_eq(i2.args.iter())
-                        .all(|(a1, a2)| has_same_structure(strip_is_valid(&a1),
-                                                           strip_is_valid(&a2)));
+                    let all_args_same_structure =
+                        i.args.iter().zip_eq(i2.args.iter()).all(|(a1, a2)| {
+                            has_same_structure(strip_is_valid(&a1), strip_is_valid(&a2))
+                        });
                     if all_args_same_structure {
                         println!("EXACT MATCH");
                         to_merge_set.push(i);
@@ -266,10 +318,10 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
                     }
                     // check all args have the same structure
                     let i2 = to_merge_set.get(0).unwrap();
-                    let some_args_same_structure = i.args.iter()
-                        .zip_eq(i2.args.iter())
-                        .any(|(a1, a2)| has_same_structure(strip_is_valid(&a1),
-                                                           strip_is_valid(&a2)));
+                    let some_args_same_structure =
+                        i.args.iter().zip_eq(i2.args.iter()).any(|(a1, a2)| {
+                            has_same_structure(strip_is_valid(&a1), strip_is_valid(&a2))
+                        });
                     if some_args_same_structure {
                         println!("PARTIAL MATCH");
                         to_merge_set.push(i);
@@ -294,11 +346,23 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
             }
         }
 
-        fn merge_args<P: IntoOpenVm>(arg1: AlgebraicExpression<P>, arg2: AlgebraicExpression<P>) -> AlgebraicExpression<P> {
+        fn merge_args<P: IntoOpenVm>(
+            arg1: AlgebraicExpression<P>,
+            arg2: AlgebraicExpression<P>,
+        ) -> AlgebraicExpression<P> {
             match arg1 {
                 // expr * is_valid
-                AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left1, op: AlgebraicBinaryOperator::Mul, right: is_valid1 }) => {
-                    let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left2, op: AlgebraicBinaryOperator::Mul, right: is_valid2 }) = arg2 else {
+                AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                    left: left1,
+                    op: AlgebraicBinaryOperator::Mul,
+                    right: is_valid1,
+                }) => {
+                    let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                        left: left2,
+                        op: AlgebraicBinaryOperator::Mul,
+                        right: is_valid2,
+                    }) = arg2
+                    else {
                         panic!("Expected binary operation for arg2, got: {arg2}");
                     };
                     if has_same_structure(&left1, &left2) {
@@ -308,9 +372,7 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
                             left: left1,
                             op: AlgebraicBinaryOperator::Mul,
                             right: is_valid1,
-                        })
-                            +
-                        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                        }) + AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
                             left: left2,
                             op: AlgebraicBinaryOperator::Mul,
                             right: is_valid2,
@@ -318,16 +380,23 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
                     }
                 }
                 // exprA * is_validA + exprB * is_validB + ...
-                AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left:_, op: AlgebraicBinaryOperator::Add, right:_ }) => {
-                    arg1 + arg2
-                }
-                _ => unreachable!()
+                AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                    left: _,
+                    op: AlgebraicBinaryOperator::Add,
+                    right: _,
+                }) => arg1 + arg2,
+                _ => unreachable!(),
             }
         }
 
         fn simplify_arg<P: IntoOpenVm>(arg: AlgebraicExpression<P>) -> AlgebraicExpression<P> {
             // expr * is_valid_expr
-            if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op: AlgebraicBinaryOperator::Mul, right: _is_valid_expr }) = arg {
+            if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                left,
+                op: AlgebraicBinaryOperator::Mul,
+                right: _is_valid_expr,
+            }) = arg
+            {
                 *left
             } else {
                 // sum of expr * is_validN, can't simplify
@@ -338,30 +407,45 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
         // merge each set of interactions
         for set in to_merge {
             let id = set[0].id;
-            let mult = simplify_expression(set.iter().map(|i| i.mult.clone()).reduce(|a, b| a + b).unwrap());
-            let args = set.into_iter().map(|i| i.args)
-                .reduce(|a, b| a.into_iter().zip_eq(b).map(|(a1, a2)| merge_args(a1, a2)).collect())
+            let mult = simplify_expression(
+                set.iter()
+                    .map(|i| i.mult.clone())
+                    .reduce(|a, b| a + b)
+                    .unwrap(),
+            );
+            let args = set
+                .into_iter()
+                .map(|i| i.args)
+                .reduce(|a, b| {
+                    a.into_iter()
+                        .zip_eq(b)
+                        .map(|(a1, a2)| merge_args(a1, a2))
+                        .collect()
+                })
                 .unwrap();
-            let args = args.into_iter().map(|arg| simplify_expression(simplify_arg(arg))).collect_vec();
-            result.push(SymbolicBusInteraction {
-                id,
-                args,
-                mult,
-            });
+            let args = args
+                .into_iter()
+                .map(|arg| simplify_expression(simplify_arg(arg)))
+                .collect_vec();
+            result.push(SymbolicBusInteraction { id, args, mult });
         }
     }
 
     result
 }
 
-fn merge_bus_interactions<P: IntoOpenVm>(interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>) -> Vec<SymbolicBusInteraction<P>> {
-    let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> = Default::default();
+fn merge_bus_interactions<P: IntoOpenVm>(
+    interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>,
+) -> Vec<SymbolicBusInteraction<P>> {
+    let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> =
+        Default::default();
 
     for interactions in interactions_by_machine {
-        let interactions_by_bus_this_machine = interactions.into_iter()
+        let interactions_by_bus_this_machine = interactions
+            .into_iter()
             // we group by bus id and number of args
             .into_group_map_by(|interaction| (interaction.id, interaction.args.len()));
-        for (k,v) in interactions_by_bus_this_machine {
+        for (k, v) in interactions_by_bus_this_machine {
             let e = interactions_by_bus.entry(k).or_default();
             e.push(v);
         }
@@ -387,33 +471,44 @@ fn merge_bus_interactions<P: IntoOpenVm>(interactions_by_machine: Vec<Vec<Symbol
         }
     }
 
-    to_merge.into_iter().map(|interactions| {
-        // assert_eq!(interactions.iter().map(|i| &i.kind).unique().count(), 1);
+    to_merge
+        .into_iter()
+        .map(|interactions| {
+            // assert_eq!(interactions.iter().map(|i| &i.kind).unique().count(), 1);
 
-        // add multiplicities together
-        let mut mult = interactions[0].mult.clone();
-        for mult2 in interactions.iter().skip(1).map(|i| i.mult.clone()) {
-            mult = mult + mult2;
-        }
+            // add multiplicities together
+            let mut mult = interactions[0].mult.clone();
+            for mult2 in interactions.iter().skip(1).map(|i| i.mult.clone()) {
+                mult = mult + mult2;
+            }
 
-        // add args together
-        let mut args = interactions[0].args.clone();
-        for args2 in interactions.iter().skip(1).map(|i| i.args.clone()) {
-            args = args.into_iter().zip_eq(args2).map(|(a1, a2)| a1 + a2).collect();
-        }
+            // add args together
+            let mut args = interactions[0].args.clone();
+            for args2 in interactions.iter().skip(1).map(|i| i.args.clone()) {
+                args = args
+                    .into_iter()
+                    .zip_eq(args2)
+                    .map(|(a1, a2)| a1 + a2)
+                    .collect();
+            }
 
-        SymbolicBusInteraction {
-            id: interactions[0].id,
-            mult,
-            args,
-        }
-    }).collect_vec()
+            SymbolicBusInteraction {
+                id: interactions[0].id,
+                mult,
+                args,
+            }
+        })
+        .collect_vec()
 }
 
 fn canonicalize_expression<P: IntoOpenVm>(expr: &mut AlgebraicExpression<P>) {
     expr.pre_visit_expressions_mut(&mut |expr| {
-        if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) = expr {
-            if *op == AlgebraicBinaryOperator::Mul || *op == AlgebraicBinaryOperator::Add && left > right {
+        if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) =
+            expr
+        {
+            if *op == AlgebraicBinaryOperator::Mul
+                || *op == AlgebraicBinaryOperator::Add && left > right
+            {
                 std::mem::swap(left, right);
             }
         }
@@ -421,13 +516,25 @@ fn canonicalize_expression<P: IntoOpenVm>(expr: &mut AlgebraicExpression<P>) {
 }
 
 /// assumes ALL constraints are of the form `Y * is_valid_expr`
-fn join_constraints<P: IntoOpenVm>(mut constraints: Vec<SymbolicConstraint<P>>) -> Vec<SymbolicConstraint<P>> {
+fn join_constraints<P: IntoOpenVm>(
+    mut constraints: Vec<SymbolicConstraint<P>>,
+) -> Vec<SymbolicConstraint<P>> {
     constraints.sort();
     let mut result: Vec<SymbolicConstraint<P>> = vec![];
     'outer: for c1 in constraints {
         for c2 in result.iter_mut() {
-            if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {left: l1, op: AlgebraicBinaryOperator::Mul, right: is_valid1}) = &c1.expr {
-                if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {left: l2, op: AlgebraicBinaryOperator::Mul, right: is_valid2}) = &mut c2.expr {
+            if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                left: l1,
+                op: AlgebraicBinaryOperator::Mul,
+                right: is_valid1,
+            }) = &c1.expr
+            {
+                if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                    left: l2,
+                    op: AlgebraicBinaryOperator::Mul,
+                    right: is_valid2,
+                }) = &mut c2.expr
+                {
                     if l1 == l2 {
                         *is_valid2 = Box::new(*is_valid1.clone() + *is_valid2.clone());
                         continue 'outer;
@@ -440,47 +547,95 @@ fn join_constraints<P: IntoOpenVm>(mut constraints: Vec<SymbolicConstraint<P>>)
     result
 }
 
-fn has_same_structure<P: IntoOpenVm>(expr1: &AlgebraicExpression<P>, expr2: &AlgebraicExpression<P>) -> bool {
+fn has_same_structure<P: IntoOpenVm>(
+    expr1: &AlgebraicExpression<P>,
+    expr2: &AlgebraicExpression<P>,
+) -> bool {
     match (expr1, expr2) {
-        (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation{ left: left1, op: op1, right: right1 }),
-         AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation{ left: left2, op: op2, right: right2 })) => {
-            op1 == op2 && has_same_structure(left1, left2) && has_same_structure(right1, right2)
-        }
-        (AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op1, expr: expr1 }),
-         AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op2, expr: expr2 })) => {
-            op1 == op2 && has_same_structure(expr1, expr2)
-        }
+        (
+            AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                left: left1,
+                op: op1,
+                right: right1,
+            }),
+            AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                left: left2,
+                op: op2,
+                right: right2,
+            }),
+        ) => op1 == op2 && has_same_structure(left1, left2) && has_same_structure(right1, right2),
+        (
+            AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation {
+                op: op1,
+                expr: expr1,
+            }),
+            AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation {
+                op: op2,
+                expr: expr2,
+            }),
+        ) => op1 == op2 && has_same_structure(expr1, expr2),
         (AlgebraicExpression::Reference(r1), AlgebraicExpression::Reference(r2)) => {
             r1.poly_id.id == r2.poly_id.id
-        },
+        }
         (AlgebraicExpression::Number(n1), AlgebraicExpression::Number(n2)) => n1 == n2,
-        _ => false
+        _ => false,
     }
 }
 
 // assumes expression is of the form `some_expr * is_valid_expr`.
 fn strip_is_valid<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> &AlgebraicExpression<P> {
-    let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op: AlgebraicBinaryOperator::Mul, right: _ }) = expr else {
+    let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+        left,
+        op: AlgebraicBinaryOperator::Mul,
+        right: _,
+    }) = expr
+    else {
         panic!("Expression is not a binary operation with Mul operator: {expr}");
     };
     left.as_ref()
 }
 
-
-fn create_mapping<P: IntoOpenVm>(from: &AlgebraicExpression<P>, to: &AlgebraicExpression<P>) -> BiMap<u64, u64> {
-    fn create_mapping_inner<P: IntoOpenVm>(from: &AlgebraicExpression<P>, to: &AlgebraicExpression<P>) -> BiMap<u64, u64> {
+fn create_mapping<P: IntoOpenVm>(
+    from: &AlgebraicExpression<P>,
+    to: &AlgebraicExpression<P>,
+) -> BiMap<u64, u64> {
+    fn create_mapping_inner<P: IntoOpenVm>(
+        from: &AlgebraicExpression<P>,
+        to: &AlgebraicExpression<P>,
+    ) -> BiMap<u64, u64> {
         match (from, to) {
-            (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left1, op: op1, right: right1 }),
-             AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left: left2, op: op2, right: right2 })) => {
+            (
+                AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                    left: left1,
+                    op: op1,
+                    right: right1,
+                }),
+                AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                    left: left2,
+                    op: op2,
+                    right: right2,
+                }),
+            ) => {
                 assert_eq!(op1, op2);
                 let mut left = create_mapping_inner(left1, left2);
                 let right = create_mapping_inner(right1, right2);
                 // assert both sides don't conflict
-                assert!(extend_if_no_conflicts(&mut left, right), "conflict in mapping: {from} => {to}");
+                assert!(
+                    extend_if_no_conflicts(&mut left, right),
+                    "conflict in mapping: {from} => {to}"
+                );
                 left
             }
-            (AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op1, expr: expr1 }),
-             AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op: op2, expr: expr2 })) => {
+            (
+                AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation {
+                    op: op1,
+                    expr: expr1,
+                }),
+                AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation {
+                    op: op2,
+                    expr: expr2,
+                }),
+            ) => {
                 assert_eq!(op1, op2);
                 create_mapping_inner(&expr1, &expr2)
             }
@@ -496,10 +651,7 @@ fn create_mapping<P: IntoOpenVm>(from: &AlgebraicExpression<P>, to: &AlgebraicEx
     create_mapping_inner(from, to)
 }
 
-fn extend_if_no_conflicts(
-    mappings: &mut BiMap<u64, u64>,
-    new_mappings: BiMap<u64, u64>,
-) -> bool {
+fn extend_if_no_conflicts(mappings: &mut BiMap<u64, u64>, new_mappings: BiMap<u64, u64>) -> bool {
     for (k, v) in &new_mappings {
         if let Some(existing_v) = mappings.get_by_left(k) {
             if existing_v != v {
@@ -519,7 +671,9 @@ fn extend_if_no_conflicts(
 /// changes poly_id to match the order in which they are encountered in the expression.
 /// This allows us to compare two expressions for the same "structure".
 /// e.g., a+a == b+b and a+a != a+b
-fn expr_poly_id_by_order<P: IntoOpenVm>(mut expr: AlgebraicExpression<P>) -> AlgebraicExpression<P> {
+fn expr_poly_id_by_order<P: IntoOpenVm>(
+    mut expr: AlgebraicExpression<P>,
+) -> AlgebraicExpression<P> {
     let mut curr_id = 0;
     let mut id_map: HashMap<u64, u64> = Default::default();
     let mut new_poly_id = |old_id: u64| {
@@ -554,10 +708,7 @@ struct ColumnAssigner<P: IntoOpenVm> {
 }
 
 impl<P: IntoOpenVm> ColumnAssigner<P> {
-    fn assign_pcp(
-        &mut self,
-        pcp: &mut PowdrPrecompile<P>,
-    ) {
+    fn assign_pcp(&mut self, pcp: &mut PowdrPrecompile<P>) {
         let idx = self.pcps.len();
         println!("Assinging PCP {}", idx);
         let mut mappings = BiMap::new();
@@ -570,11 +721,15 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
         'outer: for c in &pcp.machine.constraints {
             for pcp2 in self.pcps.iter() {
                 for c2 in &pcp2.machine.constraints {
-                    if has_same_structure(&expr_poly_id_by_order(c.expr.clone()),
-                                          &expr_poly_id_by_order(c2.expr.clone())) {
-                        println!("Found same structure: \n\t{}\n\n\t{}",
-                                 &expr_poly_id_by_order(c.expr.clone()),
-                                 &expr_poly_id_by_order(c2.expr.clone()));
+                    if has_same_structure(
+                        &expr_poly_id_by_order(c.expr.clone()),
+                        &expr_poly_id_by_order(c2.expr.clone()),
+                    ) {
+                        println!(
+                            "Found same structure: \n\t{}\n\n\t{}",
+                            &expr_poly_id_by_order(c.expr.clone()),
+                            &expr_poly_id_by_order(c2.expr.clone())
+                        );
 
                         // println!("Found same structure (ids): \n\t{}\n\n\t{}",
                         //          &expr_orig_poly_ids(c.expr.clone()),
@@ -582,10 +737,7 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
 
                         // assign the same ids
                         let new_mappings = create_mapping(&c.expr, &c2.expr);
-                        if extend_if_no_conflicts(
-                            &mut mappings,
-                            new_mappings,
-                        ) {
+                        if extend_if_no_conflicts(&mut mappings, new_mappings) {
                             // println!("\tMappings: {:#?}", mappings);
                             continue 'outer;
                         }
@@ -600,16 +752,17 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
             for pcp2 in self.pcps.iter() {
                 for b2 in &pcp2.machine.bus_interactions {
                     if b.id == b2.id && b.args.len() == b2.args.len() {
-                        let all_args_same_structure = b.args.iter()
-                            .zip_eq(b2.args.iter())
-                            .all(|(a1, a2)| has_same_structure(&expr_poly_id_by_order(a1.clone()), &expr_poly_id_by_order(a2.clone())));
+                        let all_args_same_structure =
+                            b.args.iter().zip_eq(b2.args.iter()).all(|(a1, a2)| {
+                                has_same_structure(
+                                    &expr_poly_id_by_order(a1.clone()),
+                                    &expr_poly_id_by_order(a2.clone()),
+                                )
+                            });
                         if all_args_same_structure {
                             for (arg1, arg2) in b.args.iter().zip_eq(b2.args.iter()) {
                                 let new_mappings = create_mapping(&arg1, &arg2);
-                                extend_if_no_conflicts(
-                                    &mut mappings,
-                                    new_mappings,
-                                );
+                                extend_if_no_conflicts(&mut mappings, new_mappings);
                             }
                             continue 'outer;
                         }
@@ -630,7 +783,12 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
             while mappings.get_by_right(&curr_id).is_some() {
                 curr_id += 1;
             }
-            assert!(mappings.get_by_right(&curr_id).is_none(), "New id {} already exists in mappings: {:?}", curr_id, mappings);
+            assert!(
+                mappings.get_by_right(&curr_id).is_none(),
+                "New id {} already exists in mappings: {:?}",
+                curr_id,
+                mappings
+            );
             assert!(!mappings.contains_left(&old_id));
             // println!("Assigning new id {} for old id {}", curr_id, old_id);
             let id = curr_id;
diff --git a/openvm/src/powdr_extension/chip.rs b/openvm/src/powdr_extension/chip.rs
index ae8bc24e26..f88102812d 100644
--- a/openvm/src/powdr_extension/chip.rs
+++ b/openvm/src/powdr_extension/chip.rs
@@ -7,10 +7,7 @@ use std::{
 
 use crate::{traits::OpenVmField, utils::algebraic_to_symbolic, IntoOpenVm};
 
-use super::{
-    PowdrStackedPrecompile,
-    executor::PowdrExecutor,
-};
+use super::{executor::PowdrExecutor, PowdrStackedPrecompile};
 use itertools::Itertools;
 use openvm_circuit::system::memory::MemoryController;
 use openvm_circuit::{
@@ -28,13 +25,17 @@ use openvm_stark_backend::{
     air_builders::symbolic::{
         symbolic_expression::{SymbolicEvaluator, SymbolicExpression},
         symbolic_variable::{Entry, SymbolicVariable},
-    }, interaction::BusIndex, p3_air::{Air, BaseAir}, p3_matrix::dense::RowMajorMatrix, rap::ColumnsAir
+    },
+    interaction::BusIndex,
+    p3_air::{Air, BaseAir},
+    p3_matrix::dense::RowMajorMatrix,
+    rap::ColumnsAir,
 };
 
 use openvm_stark_backend::{
     config::{StarkGenericConfig, Val},
     interaction::InteractionBuilder,
-    p3_field::{Field, PrimeField32, FieldAlgebra},
+    p3_field::{Field, FieldAlgebra, PrimeField32},
     p3_matrix::Matrix,
     prover::types::AirProofInput,
     rap::{AnyRap, BaseAirWithPublicValues, PartitionedBaseAir},
@@ -213,10 +214,7 @@ impl<P: IntoOpenVm> ChipUsageGetter for PowdrChip<P> {
     }
 
     fn current_trace_height(&self) -> usize {
-        self.executors
-            .values()
-            .map(|e| e.number_of_calls())
-            .sum()
+        self.executors.values().map(|e| e.number_of_calls()).sum()
     }
 
     fn trace_width(&self) -> usize {
diff --git a/openvm/src/powdr_extension/mod.rs b/openvm/src/powdr_extension/mod.rs
index e169f4cd7a..2c572f4ca7 100644
--- a/openvm/src/powdr_extension/mod.rs
+++ b/openvm/src/powdr_extension/mod.rs
@@ -9,4 +9,7 @@ mod executor;
 mod plonk;
 
 pub use opcode::PowdrOpcode;
-pub use vm::{OriginalInstruction, PowdrExecutor, PowdrExtension, PowdrPeriphery, PowdrPrecompile, PowdrStackedPrecompile};
+pub use vm::{
+    OriginalInstruction, PowdrExecutor, PowdrExtension, PowdrPeriphery, PowdrPrecompile,
+    PowdrStackedPrecompile,
+};
diff --git a/openvm/src/powdr_extension/vm.rs b/openvm/src/powdr_extension/vm.rs
index 52513b0187..ae433ed890 100644
--- a/openvm/src/powdr_extension/vm.rs
+++ b/openvm/src/powdr_extension/vm.rs
@@ -204,10 +204,10 @@ impl<P: IntoOpenVm> VmExtension<OpenVmField<P>> for PowdrExtension<P> {
             .first()
             .cloned();
 
-        for precompile in &self.precompiles {
+        for stacked in &self.precompiles {
             let powdr_chip: PowdrExecutor<P> = match self.implementation {
                 PrecompileImplementation::SingleRowChip => PowdrChip::new(
-                    precompile.clone(),
+                    stacked.clone(),
                     offline_memory.clone(),
                     self.base_config.clone(),
                     SharedChips::new(
@@ -217,24 +217,29 @@ impl<P: IntoOpenVm> VmExtension<OpenVmField<P>> for PowdrExtension<P> {
                     ),
                 )
                 .into(),
-                PrecompileImplementation::PlonkChip => unimplemented!("TODO")
-                // PrecompileImplementation::PlonkChip => PlonkChip::new(
-                //     precompile.clone(),
-                //     offline_memory.clone(),
-                //     self.base_config.clone(),
-                //     SharedChips::new(
-                //         bitwise_lookup.clone(),
-                //         range_checker.clone(),
-                //         tuple_range_checker.cloned(),
-                //     ),
-                //     self.bus_map.clone(),
-                // )
-                // .into(),
+                PrecompileImplementation::PlonkChip => {
+                    if stacked.precompiles.len() != 1 {
+                        panic!("Plonk chip implementation does not support chip stacking");
+                    }
+                    let precompile = stacked.precompiles.values().next().unwrap().clone();
+                    PlonkChip::new(
+                        precompile.clone(),
+                        offline_memory.clone(),
+                        self.base_config.clone(),
+                        SharedChips::new(
+                            bitwise_lookup.clone(),
+                            range_checker.clone(),
+                            tuple_range_checker.cloned(),
+                        ),
+                        self.bus_map.clone(),
+                    )
+                    .into()
+                }
             };
 
             inventory.add_executor(
                 powdr_chip,
-                precompile
+                stacked
                     .precompiles
                     .keys()
                     .map(|opcode| opcode.global_opcode()),

From 1e8382014325e052bc720faa14b0ff2455ab6397 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 16 Jun 2025 15:40:10 -0300
Subject: [PATCH 20/39] chip stacking config

---
 autoprecompiles/src/lib.rs               | 64 ++++++++++++-----------
 openvm/src/customize_exe.rs              | 32 +++++++++---
 openvm/src/customize_exe/air_stacking.rs | 66 +++++-------------------
 openvm/src/lib.rs                        |  4 ++
 openvm/src/powdr_extension/vm.rs         | 11 ++++
 5 files changed, 87 insertions(+), 90 deletions(-)

diff --git a/autoprecompiles/src/lib.rs b/autoprecompiles/src/lib.rs
index 75498c60b1..4952e92a76 100644
--- a/autoprecompiles/src/lib.rs
+++ b/autoprecompiles/src/lib.rs
@@ -217,6 +217,7 @@ pub fn build<T: FieldElement, B: BusInteractionHandler<T> + IsBusStateful<T> + C
     vm_config: VmConfig<T, B>,
     degree_bound: DegreeBound,
     opcode: u32,
+    strict_is_valid_guards: bool,
 ) -> Result<(SymbolicMachine<T>, Vec<Vec<u64>>), crate::constraint_optimizer::Error> {
     let (machine, subs) = statements_to_symbolic_machine(&program, vm_config.instruction_machines);
 
@@ -266,7 +267,7 @@ pub fn build<T: FieldElement, B: BusInteractionHandler<T> + IsBusStateful<T> + C
     );
 
     // add guards to constraints that are not satisfied by zeroes
-    let machine = add_guards(machine);
+    let machine = add_guards(machine, strict_is_valid_guards);
 
     Ok((machine, subs))
 }
@@ -315,8 +316,8 @@ fn add_guards_constraint<T: FieldElement>(
 /// Assumptions:
 /// - There are exactly one execution bus receive and one execution bus send, in this order.
 /// - There is exactly one program bus send.
-fn add_guards<T: FieldElement>(mut machine: SymbolicMachine<T>) -> SymbolicMachine<T> {
-    // let pre_degree = machine.degree();
+fn add_guards<T: FieldElement>(mut machine: SymbolicMachine<T>, strict: bool) -> SymbolicMachine<T> {
+    let pre_degree = machine.degree();
 
     let max_id = machine
         .unique_columns()
@@ -337,14 +338,18 @@ fn add_guards<T: FieldElement>(mut machine: SymbolicMachine<T>) -> SymbolicMachi
         next: false,
     });
 
-    for c in &mut machine.constraints {
-        c.expr = is_valid.clone() * c.expr.clone();
-    }
-    // machine.constraints = machine
-    //     .constraints
-    //     .into_iter()
-    //     .map(|c| add_guards_constraint(c.expr, &is_valid).into())
-    //     .collect();
+    machine.constraints = if strict {
+        machine.constraints.into_iter().map(|mut c| {
+            c.expr = is_valid.clone() * c.expr.clone();
+            c
+        }).collect()
+    } else {
+        machine
+        .constraints
+        .into_iter()
+        .map(|c| add_guards_constraint(c.expr, &is_valid).into())
+        .collect()
+    };
 
     let [execution_bus_receive, execution_bus_send] = machine
         .bus_interactions
@@ -373,6 +378,8 @@ fn add_guards<T: FieldElement>(mut machine: SymbolicMachine<T>) -> SymbolicMachi
         .unwrap();
     program_bus_send.mult = is_valid.clone();
 
+    let mut is_valid_mults = vec![];
+
     for b in &mut machine.bus_interactions {
         match b.id {
             EXECUTION_BUS_ID => {
@@ -382,29 +389,28 @@ fn add_guards<T: FieldElement>(mut machine: SymbolicMachine<T>) -> SymbolicMachi
                 // already handled
             }
             _ => {
-                b.mult = is_valid.clone() * b.mult.clone();
-                // if !satisfies_zero_witness(&b.mult) {
-                //     // guard the multiplicity by `is_valid`
-                //     b.mult = is_valid.clone() * b.mult.clone();
-                //     // TODO this would not have to be cloned if we had *=
-                //     //c.expr *= guard.clone();
-                // } else {
-                //     // if it's zero, then we do not have to change the multiplicity, but we need to force it to be zero on non-valid rows with a constraint
-                //     let one = AlgebraicExpression::Number(1u64.into());
-                //     let e = ((one - is_valid.clone()) * b.mult.clone()).into();
-                //     is_valid_mults.push(e);
-                // }
+                if strict || !satisfies_zero_witness(&b.mult) {
+                    // guard the multiplicity by `is_valid`
+                    b.mult = is_valid.clone() * b.mult.clone();
+                } else {
+                    // if it's zero, then we do not have to change the multiplicity, but we need to force it to be zero on non-valid rows with a constraint
+                    let one = AlgebraicExpression::Number(1u64.into());
+                    let e = ((one - is_valid.clone()) * b.mult.clone()).into();
+                    is_valid_mults.push(e);
+                }
             }
         }
     }
 
-    // machine.constraints.extend(is_valid_mults);
+    machine.constraints.extend(is_valid_mults);
 
-    // assert_eq!(
-    //     pre_degree,
-    //     machine.degree(),
-    //     "Degree should not change after adding guards"
-    // );
+    if !strict {
+        assert_eq!(
+            pre_degree,
+            machine.degree(),
+            "Degree should not change after adding guards"
+        );
+    }
 
     // This needs to be added after the assertion above because it's a quadratic constraint
     // so it may increase the degree of the machine.
diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 918ac7e242..a4f7bb60a5 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -1,5 +1,6 @@
 use std::collections::{BTreeMap, BTreeSet, HashMap};
 
+use crate::powdr_extension::PowdrStackedPrecompile;
 use crate::IntoOpenVm;
 use crate::OpenVmField;
 use itertools::Itertools;
@@ -214,18 +215,26 @@ pub fn customize<P: IntoOpenVm>(
             acc_block.pretty_print(openvm_instruction_formatter)
         );
 
+        let mut degree_bound = config.degree_bound.clone();
+        // chip stacking needs to guard bus arguments, so we lower the inliner bound on bus interactions
+        if config.chip_stacking_log.is_some() {
+            degree_bound.bus_interactions -= 1;
+        }
+
         // Lookup if an APC is already cached by PgoConfig::Cell and generate the APC if not
         let CachedAutoPrecompile {
             apc_opcode,
             autoprecompile,
             subs,
         } = apc_cache.remove(&acc_block.start_idx).unwrap_or_else(|| {
+            let strict_is_valid_guards = config.chip_stacking_log.is_some();
             generate_apc_cache(
                 acc_block,
                 airs,
                 POWDR_OPCODE + i,
                 config.bus_map.clone(),
-                config.degree_bound,
+                degree_bound,
+                strict_is_valid_guards,
             )
             .expect("Failed to generate autoprecompile")
         });
@@ -292,7 +301,14 @@ pub fn customize<P: IntoOpenVm>(
         ));
     }
 
-    let extensions = air_stacking(extensions);
+    let extensions = if let Some(chip_stacking_log) = config.chip_stacking_log {
+        tracing::debug!("Chip stacking enabled, grouping log: {chip_stacking_log}");
+        air_stacking(extensions, chip_stacking_log)
+    } else {
+        tracing::debug!("Chip stacking disabled");
+        extensions.into_iter().map(PowdrStackedPrecompile::new_single).collect()
+    };
+
     (
         exe,
         PowdrExtension::new(
@@ -481,9 +497,10 @@ fn generate_apc_cache<P: IntoOpenVm>(
     apc_opcode: usize,
     bus_map: BusMap,
     degree_bound: DegreeBound,
+    strict_is_valid_guards: bool,
 ) -> Result<CachedAutoPrecompile<P>, Error> {
     let (autoprecompile, subs) =
-        generate_autoprecompile(block, airs, apc_opcode, bus_map, degree_bound)?;
+        generate_autoprecompile(block, airs, apc_opcode, bus_map, degree_bound, strict_is_valid_guards)?;
 
     Ok(CachedAutoPrecompile {
         apc_opcode,
@@ -505,7 +522,9 @@ fn generate_autoprecompile<P: IntoOpenVm>(
     airs: &BTreeMap<usize, SymbolicMachine<P>>,
     apc_opcode: usize,
     bus_map: BusMap,
-    mut degree_bound: DegreeBound,
+    degree_bound: DegreeBound,
+    // chip stacking needs constraints/multiplicities fully guarded by the is_valid column
+    strict_is_valid_guards: bool,
 ) -> Result<(SymbolicMachine<P>, Vec<Vec<u64>>), Error> {
     tracing::debug!(
         "Generating autoprecompile for block at index {}",
@@ -530,10 +549,8 @@ fn generate_autoprecompile<P: IntoOpenVm>(
         bus_interaction_handler: OpenVmBusInteractionHandler::new(bus_map),
     };
 
-    // chip stacking needs to add guards to bus arguments also, so we restrict the optimizer by 1 degree here.
-    degree_bound.bus_interactions -= 1;
     let (precompile, subs) =
-        powdr_autoprecompiles::build(program, vm_config, degree_bound, apc_opcode as u32)?;
+        powdr_autoprecompiles::build(program, vm_config, degree_bound, apc_opcode as u32, strict_is_valid_guards)?;
 
     // Check that substitution values are unique over all instructions
     assert!(subs.iter().flatten().all_unique());
@@ -602,6 +619,7 @@ fn sort_blocks_by_pgo_cell_cost<P: IntoOpenVm>(
                 POWDR_OPCODE + i,
                 config.bus_map.clone(),
                 config.degree_bound,
+                false,
             )
             .ok()?;
             apc_cache.insert(acc_block.start_idx, apc_cache_entry.clone());
diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index 265dfd70ab..1e161d68d1 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -26,6 +26,7 @@ use bimap::BiMap;
 // perform air stacking of multiple precompiles into stacked precompiles.
 pub fn air_stacking<P: IntoOpenVm>(
     mut extensions: Vec<PowdrPrecompile<P>>,
+    chip_stacking_log: f32,
 ) -> Vec<PowdrStackedPrecompile<P>> {
     assert!(!extensions.is_empty());
 
@@ -41,7 +42,7 @@ pub fn air_stacking<P: IntoOpenVm>(
     // create apc groups by number of columns
     let mut groups: HashMap<usize, Vec<PowdrPrecompile<P>>> = Default::default();
     for pcp in extensions {
-        let idx = f64::log(pcp.machine.unique_columns().count() as f64, 1.1).floor() as usize;
+        let idx = f32::log(pcp.machine.unique_columns().count() as f32, chip_stacking_log).floor() as usize;
         groups.entry(idx).or_default().push(pcp);
     }
 
@@ -63,26 +64,18 @@ pub fn air_stacking<P: IntoOpenVm>(
     let mut result = vec![];
 
     for mut extensions in groups.into_values() {
-        // if there is only one precompile in the group, we can just return it as a stacked precompile
+        // group of a single precompile, no stacking
         if extensions.len() == 1 {
-            println!("not stacking precompile");
-            result.push(PowdrStackedPrecompile {
-                machine: SymbolicMachine {
-                    constraints: extensions[0].machine.constraints.clone(),
-                    bus_interactions: extensions[0].machine.bus_interactions.clone(),
-                },
-                precompiles: extensions
-                    .into_iter()
-                    .map(|p| (p.opcode.clone(), p))
-                    .collect(),
-            });
+            let precompile = extensions.pop().unwrap();
+            tracing::debug!("Precompile {} not stacked", &precompile.opcode.global_opcode().as_usize());
+            result.push(PowdrStackedPrecompile::new_single(precompile));
             continue;
         }
 
         let mut stacked_constraints = vec![];
         let mut interactions_by_machine = vec![];
 
-        println!("stacking {} precompiles", extensions.len());
+        tracing::debug!("Stacking {} precompiles", extensions.len());
 
         // take the max id in all pcps and add 1.
         let is_valid_start = 1 + extensions
@@ -102,8 +95,6 @@ pub fn air_stacking<P: IntoOpenVm>(
             // their remapping into exclusive columns.
             let is_valid_new_id = is_valid_start + idx as u64;
 
-            println!("\tpcp width: {}", pcp.machine.unique_columns().count());
-
             // remap is_valid column in constraints and interactions
             let mut remapped = pcp.machine.clone();
 
@@ -152,18 +143,10 @@ pub fn air_stacking<P: IntoOpenVm>(
             interactions_by_machine.push(remapped.bus_interactions);
         }
 
-        println!("before joining constraints: {}", stacked_constraints.len());
+        tracing::debug!("Stacked chip has {} constraints", stacked_constraints.len());
         let mut stacked_constraints = join_constraints(stacked_constraints);
         stacked_constraints.sort();
-        println!("after joining constraints: {}", stacked_constraints.len());
-        println!(
-            "max degree constraints: {}",
-            stacked_constraints
-                .iter()
-                .map(|c| c.expr.degree())
-                .max()
-                .unwrap_or(0)
-        );
+        tracing::debug!("After joining constraints: {}", stacked_constraints.len());
 
         // enforce only one is_valid is active
         let one = AlgebraicExpression::Number(P::ONE);
@@ -171,42 +154,21 @@ pub fn air_stacking<P: IntoOpenVm>(
 
         stacked_constraints.push(one_hot_is_valid.into());
 
-        println!(
-            "interaction count before: {}",
+        tracing::debug!(
+            "Stacked chip has {} bus interactions",
             interactions_by_machine.iter().flatten().count()
         );
         // let mut stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
         // let mut stacked_interactions = merge_bus_interactions(interactions_by_machine);
         let mut stacked_interactions = merge_bus_interactions2(interactions_by_machine);
         stacked_interactions.sort();
+        tracing::debug!("After merging interactions: {}", stacked_interactions.len());
 
         let machine = SymbolicMachine {
             constraints: stacked_constraints,
             bus_interactions: stacked_interactions,
         };
 
-        println!("interaction count: {}", machine.bus_interactions.len());
-        println!(
-            "max degree interaction args: {}",
-            machine
-                .bus_interactions
-                .iter()
-                .map(|i| i.args.iter().map(|a| a.degree()).max().unwrap_or(0))
-                .max()
-                .unwrap_or(0)
-        );
-        println!(
-            "max degree interaction multiplicity:{}",
-            machine
-                .bus_interactions
-                .iter()
-                .map(|i| i.mult.degree())
-                .max()
-                .unwrap_or(0)
-        );
-
-        println!("stacked width: {}", machine.unique_columns().count());
-
         result.push(PowdrStackedPrecompile {
             precompiles: extensions
                 .into_iter()
@@ -226,7 +188,6 @@ fn compact_ids<P: IntoOpenVm>(pcp: &mut PowdrPrecompile<P>) {
     let mut new_poly_id = |old_id: u64| {
         *id_map.entry(old_id).or_insert_with(|| {
             let id = curr_id;
-            // println!("remapping poly id {} to {}", old_id, curr_id);
             curr_id += 1;
             id
         })
@@ -790,7 +751,6 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
                 mappings
             );
             assert!(!mappings.contains_left(&old_id));
-            // println!("Assigning new id {} for old id {}", curr_id, old_id);
             let id = curr_id;
             mappings.insert(old_id, id);
             curr_id += 1;
@@ -804,14 +764,12 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
                 }
             }
         });
-        // println!("----- assigning instr subs ------");
         pcp.original_instructions.iter_mut().for_each(|instr| {
             instr.subs.iter_mut().for_each(|sub| {
                 *sub = new_poly_id(*sub);
             });
         });
 
-        // println!("----- assigning is valid ------");
         pcp.is_valid_column.id.id = new_poly_id(pcp.is_valid_column.id.id);
 
         self.pcps.push(pcp.clone());
diff --git a/openvm/src/lib.rs b/openvm/src/lib.rs
index 5e48a1aa92..d1914f5247 100644
--- a/openvm/src/lib.rs
+++ b/openvm/src/lib.rs
@@ -294,6 +294,9 @@ pub struct PowdrConfig {
     pub degree_bound: DegreeBound,
     /// Implementation of the precompile, i.e., how to compile it to a RAP.
     pub implementation: PrecompileImplementation,
+    /// If present, use chip stacking.
+    /// Autoprecompile chips will be grouped by the provided log of their number of witness columns.
+    pub chip_stacking_log: Option<f32>,
 }
 
 impl PowdrConfig {
@@ -307,6 +310,7 @@ impl PowdrConfig {
                 bus_interactions: customize_exe::OPENVM_DEGREE_BOUND - 1,
             },
             implementation: PrecompileImplementation::default(),
+            chip_stacking_log: None,
         }
     }
 
diff --git a/openvm/src/powdr_extension/vm.rs b/openvm/src/powdr_extension/vm.rs
index ae433ed890..f46ccb53a1 100644
--- a/openvm/src/powdr_extension/vm.rs
+++ b/openvm/src/powdr_extension/vm.rs
@@ -106,6 +106,17 @@ pub struct PowdrStackedPrecompile<P: IntoOpenVm> {
     pub machine: SymbolicMachine<P>,
 }
 
+impl<P: IntoOpenVm> PowdrStackedPrecompile<P> {
+    pub fn new_single(
+        precompile: PowdrPrecompile<P>,
+    ) -> Self {
+        Self {
+            machine: precompile.machine.clone(),
+            precompiles: BTreeMap::from([(precompile.opcode.clone(), precompile)]),
+        }
+    }
+}
+
 impl<P: IntoOpenVm> PowdrExtension<P> {
     pub fn new(
         precompiles: Vec<PowdrStackedPrecompile<P>>,

From f363529879fb8f1874ff174a9de8d65b1525013b Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 16 Jun 2025 17:24:22 -0300
Subject: [PATCH 21/39] cmdline arg

---
 autoprecompiles/src/lib.rs               | 40 +++---------------------
 cli-openvm/src/main.rs                   | 27 ++++++++++++++--
 openvm/src/customize_exe.rs              |  2 +-
 openvm/src/customize_exe/air_stacking.rs | 18 ++---------
 openvm/src/lib.rs                        |  7 +++++
 5 files changed, 39 insertions(+), 55 deletions(-)

diff --git a/autoprecompiles/src/lib.rs b/autoprecompiles/src/lib.rs
index 4952e92a76..4aca19d7b7 100644
--- a/autoprecompiles/src/lib.rs
+++ b/autoprecompiles/src/lib.rs
@@ -221,24 +221,8 @@ pub fn build<T: FieldElement, B: BusInteractionHandler<T> + IsBusStateful<T> + C
 ) -> Result<(SymbolicMachine<T>, Vec<Vec<u64>>), crate::constraint_optimizer::Error> {
     let (machine, subs) = statements_to_symbolic_machine(&program, vm_config.instruction_machines);
 
-    println!(
-        "before optimizer - machine constraints degree: {}",
-        machine
-            .constraints
-            .iter()
-            .map(|c| c.expr.degree())
-            .max()
-            .unwrap_or(0)
-    );
-    println!(
-        "before optimizer - machine bus args degree: {}",
-        machine
-            .bus_interactions
-            .iter()
-            .map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0))
-            .max()
-            .unwrap_or(0)
-    );
+    println!("before optimizer - machine constraints degree: {}", machine.constraints.iter().map(|c| c.expr.degree()).max().unwrap_or(0));
+    println!("before optimizer - machine bus args degree: {}", machine.bus_interactions.iter().map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0)).max().unwrap_or(0));
 
     let machine = optimizer::optimize(
         machine,
@@ -247,24 +231,8 @@ pub fn build<T: FieldElement, B: BusInteractionHandler<T> + IsBusStateful<T> + C
         degree_bound,
     )?;
 
-    println!(
-        "after optimizer - machine constraints degree: {}",
-        machine
-            .constraints
-            .iter()
-            .map(|c| c.expr.degree())
-            .max()
-            .unwrap_or(0)
-    );
-    println!(
-        "after optimizer - machine bus args degree: {}",
-        machine
-            .bus_interactions
-            .iter()
-            .map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0))
-            .max()
-            .unwrap_or(0)
-    );
+    println!("after optimizer - machine constraints degree: {}", machine.constraints.iter().map(|c| c.expr.degree()).max().unwrap_or(0));
+    println!("after optimizer - machine bus args degree: {}", machine.bus_interactions.iter().map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0)).max().unwrap_or(0));
 
     // add guards to constraints that are not satisfied by zeroes
     let machine = add_guards(machine, strict_is_valid_guards);
diff --git a/cli-openvm/src/main.rs b/cli-openvm/src/main.rs
index 809c851bab..09becd608e 100644
--- a/cli-openvm/src/main.rs
+++ b/cli-openvm/src/main.rs
@@ -30,6 +30,9 @@ enum Commands {
 
         #[arg(long)]
         input: Option<u32>,
+
+        #[arg(long)]
+        chip_stacking: Option<f32>,
     },
 
     Execute {
@@ -46,6 +49,9 @@ enum Commands {
 
         #[arg(long)]
         input: Option<u32>,
+
+        #[arg(long)]
+        chip_stacking: Option<f32>,
     },
 
     Pgo {
@@ -77,6 +83,9 @@ enum Commands {
 
         #[arg(long)]
         input: Option<u32>,
+
+        #[arg(long)]
+        chip_stacking: Option<f32>,
     },
 }
 
@@ -102,8 +111,12 @@ fn run_command(command: Commands) {
             skip,
             pgo,
             input,
+            chip_stacking,
         } => {
-            let powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64);
+            let mut powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64);
+            if let Some(log) = chip_stacking {
+                powdr_config = powdr_config.with_chip_stacking(log);
+            }
             let pgo_config = get_pgo_config(guest.clone(), guest_opts.clone(), pgo, input);
             let program =
                 powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config).unwrap();
@@ -116,8 +129,12 @@ fn run_command(command: Commands) {
             skip,
             pgo,
             input,
+            chip_stacking,
         } => {
-            let powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64);
+            let mut powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64);
+            if let Some(log) = chip_stacking {
+                powdr_config = powdr_config.with_chip_stacking(log);
+            }
             let pgo_config = get_pgo_config(guest.clone(), guest_opts.clone(), pgo, input);
             let program =
                 powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config).unwrap();
@@ -132,8 +149,12 @@ fn run_command(command: Commands) {
             recursion,
             pgo,
             input,
+            chip_stacking,
         } => {
-            let powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64);
+            let mut powdr_config = PowdrConfig::new(autoprecompiles as u64, skip as u64);
+            if let Some(log) = chip_stacking {
+                powdr_config = powdr_config.with_chip_stacking(log);
+            }
             let pgo_config = get_pgo_config(guest.clone(), guest_opts.clone(), pgo, input);
             let program =
                 powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config).unwrap();
diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index a4f7bb60a5..641a7e103f 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -619,7 +619,7 @@ fn sort_blocks_by_pgo_cell_cost<P: IntoOpenVm>(
                 POWDR_OPCODE + i,
                 config.bus_map.clone(),
                 config.degree_bound,
-                false,
+                config.chip_stacking_log.is_some(),
             )
             .ok()?;
             apc_cache.insert(acc_block.start_idx, apc_cache_entry.clone());
diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index 1e161d68d1..53ed402ade 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -1,21 +1,12 @@
 use std::collections::{BTreeSet, HashMap};
 
-use powdr_autoprecompiles::{
-    legacy_expression::{AlgebraicExpression, AlgebraicReference, PolyID, PolynomialType},
-    simplify_expression, SymbolicBusInteraction, SymbolicConstraint, SymbolicMachine,
-};
+use powdr_autoprecompiles::{legacy_expression::{AlgebraicExpression, AlgebraicReference, PolyID, PolynomialType}, simplify_expression, SymbolicBusInteraction, SymbolicConstraint, SymbolicMachine};
 
 use powdr_autoprecompiles::powdr::UniqueColumns;
 
-use powdr_expression::{
-    visitors::ExpressionVisitable, AlgebraicBinaryOperation, AlgebraicBinaryOperator,
-    AlgebraicUnaryOperation,
-};
+use powdr_expression::{visitors::ExpressionVisitable, AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicUnaryOperation};
 
-use crate::{
-    powdr_extension::{PowdrPrecompile, PowdrStackedPrecompile},
-    IntoOpenVm,
-};
+use crate::{powdr_extension::{PowdrPrecompile, PowdrStackedPrecompile}, IntoOpenVm};
 
 use openvm_instructions::LocalOpcode;
 
@@ -35,7 +26,6 @@ pub fn air_stacking<P: IntoOpenVm>(
             .constraints
             .iter_mut()
             .for_each(|c| canonicalize_expression(&mut c.expr));
-        // ext.machine.constraints.sort();
         // compact_ids(ext)
     });
 
@@ -145,7 +135,6 @@ pub fn air_stacking<P: IntoOpenVm>(
 
         tracing::debug!("Stacked chip has {} constraints", stacked_constraints.len());
         let mut stacked_constraints = join_constraints(stacked_constraints);
-        stacked_constraints.sort();
         tracing::debug!("After joining constraints: {}", stacked_constraints.len());
 
         // enforce only one is_valid is active
@@ -161,7 +150,6 @@ pub fn air_stacking<P: IntoOpenVm>(
         // let mut stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
         // let mut stacked_interactions = merge_bus_interactions(interactions_by_machine);
         let mut stacked_interactions = merge_bus_interactions2(interactions_by_machine);
-        stacked_interactions.sort();
         tracing::debug!("After merging interactions: {}", stacked_interactions.len());
 
         let machine = SymbolicMachine {
diff --git a/openvm/src/lib.rs b/openvm/src/lib.rs
index d1914f5247..082f2ba4e4 100644
--- a/openvm/src/lib.rs
+++ b/openvm/src/lib.rs
@@ -341,6 +341,13 @@ impl PowdrConfig {
             ..self
         }
     }
+
+    pub fn with_chip_stacking(self, chip_stacking_log: f32) -> Self {
+        Self {
+            chip_stacking_log: Some(chip_stacking_log),
+            ..self
+        }
+    }
 }
 
 pub fn compile_guest(

From 6f6ca78148491e47c3473b4856532a24dd444411 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 16 Jun 2025 17:59:26 -0300
Subject: [PATCH 22/39] remove unused

---
 openvm/src/customize_exe/air_stacking.rs | 150 +++--------------------
 1 file changed, 15 insertions(+), 135 deletions(-)

diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index 53ed402ade..dbf4b95110 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -147,9 +147,7 @@ pub fn air_stacking<P: IntoOpenVm>(
             "Stacked chip has {} bus interactions",
             interactions_by_machine.iter().flatten().count()
         );
-        // let mut stacked_interactions = merge_bus_interactions_simple(interactions_by_machine);
-        // let mut stacked_interactions = merge_bus_interactions(interactions_by_machine);
-        let mut stacked_interactions = merge_bus_interactions2(interactions_by_machine);
+        let stacked_interactions = merge_bus_interactions(interactions_by_machine);
         tracing::debug!("After merging interactions: {}", stacked_interactions.len());
 
         let machine = SymbolicMachine {
@@ -169,40 +167,9 @@ pub fn air_stacking<P: IntoOpenVm>(
     result
 }
 
-/// make PolyIDs start from zero, sequential and compact
-fn compact_ids<P: IntoOpenVm>(pcp: &mut PowdrPrecompile<P>) {
-    let mut curr_id = 0;
-    let mut id_map: HashMap<u64, u64> = Default::default();
-    let mut new_poly_id = |old_id: u64| {
-        *id_map.entry(old_id).or_insert_with(|| {
-            let id = curr_id;
-            curr_id += 1;
-            id
-        })
-    };
-    pcp.machine.pre_visit_expressions_mut(&mut |expr| {
-        if let AlgebraicExpression::Reference(r) = expr {
-            if r.poly_id.ptype == PolynomialType::Committed {
-                r.poly_id.id = new_poly_id(r.poly_id.id);
-            }
-        }
-    });
-    pcp.original_instructions.iter_mut().for_each(|instr| {
-        instr.subs.iter_mut().for_each(|sub| {
-            *sub = new_poly_id(*sub);
-        });
-    });
-
-    pcp.is_valid_column.id.id = new_poly_id(pcp.is_valid_column.id.id);
-}
-
-fn merge_bus_interactions_simple<P: IntoOpenVm>(
-    interactions: Vec<Vec<SymbolicBusInteraction<P>>>,
-) -> Vec<SymbolicBusInteraction<P>> {
-    interactions.into_iter().flatten().collect_vec()
-}
-
-fn merge_bus_interactions2<P: IntoOpenVm>(
+/// Merge bus interactions, taking into account args that have the same expression.
+/// When only compatible args are merged, the expression doesn't need an is_valid guard.
+fn merge_bus_interactions<P: IntoOpenVm>(
     interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>,
 ) -> Vec<SymbolicBusInteraction<P>> {
     // split interactions by bus/args len
@@ -235,7 +202,9 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
             .map(|i| vec![i])
             .collect_vec();
         for machine_interactions in interactions_by_machine {
+            // merge sets already used by this machine
             let mut used = BTreeSet::new();
+            // interactions that didn't find an exact match in structure
             let mut try_partial_match = vec![];
             'outer: for i in machine_interactions {
                 // try to find a call with an exact match
@@ -249,7 +218,7 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
                             has_same_structure(strip_is_valid(&a1), strip_is_valid(&a2))
                         });
                     if all_args_same_structure {
-                        println!("EXACT MATCH");
+                        // found an exact match
                         to_merge_set.push(i);
                         used.insert(idx);
                         continue 'outer;
@@ -258,6 +227,7 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
                 try_partial_match.push(i);
             }
 
+            // interactions that did not find a partial structure match
             let mut no_match = vec![];
             'outer: for i in try_partial_match {
                 // didn't find exact match, try one where some args match
@@ -272,7 +242,7 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
                             has_same_structure(strip_is_valid(&a1), strip_is_valid(&a2))
                         });
                     if some_args_same_structure {
-                        println!("PARTIAL MATCH");
+                        // found a partial match on some args
                         to_merge_set.push(i);
                         used.insert(idx);
                         continue 'outer;
@@ -285,7 +255,6 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
                 // just pick the first unused one
                 for (idx, to_merge_set) in to_merge.iter_mut().enumerate() {
                     if !used.contains(&idx) {
-                        println!("NO MATCH");
                         to_merge_set.push(i);
                         used.insert(idx);
                         continue 'outer;
@@ -383,73 +352,6 @@ fn merge_bus_interactions2<P: IntoOpenVm>(
     result
 }
 
-fn merge_bus_interactions<P: IntoOpenVm>(
-    interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>,
-) -> Vec<SymbolicBusInteraction<P>> {
-    let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> =
-        Default::default();
-
-    for interactions in interactions_by_machine {
-        let interactions_by_bus_this_machine = interactions
-            .into_iter()
-            // we group by bus id and number of args
-            .into_group_map_by(|interaction| (interaction.id, interaction.args.len()));
-        for (k, v) in interactions_by_bus_this_machine {
-            let e = interactions_by_bus.entry(k).or_default();
-            e.push(v);
-        }
-    }
-    let mut to_merge = vec![];
-    for interactions_per_machine in interactions_by_bus.into_values() {
-        // this is doing a "zip longest" among the interactions from each machine
-        // (for a given bus), and saving the elements that were picked
-        // together to be merged later
-        let mut idx = 0;
-        loop {
-            let mut items = vec![];
-            for machine_interactions in interactions_per_machine.iter() {
-                if let Some(item) = machine_interactions.get(idx) {
-                    items.push(item.clone());
-                }
-            }
-            if items.is_empty() {
-                break;
-            }
-            idx += 1;
-            to_merge.push(items);
-        }
-    }
-
-    to_merge
-        .into_iter()
-        .map(|interactions| {
-            // assert_eq!(interactions.iter().map(|i| &i.kind).unique().count(), 1);
-
-            // add multiplicities together
-            let mut mult = interactions[0].mult.clone();
-            for mult2 in interactions.iter().skip(1).map(|i| i.mult.clone()) {
-                mult = mult + mult2;
-            }
-
-            // add args together
-            let mut args = interactions[0].args.clone();
-            for args2 in interactions.iter().skip(1).map(|i| i.args.clone()) {
-                args = args
-                    .into_iter()
-                    .zip_eq(args2)
-                    .map(|(a1, a2)| a1 + a2)
-                    .collect();
-            }
-
-            SymbolicBusInteraction {
-                id: interactions[0].id,
-                mult,
-                args,
-            }
-        })
-        .collect_vec()
-}
-
 fn canonicalize_expression<P: IntoOpenVm>(expr: &mut AlgebraicExpression<P>) {
     expr.pre_visit_expressions_mut(&mut |expr| {
         if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) =
@@ -544,6 +446,7 @@ fn strip_is_valid<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> &AlgebraicExp
     left.as_ref()
 }
 
+/// Given two expressions with the same structure, creates mapping of polynomial ids from one to the other.
 fn create_mapping<P: IntoOpenVm>(
     from: &AlgebraicExpression<P>,
     to: &AlgebraicExpression<P>,
@@ -600,6 +503,7 @@ fn create_mapping<P: IntoOpenVm>(
     create_mapping_inner(from, to)
 }
 
+/// If there are no key or value clashes, extends the bimap.
 fn extend_if_no_conflicts(mappings: &mut BiMap<u64, u64>, new_mappings: BiMap<u64, u64>) -> bool {
     for (k, v) in &new_mappings {
         if let Some(existing_v) = mappings.get_by_left(k) {
@@ -641,15 +545,6 @@ fn expr_poly_id_by_order<P: IntoOpenVm>(
     expr
 }
 
-fn expr_orig_poly_ids<P: IntoOpenVm>(mut expr: AlgebraicExpression<P>) -> AlgebraicExpression<P> {
-    expr.pre_visit_expressions_mut(&mut |e| {
-        if let AlgebraicExpression::Reference(r) = e {
-            r.name = format!("col_{}", r.poly_id.id); // this one not needed, just for printing
-        }
-    });
-    expr
-}
-
 #[derive(Default)]
 struct ColumnAssigner<P: IntoOpenVm> {
     // original PCPs
@@ -658,8 +553,6 @@ struct ColumnAssigner<P: IntoOpenVm> {
 
 impl<P: IntoOpenVm> ColumnAssigner<P> {
     fn assign_pcp(&mut self, pcp: &mut PowdrPrecompile<P>) {
-        let idx = self.pcps.len();
-        println!("Assinging PCP {}", idx);
         let mut mappings = BiMap::new();
         // for each constraint, we want to find if there is a constraint in
         // another machine with the same structure that we can assign the same
@@ -674,23 +567,12 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
                         &expr_poly_id_by_order(c.expr.clone()),
                         &expr_poly_id_by_order(c2.expr.clone()),
                     ) {
-                        println!(
-                            "Found same structure: \n\t{}\n\n\t{}",
-                            &expr_poly_id_by_order(c.expr.clone()),
-                            &expr_poly_id_by_order(c2.expr.clone())
-                        );
-
-                        // println!("Found same structure (ids): \n\t{}\n\n\t{}",
-                        //          &expr_orig_poly_ids(c.expr.clone()),
-                        //          &expr_orig_poly_ids(c2.expr.clone()));
-
-                        // assign the same ids
+                        // Found the same structure, try to extend the column mapping
                         let new_mappings = create_mapping(&c.expr, &c2.expr);
                         if extend_if_no_conflicts(&mut mappings, new_mappings) {
-                            // println!("\tMappings: {:#?}", mappings);
+                            // we're done for this constraint
                             continue 'outer;
                         }
-                        println!("\tCould not extend due to conflicts!");
                     }
                 }
             }
@@ -720,15 +602,13 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
             }
         }
 
-        println!("Mappings for PCP {}: {:#?}", idx, mappings);
-
-        // now assign new ids for all references in the PCP
+        // now we assign the columns to the PCP. When an id is not present in the mapping, use some other unused column
         let mut curr_id = 0;
         let mut new_poly_id = |old_id: u64| {
             if let Some(id) = mappings.get_by_left(&old_id) {
                 return *id;
             }
-            // if the new id is a target in the mappings, find a new id not yet in the mapping
+            // find a new column not yet used in the mapping
             while mappings.get_by_right(&curr_id).is_some() {
                 curr_id += 1;
             }

From 05af4aecb83cf4d8a9109e6a6897b2cb5cce5c23 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 16 Jun 2025 18:29:36 -0300
Subject: [PATCH 23/39] cleanup

---
 openvm/src/customize_exe/air_stacking.rs | 68 ++++++++++++++++++------
 1 file changed, 51 insertions(+), 17 deletions(-)

diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index dbf4b95110..8b458bf35d 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -26,7 +26,6 @@ pub fn air_stacking<P: IntoOpenVm>(
             .constraints
             .iter_mut()
             .for_each(|c| canonicalize_expression(&mut c.expr));
-        // compact_ids(ext)
     });
 
     // create apc groups by number of columns
@@ -121,7 +120,8 @@ pub fn air_stacking<P: IntoOpenVm>(
                 .iter_mut()
                 .for_each(|interaction| {
                     interaction.args.iter_mut().for_each(|arg| {
-                        *arg = arg.clone() * is_valid.clone();
+                        *arg = is_valid.clone() * arg.clone();
+                        canonicalize_expression(arg);
                     });
                 });
 
@@ -215,7 +215,7 @@ fn merge_bus_interactions<P: IntoOpenVm>(
                     let i2 = to_merge_set.get(0).unwrap();
                     let all_args_same_structure =
                         i.args.iter().zip_eq(i2.args.iter()).all(|(a1, a2)| {
-                            has_same_structure(strip_is_valid(&a1), strip_is_valid(&a2))
+                            has_same_structure(strip_guard(&a1), strip_guard(&a2))
                         });
                     if all_args_same_structure {
                         // found an exact match
@@ -239,7 +239,7 @@ fn merge_bus_interactions<P: IntoOpenVm>(
                     let i2 = to_merge_set.get(0).unwrap();
                     let some_args_same_structure =
                         i.args.iter().zip_eq(i2.args.iter()).any(|(a1, a2)| {
-                            has_same_structure(strip_is_valid(&a1), strip_is_valid(&a2))
+                            has_same_structure(strip_guard(&a1), strip_guard(&a2))
                         });
                     if some_args_same_structure {
                         // found a partial match on some args
@@ -352,8 +352,40 @@ fn merge_bus_interactions<P: IntoOpenVm>(
     result
 }
 
+/// true if expression is of the form `is_valid * some_expr`.
+fn is_valid_guarded<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> bool {
+    match expr {
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left,
+            op: AlgebraicBinaryOperator::Mul,
+            ..
+        }) => {
+            match left.as_ref() {
+                AlgebraicExpression::Reference(AlgebraicReference {
+                    name,
+                    ..
+                }) => {
+                    name.starts_with("is_valid")
+                }
+                _ => false,
+            }
+        },
+        _ => false,
+    }
+}
+
+/// takes an expression `is_valid * expr` and canonicalizes the right side
 fn canonicalize_expression<P: IntoOpenVm>(expr: &mut AlgebraicExpression<P>) {
-    expr.pre_visit_expressions_mut(&mut |expr| {
+    assert!(is_valid_guarded(expr), "not left guarded by is_valid: {expr}");
+    let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+        left: _is_valid,
+        op: AlgebraicBinaryOperator::Mul,
+        right,
+    }) = expr
+    else {
+        panic!("expected a Mul operation, got: {expr}");
+    };
+    right.pre_visit_expressions_mut(&mut |expr| {
         if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) =
             expr
         {
@@ -366,7 +398,7 @@ fn canonicalize_expression<P: IntoOpenVm>(expr: &mut AlgebraicExpression<P>) {
     });
 }
 
-/// assumes ALL constraints are of the form `Y * is_valid_expr`
+/// assumes constraints are of the form `is_valid_expr * some_expr`
 fn join_constraints<P: IntoOpenVm>(
     mut constraints: Vec<SymbolicConstraint<P>>,
 ) -> Vec<SymbolicConstraint<P>> {
@@ -375,18 +407,18 @@ fn join_constraints<P: IntoOpenVm>(
     'outer: for c1 in constraints {
         for c2 in result.iter_mut() {
             if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                left: l1,
+                left: is_valid1,
                 op: AlgebraicBinaryOperator::Mul,
-                right: is_valid1,
+                right: r1,
             }) = &c1.expr
             {
                 if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                    left: l2,
+                    left: is_valid2,
                     op: AlgebraicBinaryOperator::Mul,
-                    right: is_valid2,
+                    right: r2,
                 }) = &mut c2.expr
                 {
-                    if l1 == l2 {
+                    if r1 == r2 {
                         *is_valid2 = Box::new(*is_valid1.clone() + *is_valid2.clone());
                         continue 'outer;
                     }
@@ -433,17 +465,18 @@ fn has_same_structure<P: IntoOpenVm>(
     }
 }
 
-// assumes expression is of the form `some_expr * is_valid_expr`.
-fn strip_is_valid<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> &AlgebraicExpression<P> {
+/// assumes expression is of the form `is_valid_expr * some_expr`.
+/// Returns the right side.
+fn strip_guard<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> &AlgebraicExpression<P> {
     let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-        left,
+        left: _,
         op: AlgebraicBinaryOperator::Mul,
-        right: _,
+        right,
     }) = expr
     else {
         panic!("Expression is not a binary operation with Mul operator: {expr}");
     };
-    left.as_ref()
+    right.as_ref()
 }
 
 /// Given two expressions with the same structure, creates mapping of polynomial ids from one to the other.
@@ -539,7 +572,8 @@ fn expr_poly_id_by_order<P: IntoOpenVm>(
     expr.pre_visit_expressions_mut(&mut |e| {
         if let AlgebraicExpression::Reference(r) = e {
             r.poly_id.id = new_poly_id(r.poly_id.id);
-            r.name = format!("col_{}", r.poly_id.id); // this one not needed, just for printing
+            // this is just useful for printing
+            r.name = format!("col_{}", r.poly_id.id);
         }
     });
     expr

From 1075ba20b23b51d3468fbff0676c1874e184de2b Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Tue, 17 Jun 2025 11:23:53 -0300
Subject: [PATCH 24/39] fmt

---
 openvm/src/customize_exe.rs              | 24 +++++++--
 openvm/src/customize_exe/air_stacking.rs | 65 +++++++++++++++---------
 openvm/src/powdr_extension/vm.rs         |  4 +-
 3 files changed, 61 insertions(+), 32 deletions(-)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 641a7e103f..663103dea2 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -306,7 +306,10 @@ pub fn customize<P: IntoOpenVm>(
         air_stacking(extensions, chip_stacking_log)
     } else {
         tracing::debug!("Chip stacking disabled");
-        extensions.into_iter().map(PowdrStackedPrecompile::new_single).collect()
+        extensions
+            .into_iter()
+            .map(PowdrStackedPrecompile::new_single)
+            .collect()
     };
 
     (
@@ -499,8 +502,14 @@ fn generate_apc_cache<P: IntoOpenVm>(
     degree_bound: DegreeBound,
     strict_is_valid_guards: bool,
 ) -> Result<CachedAutoPrecompile<P>, Error> {
-    let (autoprecompile, subs) =
-        generate_autoprecompile(block, airs, apc_opcode, bus_map, degree_bound, strict_is_valid_guards)?;
+    let (autoprecompile, subs) = generate_autoprecompile(
+        block,
+        airs,
+        apc_opcode,
+        bus_map,
+        degree_bound,
+        strict_is_valid_guards,
+    )?;
 
     Ok(CachedAutoPrecompile {
         apc_opcode,
@@ -549,8 +558,13 @@ fn generate_autoprecompile<P: IntoOpenVm>(
         bus_interaction_handler: OpenVmBusInteractionHandler::new(bus_map),
     };
 
-    let (precompile, subs) =
-        powdr_autoprecompiles::build(program, vm_config, degree_bound, apc_opcode as u32, strict_is_valid_guards)?;
+    let (precompile, subs) = powdr_autoprecompiles::build(
+        program,
+        vm_config,
+        degree_bound,
+        apc_opcode as u32,
+        strict_is_valid_guards,
+    )?;
 
     // Check that substitution values are unique over all instructions
     assert!(subs.iter().flatten().all_unique());
diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index 8b458bf35d..c5e3c5ecfa 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -1,12 +1,21 @@
 use std::collections::{BTreeSet, HashMap};
 
-use powdr_autoprecompiles::{legacy_expression::{AlgebraicExpression, AlgebraicReference, PolyID, PolynomialType}, simplify_expression, SymbolicBusInteraction, SymbolicConstraint, SymbolicMachine};
+use powdr_autoprecompiles::{
+    legacy_expression::{AlgebraicExpression, AlgebraicReference, PolyID, PolynomialType},
+    simplify_expression, SymbolicBusInteraction, SymbolicConstraint, SymbolicMachine,
+};
 
 use powdr_autoprecompiles::powdr::UniqueColumns;
 
-use powdr_expression::{visitors::ExpressionVisitable, AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicUnaryOperation};
+use powdr_expression::{
+    visitors::ExpressionVisitable, AlgebraicBinaryOperation, AlgebraicBinaryOperator,
+    AlgebraicUnaryOperation,
+};
 
-use crate::{powdr_extension::{PowdrPrecompile, PowdrStackedPrecompile}, IntoOpenVm};
+use crate::{
+    powdr_extension::{PowdrPrecompile, PowdrStackedPrecompile},
+    IntoOpenVm,
+};
 
 use openvm_instructions::LocalOpcode;
 
@@ -14,7 +23,8 @@ use itertools::Itertools;
 
 use bimap::BiMap;
 
-// perform air stacking of multiple precompiles into stacked precompiles.
+/// Perform air stacking on the set of precompiles, grouping them by the log of
+/// their number of columns.
 pub fn air_stacking<P: IntoOpenVm>(
     mut extensions: Vec<PowdrPrecompile<P>>,
     chip_stacking_log: f32,
@@ -31,7 +41,11 @@ pub fn air_stacking<P: IntoOpenVm>(
     // create apc groups by number of columns
     let mut groups: HashMap<usize, Vec<PowdrPrecompile<P>>> = Default::default();
     for pcp in extensions {
-        let idx = f32::log(pcp.machine.unique_columns().count() as f32, chip_stacking_log).floor() as usize;
+        let idx = f32::log(
+            pcp.machine.unique_columns().count() as f32,
+            chip_stacking_log,
+        )
+        .floor() as usize;
         groups.entry(idx).or_default().push(pcp);
     }
 
@@ -56,7 +70,10 @@ pub fn air_stacking<P: IntoOpenVm>(
         // group of a single precompile, no stacking
         if extensions.len() == 1 {
             let precompile = extensions.pop().unwrap();
-            tracing::debug!("Precompile {} not stacked", &precompile.opcode.global_opcode().as_usize());
+            tracing::debug!(
+                "Precompile {} not stacked",
+                &precompile.opcode.global_opcode().as_usize()
+            );
             result.push(PowdrStackedPrecompile::new_single(precompile));
             continue;
         }
@@ -213,10 +230,11 @@ fn merge_bus_interactions<P: IntoOpenVm>(
                         continue;
                     }
                     let i2 = to_merge_set.get(0).unwrap();
-                    let all_args_same_structure =
-                        i.args.iter().zip_eq(i2.args.iter()).all(|(a1, a2)| {
-                            has_same_structure(strip_guard(&a1), strip_guard(&a2))
-                        });
+                    let all_args_same_structure = i
+                        .args
+                        .iter()
+                        .zip_eq(i2.args.iter())
+                        .all(|(a1, a2)| has_same_structure(strip_guard(&a1), strip_guard(&a2)));
                     if all_args_same_structure {
                         // found an exact match
                         to_merge_set.push(i);
@@ -237,10 +255,11 @@ fn merge_bus_interactions<P: IntoOpenVm>(
                     }
                     // check all args have the same structure
                     let i2 = to_merge_set.get(0).unwrap();
-                    let some_args_same_structure =
-                        i.args.iter().zip_eq(i2.args.iter()).any(|(a1, a2)| {
-                            has_same_structure(strip_guard(&a1), strip_guard(&a2))
-                        });
+                    let some_args_same_structure = i
+                        .args
+                        .iter()
+                        .zip_eq(i2.args.iter())
+                        .any(|(a1, a2)| has_same_structure(strip_guard(&a1), strip_guard(&a2)));
                     if some_args_same_structure {
                         // found a partial match on some args
                         to_merge_set.push(i);
@@ -359,16 +378,11 @@ fn is_valid_guarded<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> bool {
             left,
             op: AlgebraicBinaryOperator::Mul,
             ..
-        }) => {
-            match left.as_ref() {
-                AlgebraicExpression::Reference(AlgebraicReference {
-                    name,
-                    ..
-                }) => {
-                    name.starts_with("is_valid")
-                }
-                _ => false,
+        }) => match left.as_ref() {
+            AlgebraicExpression::Reference(AlgebraicReference { name, .. }) => {
+                name.starts_with("is_valid")
             }
+            _ => false,
         },
         _ => false,
     }
@@ -376,7 +390,10 @@ fn is_valid_guarded<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> bool {
 
 /// takes an expression `is_valid * expr` and canonicalizes the right side
 fn canonicalize_expression<P: IntoOpenVm>(expr: &mut AlgebraicExpression<P>) {
-    assert!(is_valid_guarded(expr), "not left guarded by is_valid: {expr}");
+    assert!(
+        is_valid_guarded(expr),
+        "not left guarded by is_valid: {expr}"
+    );
     let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
         left: _is_valid,
         op: AlgebraicBinaryOperator::Mul,
diff --git a/openvm/src/powdr_extension/vm.rs b/openvm/src/powdr_extension/vm.rs
index f46ccb53a1..b907014573 100644
--- a/openvm/src/powdr_extension/vm.rs
+++ b/openvm/src/powdr_extension/vm.rs
@@ -107,9 +107,7 @@ pub struct PowdrStackedPrecompile<P: IntoOpenVm> {
 }
 
 impl<P: IntoOpenVm> PowdrStackedPrecompile<P> {
-    pub fn new_single(
-        precompile: PowdrPrecompile<P>,
-    ) -> Self {
+    pub fn new_single(precompile: PowdrPrecompile<P>) -> Self {
         Self {
             machine: precompile.machine.clone(),
             precompiles: BTreeMap::from([(precompile.opcode.clone(), precompile)]),

From 428550b769bc234da3f2f521d546ba6298a5e8bb Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Tue, 17 Jun 2025 15:43:15 -0300
Subject: [PATCH 25/39] cleanup

---
 openvm/src/customize_exe/air_stacking.rs | 83 ++++++++----------------
 1 file changed, 27 insertions(+), 56 deletions(-)

diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index c5e3c5ecfa..83295af9dc 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -32,10 +32,7 @@ pub fn air_stacking<P: IntoOpenVm>(
     assert!(!extensions.is_empty());
 
     extensions.iter_mut().for_each(|ext| {
-        ext.machine
-            .constraints
-            .iter_mut()
-            .for_each(|c| canonicalize_expression(&mut c.expr));
+        assert!(ext.machine.constraints.iter_mut().all(|c| is_valid_guarded(&c.expr)))
     });
 
     // create apc groups by number of columns
@@ -60,7 +57,7 @@ pub fn air_stacking<P: IntoOpenVm>(
         });
         let mut column_assigner = ColumnAssigner::default();
         g.iter_mut().for_each(|ext| {
-            column_assigner.assign_pcp(ext);
+            column_assigner.assign_precompile(ext);
         });
     });
 
@@ -136,10 +133,7 @@ pub fn air_stacking<P: IntoOpenVm>(
                 .bus_interactions
                 .iter_mut()
                 .for_each(|interaction| {
-                    interaction.args.iter_mut().for_each(|arg| {
-                        *arg = is_valid.clone() * arg.clone();
-                        canonicalize_expression(arg);
-                    });
+                    interaction.args.iter_mut().for_each(|arg| { *arg = is_valid.clone() * arg.clone(); });
                 });
 
             is_valid_sum = is_valid_sum
@@ -288,35 +282,36 @@ fn merge_bus_interactions<P: IntoOpenVm>(
             arg2: AlgebraicExpression<P>,
         ) -> AlgebraicExpression<P> {
             match arg1 {
-                // expr * is_valid
+                // is_valid * expr
                 AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                    left: left1,
+                    left: is_valid1,
                     op: AlgebraicBinaryOperator::Mul,
-                    right: is_valid1,
+                    right: right1,
                 }) => {
                     let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                        left: left2,
+                        left: is_valid2,
                         op: AlgebraicBinaryOperator::Mul,
-                        right: is_valid2,
+                        right: right2,
                     }) = arg2
                     else {
                         panic!("Expected binary operation for arg2, got: {arg2}");
                     };
-                    if has_same_structure(&left1, &left2) {
-                        *left1 * (*is_valid1 + *is_valid2)
+                    if has_same_structure(&right1, &right2) {
+                        // when the expressions match, we just add the new guard
+                        (*is_valid1 + *is_valid2) * *right1
                     } else {
                         AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                            left: left1,
+                            left: is_valid1,
                             op: AlgebraicBinaryOperator::Mul,
-                            right: is_valid1,
+                            right: right1,
                         }) + AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                            left: left2,
+                            left: is_valid2,
                             op: AlgebraicBinaryOperator::Mul,
-                            right: is_valid2,
+                            right: right2,
                         })
                     }
                 }
-                // exprA * is_validA + exprB * is_validB + ...
+                // is_validA * exprA + is_validB * exprB ... we can't remove the guards
                 AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
                     left: _,
                     op: AlgebraicBinaryOperator::Add,
@@ -326,17 +321,18 @@ fn merge_bus_interactions<P: IntoOpenVm>(
             }
         }
 
+        /// if the bus arg is of the form `is_valid * expr`, we can remove the is_valid part
         fn simplify_arg<P: IntoOpenVm>(arg: AlgebraicExpression<P>) -> AlgebraicExpression<P> {
-            // expr * is_valid_expr
+            // is_valid_expr * expr
             if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                left,
+                left: _is_valid_expr,
                 op: AlgebraicBinaryOperator::Mul,
-                right: _is_valid_expr,
+                right,
             }) = arg
             {
-                *left
+                *right
             } else {
-                // sum of expr * is_validN, can't simplify
+                // sum of different is_valid * expr, can't simplify
                 arg
             }
         }
@@ -371,7 +367,7 @@ fn merge_bus_interactions<P: IntoOpenVm>(
     result
 }
 
-/// true if expression is of the form `is_valid * some_expr`.
+/// true if expression is of the form `is_valid_col * some_expr`.
 fn is_valid_guarded<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> bool {
     match expr {
         AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
@@ -388,33 +384,6 @@ fn is_valid_guarded<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> bool {
     }
 }
 
-/// takes an expression `is_valid * expr` and canonicalizes the right side
-fn canonicalize_expression<P: IntoOpenVm>(expr: &mut AlgebraicExpression<P>) {
-    assert!(
-        is_valid_guarded(expr),
-        "not left guarded by is_valid: {expr}"
-    );
-    let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-        left: _is_valid,
-        op: AlgebraicBinaryOperator::Mul,
-        right,
-    }) = expr
-    else {
-        panic!("expected a Mul operation, got: {expr}");
-    };
-    right.pre_visit_expressions_mut(&mut |expr| {
-        if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) =
-            expr
-        {
-            if *op == AlgebraicBinaryOperator::Mul
-                || *op == AlgebraicBinaryOperator::Add && left > right
-            {
-                std::mem::swap(left, right);
-            }
-        }
-    });
-}
-
 /// assumes constraints are of the form `is_valid_expr * some_expr`
 fn join_constraints<P: IntoOpenVm>(
     mut constraints: Vec<SymbolicConstraint<P>>,
@@ -597,13 +566,15 @@ fn expr_poly_id_by_order<P: IntoOpenVm>(
 }
 
 #[derive(Default)]
+/// Remaps the columns of precompiles trying to take into account the structure
+/// of constraints and bus interactions of other precompiles
 struct ColumnAssigner<P: IntoOpenVm> {
-    // original PCPs
     pcps: Vec<PowdrPrecompile<P>>,
 }
 
 impl<P: IntoOpenVm> ColumnAssigner<P> {
-    fn assign_pcp(&mut self, pcp: &mut PowdrPrecompile<P>) {
+    /// Assigns columns to the given precompile and updates the assigner
+    fn assign_precompile(&mut self, pcp: &mut PowdrPrecompile<P>) {
         let mut mappings = BiMap::new();
         // for each constraint, we want to find if there is a constraint in
         // another machine with the same structure that we can assign the same

From 5193df61f3b30fb4e006bdda722e1dfe769b4638 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 18 Jun 2025 11:04:59 -0300
Subject: [PATCH 26/39] normalize expr

---
 expression/src/lib.rs                    | 162 +++++++++++++++++++++++
 openvm/src/customize_exe/air_stacking.rs |  20 ++-
 2 files changed, 180 insertions(+), 2 deletions(-)

diff --git a/expression/src/lib.rs b/expression/src/lib.rs
index 9b094243cf..57afa95dee 100644
--- a/expression/src/lib.rs
+++ b/expression/src/lib.rs
@@ -127,6 +127,13 @@ impl<T, R> AlgebraicExpression<T, R> {
     }
 }
 
+impl<T: Clone + Ord, R: Clone + Ord> AlgebraicExpression<T, R> {
+    /// Get a canonical form by expanding the expression and reordering operands.
+    pub fn normalize(self) -> Self {
+        normalize(expand(self))
+    }
+}
+
 impl<T, R> ops::Add for AlgebraicExpression<T, R> {
     type Output = Self;
 
@@ -164,3 +171,158 @@ impl<T, R> From<T> for AlgebraicExpression<T, R> {
         AlgebraicExpression::Number(value)
     }
 }
+
+// normalizing an expression
+
+// Helper to flatten Add trees
+fn flatten_add<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>, acc: &mut Vec<AlgebraicExpression<T, R>>) {
+    match expr {
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left,
+            op: AlgebraicBinaryOperator::Add,
+            right,
+        }) => {
+            flatten_add(*left, acc);
+            flatten_add(*right, acc);
+        }
+        _ => acc.push(normalize(expr)),
+    }
+}
+
+// Helper to flatten Mul trees
+fn flatten_mul<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>, acc: &mut Vec<AlgebraicExpression<T, R>>) {
+    match expr {
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left,
+            op: AlgebraicBinaryOperator::Mul,
+            right,
+        }) => {
+            flatten_mul(*left, acc);
+            flatten_mul(*right, acc);
+        }
+        _ => acc.push(normalize(expr)),
+    }
+}
+
+fn normalize<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>) -> AlgebraicExpression<T, R> {
+    match expr {
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left,
+            op: AlgebraicBinaryOperator::Add,
+            right,
+        }) => {
+            let mut terms = Vec::new();
+            flatten_add(*left, &mut terms);
+            flatten_add(*right, &mut terms);
+            terms.sort();
+            terms.into_iter().reduce(|a, b| a + b).unwrap()
+        }
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left,
+            op: AlgebraicBinaryOperator::Mul,
+            right,
+        }) => {
+            let mut factors = Vec::new();
+            flatten_mul(*left, &mut factors);
+            flatten_mul(*right, &mut factors);
+            factors.sort();
+            factors.into_iter().reduce(|a, b| a * b).unwrap()
+        }
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left,
+            op: AlgebraicBinaryOperator::Sub,
+            right,
+        }) => {
+            normalize(*left) - normalize(*right)
+        }
+        _ => expr,
+    }
+}
+
+fn expand<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>) -> AlgebraicExpression<T, R> {
+    match expr {
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left,
+            op: AlgebraicBinaryOperator::Mul,
+            right,
+        }) => {
+            let left = expand(*left);
+            let right = expand(*right);
+            match (left, right) {
+                // (a + b)*c ==> a*c + b*c
+                (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                    left: left1,
+                    op: AlgebraicBinaryOperator::Add,
+                    right: right1,
+                }), right) => expand(*left1 * right.clone()) + expand(*right1 * right),
+                // a*(b + c) ==> a*b + a*c
+                (left, AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                    left: left1,
+                    op: AlgebraicBinaryOperator::Add,
+                    right: right1,
+                })) => expand(left.clone() * *left1) + expand(left.clone() * *right1),
+                (left, right) => left * right,
+            }
+        }
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left,
+            op: AlgebraicBinaryOperator::Add,
+            right,
+        }) => expand(*left) + expand(*right),
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left,
+            op: AlgebraicBinaryOperator::Sub,
+            right,
+        }) => expand(*left) - expand(*right),
+        _ => expr,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    type Expr = AlgebraicExpression<u32, String>;
+
+    fn v(name: &str) -> Expr {
+        Expr::Reference(name.to_string())
+    }
+
+    #[test]
+    fn test_normalize() {
+        let expr1 = normalize(v("c") + v("b") + v("a"));
+        let expr2 = normalize(v("a") + v("b") + v("c"));
+        assert_eq!(expr1, expr2);
+
+        let expr1 = normalize(v("c") * v("b") * v("a"));
+        let expr2 = normalize(v("a") * v("b") * v("c"));
+        assert_eq!(expr1, expr2);
+    }
+
+    #[test]
+    fn test_expand() {
+        let e1 = expand(v("a") * (v("b") + v("c")));
+        let e2 = v("a") * v("b") + v("a") * v("c");
+        assert_eq!(e1, e2);
+
+        let e1 = expand((v("a") + v("b")) * v("c"));
+        let e2 = v("a") * v("c") + v("b") * v("c");
+        assert_eq!(e1, e2);
+
+        // (a + b) * (c + d)
+        let e1 = expand((v("a") + v("b")) * (v("c") + v("d")));
+        let e2 = v("a") * v("c") + v("a") * v("d") + v("b") * v("c") + v("b") * v("d");
+        assert_eq!(normalize(e1), normalize(e2));
+
+        // (a + b) * (c + d) * (e + f)
+        let e1 = expand((v("a") + v("b")) * (v("c") + v("d")) * (v("e") + v("f")));
+        let e2 = v("a") * v("c") * v("e")
+            + v("a") * v("c") * v("f")
+            + v("a") * v("d") * v("e")
+            + v("a") * v("d") * v("f")
+            + v("b") * v("c") * v("e")
+            + v("b") * v("c") * v("f")
+            + v("b") * v("d") * v("e")
+            + v("b") * v("d") * v("f");
+        assert_eq!(normalize(e1), normalize(e2));
+    }
+}
diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index dc63f792a7..dc342debf1 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -32,7 +32,9 @@ pub fn air_stacking<P: IntoOpenVm>(
     assert!(!extensions.is_empty());
 
     extensions.iter_mut().for_each(|ext| {
-        assert!(ext.machine.constraints.iter_mut().all(|c| is_valid_guarded(&c.expr)))
+        ext.machine.constraints.iter_mut().for_each(|c| {
+            normalize_guarded(&mut c.expr);
+        })
     });
 
     // create apc groups by number of columns
@@ -128,7 +130,7 @@ pub fn air_stacking<P: IntoOpenVm>(
                 .bus_interactions
                 .iter_mut()
                 .for_each(|interaction| {
-                    interaction.args.iter_mut().for_each(|arg| { *arg = is_valid.clone() * arg.clone(); });
+                    interaction.args.iter_mut().for_each(|arg| { *arg = is_valid.clone() * arg.clone().normalize(); });
                 });
 
             is_valid_sum = is_valid_sum
@@ -379,6 +381,20 @@ fn is_valid_guarded<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> bool {
     }
 }
 
+fn normalize_guarded<P: IntoOpenVm>(expr:&mut AlgebraicExpression<P>) {
+    assert!(is_valid_guarded(expr), "not left guarded by is_valid: {expr}");
+    match expr {
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left: _is_valid,
+            op: AlgebraicBinaryOperator::Mul,
+            right,
+        }) => {
+            *right.as_mut() = (*right.clone()).normalize();
+        },
+        _ => unreachable!(),
+    }
+}
+
 /// assumes constraints are of the form `is_valid_expr * some_expr`
 fn join_constraints<P: IntoOpenVm>(
     mut constraints: Vec<SymbolicConstraint<P>>,

From 7ea0748d893382f83dbfed20f6fb99090b327cf2 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 18 Jun 2025 14:52:08 -0300
Subject: [PATCH 27/39] cleanup

---
 autoprecompiles/src/lib.rs         |  9 +--------
 expression/src/lib.rs              |  6 ++++--
 openvm/Cargo.toml                  |  1 -
 openvm/src/powdr_extension/chip.rs | 25 ++++++++++++++++---------
 4 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/autoprecompiles/src/lib.rs b/autoprecompiles/src/lib.rs
index af9fc680c2..c56a69043a 100644
--- a/autoprecompiles/src/lib.rs
+++ b/autoprecompiles/src/lib.rs
@@ -254,9 +254,6 @@ pub fn build<T: FieldElement, B: BusInteractionHandler<T> + IsBusStateful<T> + C
         &vm_config.bus_map,
     );
 
-    println!("before optimizer - machine constraints degree: {}", machine.constraints.iter().map(|c| c.expr.degree()).max().unwrap_or(0));
-    println!("before optimizer - machine bus args degree: {}", machine.bus_interactions.iter().map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0)).max().unwrap_or(0));
-
     let machine = optimizer::optimize(
         machine,
         vm_config.bus_interaction_handler,
@@ -265,9 +262,6 @@ pub fn build<T: FieldElement, B: BusInteractionHandler<T> + IsBusStateful<T> + C
         &vm_config.bus_map,
     )?;
 
-    println!("after optimizer - machine constraints degree: {}", machine.constraints.iter().map(|c| c.expr.degree()).max().unwrap_or(0));
-    println!("after optimizer - machine bus args degree: {}", machine.bus_interactions.iter().map(|b| b.args.iter().map(|a| a.degree()).max().unwrap_or(0)).max().unwrap_or(0));
-
     // add guards to constraints that are not satisfied by zeroes
     let machine = add_guards(machine, vm_config.bus_map, strict_is_valid_guards);
 
@@ -368,8 +362,7 @@ fn add_guards<T: FieldElement>(
         .unwrap();
     program_bus_send.mult = is_valid.clone();
 
-    let mut is_valid_mults = vec![];
-
+    let mut is_valid_mults: Vec<SymbolicConstraint<T>> = Vec::new();
     for b in &mut machine.bus_interactions {
         // already handled exec and pc lookup bus types
         if b.id != exec_bus_id && b.id != pc_lookup_bus_id {
diff --git a/expression/src/lib.rs b/expression/src/lib.rs
index 57afa95dee..f33e29502a 100644
--- a/expression/src/lib.rs
+++ b/expression/src/lib.rs
@@ -172,8 +172,6 @@ impl<T, R> From<T> for AlgebraicExpression<T, R> {
     }
 }
 
-// normalizing an expression
-
 // Helper to flatten Add trees
 fn flatten_add<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>, acc: &mut Vec<AlgebraicExpression<T, R>>) {
     match expr {
@@ -296,6 +294,10 @@ mod tests {
         let expr1 = normalize(v("c") * v("b") * v("a"));
         let expr2 = normalize(v("a") * v("b") * v("c"));
         assert_eq!(expr1, expr2);
+
+        let expr1 = normalize(v("a") + v("b") * v("c"));
+        let expr2 = normalize(v("c") * v("b") + v("a"));
+        assert_eq!(expr1, expr2);
     }
 
     #[test]
diff --git a/openvm/Cargo.toml b/openvm/Cargo.toml
index 53af84274b..944ede6437 100644
--- a/openvm/Cargo.toml
+++ b/openvm/Cargo.toml
@@ -43,7 +43,6 @@ powdr-number.workspace = true
 powdr-riscv-elf.workspace = true
 powdr-autoprecompiles.workspace = true
 powdr-constraint-solver.workspace = true
-powdr-pilopt.workspace = true
 
 eyre = "0.6.12"
 serde = "1.0.217"
diff --git a/openvm/src/powdr_extension/chip.rs b/openvm/src/powdr_extension/chip.rs
index 55b6b3cb5b..183e602541 100644
--- a/openvm/src/powdr_extension/chip.rs
+++ b/openvm/src/powdr_extension/chip.rs
@@ -62,14 +62,21 @@ impl<P: IntoOpenVm> PowdrChip<P> {
         periphery: PowdrPeripheryInstances,
     ) -> Self {
         let air = PowdrAir::new(precompile.machine);
-        let name = format!(
-            "StackedPrecompile_{}",
-            precompile
-                .precompiles
-                .keys()
-                .map(|o| o.global_opcode())
-                .join("_")
-        );
+
+        let name = if precompile.precompiles.len() == 1 {
+            // single precompile, just use its name
+            precompile.precompiles.values().next().unwrap().name.clone()
+        } else {
+            // TODO: this name can be quite big depending on the number of precompiles
+            format!(
+                "StackedPrecompile_{}",
+                precompile
+                    .precompiles
+                    .keys()
+                    .map(|o| o.global_opcode())
+                    .join("_")
+            )
+        };
 
         let executors = precompile
             .precompiles
@@ -88,7 +95,6 @@ impl<P: IntoOpenVm> PowdrChip<P> {
             .collect();
 
         Self {
-            // TODO: proper name
             name,
             air: Arc::new(air),
             executors,
@@ -121,6 +127,7 @@ impl<P: IntoOpenVm> InstructionExecutor<OpenVmField<P>> for PowdrChip<P> {
 
 impl<P: IntoOpenVm> ChipUsageGetter for PowdrChip<P> {
     fn air_name(&self) -> String {
+        // TODO: this name can be quite big depending on the number of stacked precompiles
         format!("powdr_air_for_opcodes_{}", self.executors.keys().join("_"))
     }
 

From 4473e9a8a9ebe309e9f2a88702165cecdb2574eb Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 18 Jun 2025 15:04:54 -0300
Subject: [PATCH 28/39] witgen cleanup

---
 openvm/src/powdr_extension/chip.rs | 103 +++++++++++++++++------------
 openvm/src/powdr_extension/vm.rs   |   2 +-
 2 files changed, 61 insertions(+), 44 deletions(-)

diff --git a/openvm/src/powdr_extension/chip.rs b/openvm/src/powdr_extension/chip.rs
index 183e602541..b95edfd359 100644
--- a/openvm/src/powdr_extension/chip.rs
+++ b/openvm/src/powdr_extension/chip.rs
@@ -152,54 +152,71 @@ where
         tracing::trace!("Generating air proof input for PowdrChip {}", self.name);
 
         let num_records = self.current_trace_height();
-        let height = next_power_of_two_or_zero(num_records);
         let width = self.trace_width();
-        let mut values = Val::<SC>::zero_vec(height * width);
-        let mut values_curr_record = 0;
-        // this is just for sanity checking later
-        let all_is_valid_ids = self
-            .executors
-            .values()
-            .map(|executor| executor.is_valid_poly_id)
-            .collect::<Vec<_>>();
-
-        for (_opcode, executor) in self.executors {
-            let executor_is_valid_id = executor.is_valid_poly_id;
-            let executor_calls = executor.number_of_calls();
-            let mut trace = executor.generate_witness::<SC>(
+        let height = next_power_of_two_or_zero(num_records);
+
+        let trace = if self.executors.len() == 1 {
+            // non-stacked precompile
+            let executor = self.executors.into_values().next().unwrap();
+            executor.generate_witness::<SC>(
                 &self.air.column_index_by_poly_id,
                 &self.air.machine.bus_interactions,
-            );
-
-            // copy to main trace
-            values
-                .chunks_mut(width)
-                .skip(values_curr_record)
-                .zip(trace.rows_mut().take(executor_calls))
-                .for_each(|(values_row, trace_row)| {
-                    assert!(values_row.len() >= trace_row.len());
-                    // copy the trace row to the main trace row
-                    values_row.copy_from_slice(trace_row);
-
-                    // check that only the correct is valid is set to ONE
-                    for id in all_is_valid_ids.iter() {
-                        if *id == executor_is_valid_id {
-                            assert_eq!(
-                                values_row[self.air.column_index_by_poly_id[id]],
-                                <Val<SC>>::ONE
-                            );
-                        } else {
-                            assert_eq!(
-                                values_row[self.air.column_index_by_poly_id[id]],
-                                <Val<SC>>::ZERO
-                            );
+            )
+        } else {
+            // stacked precompiles, reserve space for the full trace and call each executor in turn
+
+            // this is just for sanity checking later
+            let all_is_valid_ids = self
+                .executors
+                .values()
+                .map(|executor| executor.is_valid_poly_id)
+                .collect::<Vec<_>>();
+
+            // TODO: we could avoid copying by having witgen take this vec as input to write the values in
+            let mut values = Val::<SC>::zero_vec(height * width);
+            let mut values_curr_record = 0;
+
+            for (_opcode, executor) in self.executors {
+                let executor_is_valid_id = executor.is_valid_poly_id;
+                let executor_calls = executor.number_of_calls();
+                let mut trace = executor.generate_witness::<SC>(
+                    &self.air.column_index_by_poly_id,
+                    &self.air.machine.bus_interactions,
+                );
+
+                // copy to main trace
+                values
+                    .chunks_mut(width)
+                    .skip(values_curr_record)
+                    .zip(trace.rows_mut().take(executor_calls))
+                    .for_each(|(values_row, trace_row)| {
+                        assert!(values_row.len() >= trace_row.len());
+                        // copy the trace row to the main trace row
+                        values_row.copy_from_slice(trace_row);
+
+                        // check that only the correct is valid is set to ONE
+                        for id in all_is_valid_ids.iter() {
+                            if *id == executor_is_valid_id {
+                                assert_eq!(
+                                    values_row[self.air.column_index_by_poly_id[id]],
+                                    <Val<SC>>::ONE
+                                );
+                            } else {
+                                assert_eq!(
+                                    values_row[self.air.column_index_by_poly_id[id]],
+                                    <Val<SC>>::ZERO
+                                );
+                            }
                         }
-                    }
-                });
-            values_curr_record += executor_calls;
-        }
+                    });
+                values_curr_record += executor_calls;
+            }
+
+            RowMajorMatrix::new(values, width)
+        };
+
+        assert_eq!(trace.width(), width);
 
-        let trace = RowMajorMatrix::new(values, width);
         AirProofInput::simple(trace, vec![])
     }
 }
diff --git a/openvm/src/powdr_extension/vm.rs b/openvm/src/powdr_extension/vm.rs
index 3ea99c1020..918749ebc4 100644
--- a/openvm/src/powdr_extension/vm.rs
+++ b/openvm/src/powdr_extension/vm.rs
@@ -100,7 +100,7 @@ impl<P: IntoOpenVm> PowdrPrecompile<P> {
 pub struct PowdrStackedPrecompile<P: IntoOpenVm> {
     /// stacked precompiles by opcode
     pub precompiles: BTreeMap<PowdrOpcode, PowdrPrecompile<P>>,
-    /// constraints
+    /// joined constraints / interactions
     pub machine: SymbolicMachine<P>,
 }
 

From 35c8ba02634d214fa45689054a523ad7e5e49598 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 18 Jun 2025 15:38:25 -0300
Subject: [PATCH 29/39] cleanup

---
 openvm/src/customize_exe/air_stacking.rs | 146 ++++++++++++-----------
 1 file changed, 77 insertions(+), 69 deletions(-)

diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index dc342debf1..eeaf76865b 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -1,8 +1,7 @@
 use std::collections::{BTreeSet, HashMap};
 
 use powdr_autoprecompiles::{
-    expression::{AlgebraicExpression, AlgebraicReference},
-    simplify_expression, SymbolicBusInteraction, SymbolicConstraint, SymbolicMachine,
+    expression::{AlgebraicExpression, AlgebraicReference}, powdr::make_bool, simplify_expression, SymbolicBusInteraction, SymbolicConstraint, SymbolicMachine
 };
 
 use powdr_autoprecompiles::powdr::UniqueReferences;
@@ -146,10 +145,7 @@ pub fn air_stacking<P: IntoOpenVm>(
         tracing::debug!("After joining constraints: {}", stacked_constraints.len());
 
         // enforce only one is_valid is active
-        let one = AlgebraicExpression::Number(P::ONE);
-        let one_hot_is_valid = (one - is_valid_sum.clone().unwrap()) * is_valid_sum.unwrap();
-
-        stacked_constraints.push(one_hot_is_valid.into());
+        stacked_constraints.push(make_bool(is_valid_sum.unwrap()).into());
 
         tracing::debug!(
             "Stacked chip has {} bus interactions",
@@ -180,6 +176,66 @@ pub fn air_stacking<P: IntoOpenVm>(
 fn merge_bus_interactions<P: IntoOpenVm>(
     interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>,
 ) -> Vec<SymbolicBusInteraction<P>> {
+    /// helper: combine args taking into acount if they have the same guarded expression
+    fn merge_args<P: IntoOpenVm>(
+        arg1: AlgebraicExpression<P>,
+        arg2: AlgebraicExpression<P>,
+    ) -> AlgebraicExpression<P> {
+        match arg1 {
+            // is_valid * expr
+            AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                left: is_valid1,
+                op: AlgebraicBinaryOperator::Mul,
+                right: right1,
+            }) => {
+                let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                    left: is_valid2,
+                    op: AlgebraicBinaryOperator::Mul,
+                    right: right2,
+                }) = arg2
+                else {
+                    panic!("Expected binary operation for arg2, got: {arg2}");
+                };
+                if has_same_structure(&right1, &right2) {
+                    // when the expressions match, we just add the new guard
+                    (*is_valid1 + *is_valid2) * *right1
+                } else {
+                    AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                        left: is_valid1,
+                        op: AlgebraicBinaryOperator::Mul,
+                        right: right1,
+                    }) + AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                        left: is_valid2,
+                        op: AlgebraicBinaryOperator::Mul,
+                        right: right2,
+                    })
+                }
+            }
+            // is_validA * exprA + is_validB * exprB ... we can't remove the guards
+            AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                left: _,
+                op: AlgebraicBinaryOperator::Add,
+                right: _,
+            }) => arg1 + arg2,
+            _ => unreachable!(),
+        }
+    }
+
+    /// helper: if the merged bus arg is a multiplication, it is of the form `is_valid_expr * expr`, and we can remove the guard
+    fn simplify_arg<P: IntoOpenVm>(arg: AlgebraicExpression<P>) -> AlgebraicExpression<P> {
+        if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left: _is_valid_expr,
+            op: AlgebraicBinaryOperator::Mul,
+            right,
+        }) = arg
+        {
+            *right
+        } else {
+            // sum of different is_valid * expr, can't simplify
+            arg
+        }
+    }
+
     // split interactions by bus/args len
     let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> =
         Default::default();
@@ -209,6 +265,10 @@ fn merge_bus_interactions<P: IntoOpenVm>(
             .into_iter()
             .map(|i| vec![i])
             .collect_vec();
+        // In the following, we go through each machine and try to find for each
+        // of its interactions, an existing interaction to merge with. First we
+        // try to find one where all args are compatible, if not found, one with
+        // a partial match, and finally just merge with any.
         for machine_interactions in interactions_by_machine {
             // merge sets already used by this machine
             let mut used = BTreeSet::new();
@@ -233,6 +293,7 @@ fn merge_bus_interactions<P: IntoOpenVm>(
                         continue 'outer;
                     }
                 }
+                // no exact match for the args found, try a partial match
                 try_partial_match.push(i);
             }
 
@@ -274,66 +335,6 @@ fn merge_bus_interactions<P: IntoOpenVm>(
             }
         }
 
-        fn merge_args<P: IntoOpenVm>(
-            arg1: AlgebraicExpression<P>,
-            arg2: AlgebraicExpression<P>,
-        ) -> AlgebraicExpression<P> {
-            match arg1 {
-                // is_valid * expr
-                AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                    left: is_valid1,
-                    op: AlgebraicBinaryOperator::Mul,
-                    right: right1,
-                }) => {
-                    let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                        left: is_valid2,
-                        op: AlgebraicBinaryOperator::Mul,
-                        right: right2,
-                    }) = arg2
-                    else {
-                        panic!("Expected binary operation for arg2, got: {arg2}");
-                    };
-                    if has_same_structure(&right1, &right2) {
-                        // when the expressions match, we just add the new guard
-                        (*is_valid1 + *is_valid2) * *right1
-                    } else {
-                        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                            left: is_valid1,
-                            op: AlgebraicBinaryOperator::Mul,
-                            right: right1,
-                        }) + AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                            left: is_valid2,
-                            op: AlgebraicBinaryOperator::Mul,
-                            right: right2,
-                        })
-                    }
-                }
-                // is_validA * exprA + is_validB * exprB ... we can't remove the guards
-                AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                    left: _,
-                    op: AlgebraicBinaryOperator::Add,
-                    right: _,
-                }) => arg1 + arg2,
-                _ => unreachable!(),
-            }
-        }
-
-        /// if the bus arg is of the form `is_valid * expr`, we can remove the is_valid part
-        fn simplify_arg<P: IntoOpenVm>(arg: AlgebraicExpression<P>) -> AlgebraicExpression<P> {
-            // is_valid_expr * expr
-            if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                left: _is_valid_expr,
-                op: AlgebraicBinaryOperator::Mul,
-                right,
-            }) = arg
-            {
-                *right
-            } else {
-                // sum of different is_valid * expr, can't simplify
-                arg
-            }
-        }
-
         // merge each set of interactions
         for set in to_merge {
             let id = set[0].id;
@@ -401,7 +402,9 @@ fn join_constraints<P: IntoOpenVm>(
 ) -> Vec<SymbolicConstraint<P>> {
     constraints.sort();
     let mut result: Vec<SymbolicConstraint<P>> = vec![];
-    'outer: for c1 in constraints {
+    for c1 in constraints {
+        // try to find an existing constraint with the same expression but guarded by a different is valid to join with
+        let mut found_match = false;
         for c2 in result.iter_mut() {
             if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
                 left: is_valid1,
@@ -417,16 +420,21 @@ fn join_constraints<P: IntoOpenVm>(
                 {
                     if r1 == r2 {
                         *is_valid2 = Box::new(*is_valid1.clone() + *is_valid2.clone());
-                        continue 'outer;
+                        found_match = true;
+                        break;
                     }
                 }
             }
         }
-        result.push(c1);
+        if !found_match {
+            // no matching constraint to join with, just add it as is
+            result.push(c1);
+        }
     }
     result
 }
 
+/// Compare two expressions for equality disregarding the reference name.
 fn has_same_structure<P: IntoOpenVm>(
     expr1: &AlgebraicExpression<P>,
     expr2: &AlgebraicExpression<P>,

From c13951485f12544a37b8f758b82be17ef8761d88 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 18 Jun 2025 15:52:44 -0300
Subject: [PATCH 30/39] cleanup

---
 openvm/src/customize_exe/air_stacking.rs | 135 +++++++++++------------
 1 file changed, 63 insertions(+), 72 deletions(-)

diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index eeaf76865b..62010b258d 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -176,66 +176,6 @@ pub fn air_stacking<P: IntoOpenVm>(
 fn merge_bus_interactions<P: IntoOpenVm>(
     interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>,
 ) -> Vec<SymbolicBusInteraction<P>> {
-    /// helper: combine args taking into acount if they have the same guarded expression
-    fn merge_args<P: IntoOpenVm>(
-        arg1: AlgebraicExpression<P>,
-        arg2: AlgebraicExpression<P>,
-    ) -> AlgebraicExpression<P> {
-        match arg1 {
-            // is_valid * expr
-            AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                left: is_valid1,
-                op: AlgebraicBinaryOperator::Mul,
-                right: right1,
-            }) => {
-                let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                    left: is_valid2,
-                    op: AlgebraicBinaryOperator::Mul,
-                    right: right2,
-                }) = arg2
-                else {
-                    panic!("Expected binary operation for arg2, got: {arg2}");
-                };
-                if has_same_structure(&right1, &right2) {
-                    // when the expressions match, we just add the new guard
-                    (*is_valid1 + *is_valid2) * *right1
-                } else {
-                    AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                        left: is_valid1,
-                        op: AlgebraicBinaryOperator::Mul,
-                        right: right1,
-                    }) + AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                        left: is_valid2,
-                        op: AlgebraicBinaryOperator::Mul,
-                        right: right2,
-                    })
-                }
-            }
-            // is_validA * exprA + is_validB * exprB ... we can't remove the guards
-            AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                left: _,
-                op: AlgebraicBinaryOperator::Add,
-                right: _,
-            }) => arg1 + arg2,
-            _ => unreachable!(),
-        }
-    }
-
-    /// helper: if the merged bus arg is a multiplication, it is of the form `is_valid_expr * expr`, and we can remove the guard
-    fn simplify_arg<P: IntoOpenVm>(arg: AlgebraicExpression<P>) -> AlgebraicExpression<P> {
-        if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-            left: _is_valid_expr,
-            op: AlgebraicBinaryOperator::Mul,
-            right,
-        }) = arg
-        {
-            *right
-        } else {
-            // sum of different is_valid * expr, can't simplify
-            arg
-        }
-    }
-
     // split interactions by bus/args len
     let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> =
         Default::default();
@@ -266,16 +206,13 @@ fn merge_bus_interactions<P: IntoOpenVm>(
             .map(|i| vec![i])
             .collect_vec();
         // In the following, we go through each machine and try to find for each
-        // of its interactions, an existing interaction to merge with. First we
-        // try to find one where all args are compatible, if not found, one with
-        // a partial match, and finally just merge with any.
+        // of its interactions an existing interaction to merge it with.
         for machine_interactions in interactions_by_machine {
             // merge sets already used by this machine
             let mut used = BTreeSet::new();
-            // interactions that didn't find an exact match in structure
+            // first: try to find an exact match, where all args have the same expression
             let mut try_partial_match = vec![];
             'outer: for i in machine_interactions {
-                // try to find a call with an exact match
                 for (idx, to_merge_set) in to_merge.iter_mut().enumerate() {
                     if used.contains(&idx) {
                         continue;
@@ -297,15 +234,13 @@ fn merge_bus_interactions<P: IntoOpenVm>(
                 try_partial_match.push(i);
             }
 
-            // interactions that did not find a partial structure match
+            // then: try to find a partial match, where some of the args have the same expression
             let mut no_match = vec![];
             'outer: for i in try_partial_match {
-                // didn't find exact match, try one where some args match
                 for (idx, to_merge_set) in to_merge.iter_mut().enumerate() {
                     if used.contains(&idx) {
                         continue;
                     }
-                    // check all args have the same structure
                     let i2 = to_merge_set.get(0).unwrap();
                     let some_args_same_structure = i
                         .args
@@ -322,6 +257,7 @@ fn merge_bus_interactions<P: IntoOpenVm>(
                 no_match.push(i);
             }
 
+            // finally: just pick the first unused one to merge with
             'outer: for i in no_match {
                 // just pick the first unused one
                 for (idx, to_merge_set) in to_merge.iter_mut().enumerate() {
@@ -335,9 +271,11 @@ fn merge_bus_interactions<P: IntoOpenVm>(
             }
         }
 
-        // merge each set of interactions
+        // do the actual merging of the grouped interactions
         for set in to_merge {
+            assert!(set.iter().map(|i| i.id).unique().count() == 1, "grouped interactions should have the same bus");
             let id = set[0].id;
+            // multiplicites are just added
             let mult = simplify_expression(
                 set.iter()
                     .map(|i| i.mult.clone())
@@ -350,13 +288,14 @@ fn merge_bus_interactions<P: IntoOpenVm>(
                 .reduce(|a, b| {
                     a.into_iter()
                         .zip_eq(b)
-                        .map(|(a1, a2)| merge_args(a1, a2))
+                        .map(|(a1, a2)| bus_arg_merge(a1, a2))
                         .collect()
                 })
                 .unwrap();
+            // remove the is_valid guard from the args when possible
             let args = args
                 .into_iter()
-                .map(|arg| simplify_expression(simplify_arg(arg)))
+                .map(|arg| simplify_expression(bus_arg_simplify(arg)))
                 .collect_vec();
             result.push(SymbolicBusInteraction { id, args, mult });
         }
@@ -365,6 +304,58 @@ fn merge_bus_interactions<P: IntoOpenVm>(
     result
 }
 
+/// helper to combine args taking into acount if they have the same guarded expression
+fn bus_arg_merge<P: IntoOpenVm>(
+    arg1: AlgebraicExpression<P>,
+    arg2: AlgebraicExpression<P>,
+) -> AlgebraicExpression<P> {
+    match arg1 {
+        // is_valid * expr
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left: is_valid1,
+            op: AlgebraicBinaryOperator::Mul,
+            right: right1,
+        }) => {
+            let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                left: is_valid2,
+                op: AlgebraicBinaryOperator::Mul,
+                right: right2,
+            }) = arg2
+            else {
+                panic!("Expected binary operation for arg2, got: {arg2}");
+            };
+            if has_same_structure(&right1, &right2) {
+                // when the expressions match, we just add the new guard
+                (*is_valid1 + *is_valid2) * *right1
+            } else {
+                (*is_valid1 * *right1) + (*is_valid2 * *right2)
+            }
+        }
+        // is_validA * exprA + is_validB * exprB ... we can't remove the guards
+        AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+            left: _,
+            op: AlgebraicBinaryOperator::Add,
+            right: _,
+        }) => arg1 + arg2,
+        _ => unreachable!(),
+    }
+}
+
+/// helper: if the merged bus arg is a multiplication, it is of the form `is_valid_expr * expr`, and we can remove the guard
+fn bus_arg_simplify<P: IntoOpenVm>(arg: AlgebraicExpression<P>) -> AlgebraicExpression<P> {
+    if let AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+        left: _is_valid_expr,
+        op: AlgebraicBinaryOperator::Mul,
+        right,
+    }) = arg
+    {
+        *right
+    } else {
+        // sum of different is_valid * expr, can't simplify
+        arg
+    }
+}
+
 /// true if expression is of the form `is_valid_col * some_expr`.
 fn is_valid_guarded<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> bool {
     match expr {
@@ -383,7 +374,7 @@ fn is_valid_guarded<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> bool {
 }
 
 fn normalize_guarded<P: IntoOpenVm>(expr:&mut AlgebraicExpression<P>) {
-    assert!(is_valid_guarded(expr), "not left guarded by is_valid: {expr}");
+    assert!(is_valid_guarded(expr), "not left guarded by is_valid col: {expr}");
     match expr {
         AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
             left: _is_valid,

From 44c8c58499dd2aedff43bee8aabc65433c2da58f Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 18 Jun 2025 16:27:20 -0300
Subject: [PATCH 31/39] clenaup

---
 openvm/src/customize_exe/air_stacking.rs | 72 ++++++++++++++----------
 1 file changed, 41 insertions(+), 31 deletions(-)

diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index 62010b258d..065753bb52 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -151,7 +151,7 @@ pub fn air_stacking<P: IntoOpenVm>(
             "Stacked chip has {} bus interactions",
             interactions_by_machine.iter().flatten().count()
         );
-        let stacked_interactions = merge_bus_interactions(interactions_by_machine);
+        let stacked_interactions = join_bus_interactions(interactions_by_machine);
         tracing::debug!("After merging interactions: {}", stacked_interactions.len());
 
         let machine = SymbolicMachine {
@@ -173,10 +173,10 @@ pub fn air_stacking<P: IntoOpenVm>(
 
 /// Merge bus interactions, taking into account args that have the same expression.
 /// When only compatible args are merged, the expression doesn't need an is_valid guard.
-fn merge_bus_interactions<P: IntoOpenVm>(
+fn join_bus_interactions<P: IntoOpenVm>(
     interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>,
 ) -> Vec<SymbolicBusInteraction<P>> {
-    // split interactions by bus/args len
+    // split interactions by bus/args len, and then by machine
     let mut interactions_by_bus: HashMap<_, Vec<Vec<SymbolicBusInteraction<P>>>> =
         Default::default();
 
@@ -190,6 +190,8 @@ fn merge_bus_interactions<P: IntoOpenVm>(
             e.push(v);
         }
     }
+
+    // for each target bus, we sort the machines by number of interactions to that bus
     interactions_by_bus
         .values_mut()
         .for_each(|interactions_by_machine| {
@@ -197,8 +199,11 @@ fn merge_bus_interactions<P: IntoOpenVm>(
         });
 
     let mut result = vec![];
+    // Go through each bus, merging interactions to that bus
     for mut interactions_by_machine in interactions_by_bus.into_values() {
-        // to_merge is a vec of vecs, each inner vec is a set of interactions to be merged
+        // to_merge is a vec of vecs, each inner vec is a set of interactions to be merged.
+        // We initialize it using the machine with the largest number of interactions to that bus.
+        // This is so each subsequent machine can find at least one interaction to merge with for each of its own.
         let mut to_merge = interactions_by_machine
             .pop()
             .unwrap()
@@ -272,38 +277,43 @@ fn merge_bus_interactions<P: IntoOpenVm>(
         }
 
         // do the actual merging of the grouped interactions
-        for set in to_merge {
-            assert!(set.iter().map(|i| i.id).unique().count() == 1, "grouped interactions should have the same bus");
-            let id = set[0].id;
-            // multiplicites are just added
-            let mult = simplify_expression(
-                set.iter()
-                    .map(|i| i.mult.clone())
-                    .reduce(|a, b| a + b)
-                    .unwrap(),
-            );
-            let args = set
-                .into_iter()
-                .map(|i| i.args)
-                .reduce(|a, b| {
-                    a.into_iter()
-                        .zip_eq(b)
-                        .map(|(a1, a2)| bus_arg_merge(a1, a2))
-                        .collect()
-                })
-                .unwrap();
-            // remove the is_valid guard from the args when possible
-            let args = args
-                .into_iter()
-                .map(|arg| simplify_expression(bus_arg_simplify(arg)))
-                .collect_vec();
-            result.push(SymbolicBusInteraction { id, args, mult });
-        }
+        result.extend(to_merge.into_iter().map(merge_bus_interactions));
     }
 
     result
 }
 
+/// helper to combine a set of bus interactions into a single interaction
+fn merge_bus_interactions<P: IntoOpenVm>(
+    interactions: Vec<SymbolicBusInteraction<P>>,
+) -> SymbolicBusInteraction<P> {
+    assert!(interactions.iter().map(|i| i.id).unique().count() == 1, "grouped interactions should have the same bus");
+    let id = interactions[0].id;
+    // multiplicites are just added
+    let mult = simplify_expression(
+        interactions.iter()
+            .map(|i| i.mult.clone())
+            .reduce(|a, b| a + b)
+            .unwrap(),
+    );
+    let args = interactions
+        .into_iter()
+        .map(|i| i.args)
+        .reduce(|a, b| {
+            a.into_iter()
+                .zip_eq(b)
+                .map(|(a1, a2)| bus_arg_merge(a1, a2))
+                .collect()
+        })
+        .unwrap();
+    // remove the is_valid guard from the args when possible
+    let args = args
+        .into_iter()
+        .map(|arg| simplify_expression(bus_arg_simplify(arg)))
+        .collect_vec();
+    SymbolicBusInteraction { id, args, mult }
+}
+
 /// helper to combine args taking into acount if they have the same guarded expression
 fn bus_arg_merge<P: IntoOpenVm>(
     arg1: AlgebraicExpression<P>,

From a5483e9dd4691a9bf08ed8a7b40ec1baf35df612 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 18 Jun 2025 17:51:58 -0300
Subject: [PATCH 32/39] cleanup

---
 openvm/src/customize_exe/air_stacking.rs | 165 +++++++++++------------
 1 file changed, 81 insertions(+), 84 deletions(-)

diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index 065753bb52..5747a9657b 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -36,7 +36,7 @@ pub fn air_stacking<P: IntoOpenVm>(
         })
     });
 
-    // create apc groups by number of columns
+    // group precompiles by number of columns
     let mut groups: HashMap<usize, Vec<PowdrPrecompile<P>>> = Default::default();
     for pcp in extensions {
         let idx = f32::log(
@@ -47,25 +47,11 @@ pub fn air_stacking<P: IntoOpenVm>(
         groups.entry(idx).or_default().push(pcp);
     }
 
-    // sort each group by number of columns
-    groups.values_mut().for_each(|g| {
-        // assign largest pcp first
-        g.sort_by(|pcp1, pcp2| {
-            pcp2.machine
-                .unique_references()
-                .count()
-                .cmp(&pcp1.machine.unique_references().count())
-        });
-        let mut column_assigner = ColumnAssigner::default();
-        g.iter_mut().for_each(|ext| {
-            column_assigner.assign_precompile(ext);
-        });
-    });
-
     let mut result = vec![];
 
+    // transform each group into a stacked precompile
     for mut extensions in groups.into_values() {
-        // group of a single precompile, no stacking
+        // group has a single precompile, no stacking
         if extensions.len() == 1 {
             let precompile = extensions.pop().unwrap();
             tracing::debug!(
@@ -76,6 +62,20 @@ pub fn air_stacking<P: IntoOpenVm>(
             continue;
         }
 
+        // handle largest precompile first
+        extensions.sort_by(|pcp1, pcp2| {
+            pcp2.machine
+                .unique_references()
+                .count()
+                .cmp(&pcp1.machine.unique_references().count())
+        });
+
+        // assign columns to the precompiles in this group
+        let mut column_assigner = ColumnAssigner::default();
+        extensions.iter_mut().for_each(|ext| {
+            column_assigner.assign_precompile(ext);
+        });
+
         let mut stacked_constraints = vec![];
         let mut interactions_by_machine = vec![];
 
@@ -101,7 +101,6 @@ pub fn air_stacking<P: IntoOpenVm>(
 
             // remap is_valid column in constraints and interactions
             let mut remapped = pcp.machine.clone();
-
             remapped.pre_visit_expressions_mut(&mut |expr| {
                 if let AlgebraicExpression::Reference(r) = expr {
                     assert!(r.id <= is_valid_start);
@@ -562,7 +561,6 @@ fn extend_if_no_conflicts(mappings: &mut BiMap<u64, u64>, new_mappings: BiMap<u6
 
 /// changes poly_id to match the order in which they are encountered in the expression.
 /// This allows us to compare two expressions for the same "structure".
-/// e.g., a+a == b+b and a+a != a+b
 fn expr_poly_id_by_order<P: IntoOpenVm>(
     mut expr: AlgebraicExpression<P>,
 ) -> AlgebraicExpression<P> {
@@ -601,85 +599,84 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
         // ids.
         // There may be multiple such constraints, so we try them all (as some
         // of the ids in the current constraint may already be assigned)
-
-        'outer: for c in &pcp.machine.constraints {
-            for pcp2 in self.pcps.iter() {
-                for c2 in &pcp2.machine.constraints {
-                    if has_same_structure(
-                        &expr_poly_id_by_order(c.expr.clone()),
-                        &expr_poly_id_by_order(c2.expr.clone()),
-                    ) {
-                        // Found the same structure, try to extend the column mapping
-                        let new_mappings = create_mapping(&c.expr, &c2.expr);
-                        if extend_if_no_conflicts(&mut mappings, new_mappings) {
-                            // we're done for this constraint
-                            continue 'outer;
-                        }
+        for c in &pcp.machine.constraints {
+            for c2 in self.pcps.iter().map(|pcp| pcp.machine.constraints.iter()).flatten() {
+                if has_same_structure(
+                    &expr_poly_id_by_order(c.expr.clone()),
+                    &expr_poly_id_by_order(c2.expr.clone()),
+                ) {
+                    // Found the same structure, try to extend the column mapping
+                    let new_mappings = create_mapping(&c.expr, &c2.expr);
+                    if extend_if_no_conflicts(&mut mappings, new_mappings) {
+                        // we're done for this constraint
+                        break;
                     }
                 }
             }
         }
 
         // do the same for bus interactions
-        'outer: for b in &pcp.machine.bus_interactions {
-            for pcp2 in self.pcps.iter() {
-                for b2 in &pcp2.machine.bus_interactions {
-                    if b.id == b2.id && b.args.len() == b2.args.len() {
-                        let all_args_same_structure =
-                            b.args.iter().zip_eq(b2.args.iter()).all(|(a1, a2)| {
-                                has_same_structure(
-                                    &expr_poly_id_by_order(a1.clone()),
-                                    &expr_poly_id_by_order(a2.clone()),
-                                )
-                            });
-                        if all_args_same_structure {
-                            for (arg1, arg2) in b.args.iter().zip_eq(b2.args.iter()) {
-                                let new_mappings = create_mapping(&arg1, &arg2);
-                                extend_if_no_conflicts(&mut mappings, new_mappings);
-                            }
-                            continue 'outer;
+        for b in &pcp.machine.bus_interactions {
+            for b2 in self.pcps.iter().map(|pcp| pcp.machine.bus_interactions.iter()).flatten() {
+                if b.id == b2.id && b.args.len() == b2.args.len() {
+                    let all_args_same_structure =
+                        b.args.iter().zip_eq(b2.args.iter()).all(|(a1, a2)| {
+                            has_same_structure(
+                                &expr_poly_id_by_order(a1.clone()),
+                                &expr_poly_id_by_order(a2.clone()),
+                            )
+                        });
+                    if all_args_same_structure {
+                        for (arg1, arg2) in b.args.iter().zip_eq(b2.args.iter()) {
+                            let new_mappings = create_mapping(&arg1, &arg2);
+                            extend_if_no_conflicts(&mut mappings, new_mappings);
                         }
+                        break;
                     }
                 }
             }
         }
 
-        // now we assign the columns to the PCP. When an id is not present in the mapping, use some other unused column
-        let mut curr_id = 0;
-        let mut new_poly_id = |old_id: u64| {
-            if let Some(id) = mappings.get_by_left(&old_id) {
-                return *id;
-            }
-            // find a new column not yet used in the mapping
-            while mappings.get_by_right(&curr_id).is_some() {
-                curr_id += 1;
-            }
-            assert!(
-                mappings.get_by_right(&curr_id).is_none(),
-                "New id {} already exists in mappings: {:?}",
-                curr_id,
-                mappings
-            );
-            assert!(!mappings.contains_left(&old_id));
-            let id = curr_id;
-            mappings.insert(old_id, id);
+        assign_columns_from_mapping(pcp, mappings);
+        self.pcps.push(pcp.clone());
+    }
+}
+
+/// Reasign precompile columns taking into acount the given mapping.
+/// When not present, use an unused column id (starting from 0).
+fn assign_columns_from_mapping<P: IntoOpenVm>(pcp: &mut PowdrPrecompile<P>, mut mapping: BiMap<u64, u64>) {
+    let mut curr_id = 0;
+    let mut new_poly_id = |old_id: u64| {
+        if let Some(id) = mapping.get_by_left(&old_id) {
+            return *id;
+        }
+        // find a new column not yet used in the mapping
+        while mapping.get_by_right(&curr_id).is_some() {
             curr_id += 1;
-            id
-        };
+        }
+        assert!(
+            mapping.get_by_right(&curr_id).is_none(),
+            "New id {} already exists in mapping: {:?}",
+            curr_id,
+            mapping
+        );
+        assert!(!mapping.contains_left(&old_id));
+        let id = curr_id;
+        mapping.insert(old_id, id);
+        curr_id += 1;
+        id
+    };
 
-        pcp.machine.pre_visit_expressions_mut(&mut |expr| {
-            if let AlgebraicExpression::Reference(r) = expr {
-                r.id = new_poly_id(r.id);
-            }
-        });
-        pcp.original_instructions.iter_mut().for_each(|instr| {
-            instr.subs.iter_mut().for_each(|sub| {
-                *sub = new_poly_id(*sub);
-            });
+    pcp.machine.pre_visit_expressions_mut(&mut |expr| {
+        if let AlgebraicExpression::Reference(r) = expr {
+            r.id = new_poly_id(r.id);
+        }
+    });
+    pcp.original_instructions.iter_mut().for_each(|instr| {
+        instr.subs.iter_mut().for_each(|sub| {
+            *sub = new_poly_id(*sub);
         });
+    });
 
-        pcp.is_valid_column.id = new_poly_id(pcp.is_valid_column.id);
-
-        self.pcps.push(pcp.clone());
-    }
+    pcp.is_valid_column.id = new_poly_id(pcp.is_valid_column.id);
 }

From c01c5cadace281c0e384f9c52f7d71e2590614c8 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 18 Jun 2025 18:19:48 -0300
Subject: [PATCH 33/39] test

---
 expression/src/lib.rs | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/expression/src/lib.rs b/expression/src/lib.rs
index f33e29502a..413f246143 100644
--- a/expression/src/lib.rs
+++ b/expression/src/lib.rs
@@ -249,16 +249,16 @@ fn expand<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>) -> Al
             match (left, right) {
                 // (a + b)*c ==> a*c + b*c
                 (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                    left: left1,
+                    left: a,
                     op: AlgebraicBinaryOperator::Add,
-                    right: right1,
-                }), right) => expand(*left1 * right.clone()) + expand(*right1 * right),
+                    right: b,
+                }), c) => expand(*a * c.clone()) + expand(*b * c),
                 // a*(b + c) ==> a*b + a*c
-                (left, AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                    left: left1,
+                (a, AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                    left: b,
                     op: AlgebraicBinaryOperator::Add,
-                    right: right1,
-                })) => expand(left.clone() * *left1) + expand(left.clone() * *right1),
+                    right: c,
+                })) => expand(a.clone() * *b) + expand(a.clone() * *c),
                 (left, right) => left * right,
             }
         }
@@ -326,5 +326,14 @@ mod tests {
             + v("b") * v("d") * v("e")
             + v("b") * v("d") * v("f");
         assert_eq!(normalize(e1), normalize(e2));
+
+        // a*(b + c) + d + b*(a + b)
+        let expr1 = normalize(expand(v("a") * (v("b") + v("c")) + v("d") + v("b") * (v("a") + v("b"))));
+        // a*b + a*c + (d + b*a + b*b)
+        let expr2 = normalize(expand(v("a") * v("b") + v("a") * v("c") + (v("d") + v("b") * v("a") + v("b") * v("b"))));
+        // (d + b*a) + b*b + (a*b + a*c)
+        let expr3 = normalize(expand((v("d") + v("b") * v("a")) + v("b") * v("b") + (v("a") * v("b") + v("a") * v("c"))));
+        assert_eq!(expr1, expr2);
+        assert_eq!(expr1, expr3);
     }
 }

From 54dd766286da46f45e472fd5cd57d71064c71909 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 18 Jun 2025 18:42:58 -0300
Subject: [PATCH 34/39] small change

---
 openvm/src/customize_exe/air_stacking.rs | 216 +++++++++++------------
 1 file changed, 106 insertions(+), 110 deletions(-)

diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index 5747a9657b..f84b7307b5 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -47,131 +47,126 @@ pub fn air_stacking<P: IntoOpenVm>(
         groups.entry(idx).or_default().push(pcp);
     }
 
-    let mut result = vec![];
+    groups.into_values().map(group_into_stacked).collect_vec()
+}
 
-    // transform each group into a stacked precompile
-    for mut extensions in groups.into_values() {
-        // group has a single precompile, no stacking
-        if extensions.len() == 1 {
-            let precompile = extensions.pop().unwrap();
-            tracing::debug!(
-                "Precompile {} not stacked",
-                &precompile.opcode.global_opcode().as_usize()
-            );
-            result.push(PowdrStackedPrecompile::new_single(precompile));
-            continue;
-        }
+/// Takes a group of precompiles and stacks them into a single stacked precompile.
+fn group_into_stacked<P: IntoOpenVm>(mut group: Vec<PowdrPrecompile<P>>) -> PowdrStackedPrecompile<P>{
+    // group has a single precompile, no stacking
+    if group.len() == 1 {
+        let precompile = group.pop().unwrap();
+        tracing::debug!(
+            "Precompile {} not stacked",
+            &precompile.opcode.global_opcode().as_usize()
+        );
+        return PowdrStackedPrecompile::new_single(precompile);
+    }
 
-        // handle largest precompile first
-        extensions.sort_by(|pcp1, pcp2| {
-            pcp2.machine
-                .unique_references()
-                .count()
-                .cmp(&pcp1.machine.unique_references().count())
-        });
+    // handle largest precompile first
+    group.sort_by(|pcp1, pcp2| {
+        pcp2.machine
+            .unique_references()
+            .count()
+            .cmp(&pcp1.machine.unique_references().count())
+    });
 
-        // assign columns to the precompiles in this group
-        let mut column_assigner = ColumnAssigner::default();
-        extensions.iter_mut().for_each(|ext| {
-            column_assigner.assign_precompile(ext);
-        });
+    // assign columns to the precompiles in this group
+    let mut column_assigner = ColumnAssigner::default();
+    group.iter_mut().for_each(|ext| {
+        column_assigner.assign_precompile(ext);
+    });
 
-        let mut stacked_constraints = vec![];
-        let mut interactions_by_machine = vec![];
-
-        tracing::debug!("Stacking {} precompiles", extensions.len());
-
-        // take the max id in all pcps and add 1.
-        let is_valid_start = 1 + extensions
-            .iter()
-            .flat_map(|pcp| {
-                pcp.original_instructions
-                    .iter()
-                    .flat_map(|instr| instr.subs.iter())
-            })
-            .max()
-            .unwrap();
-
-        let mut is_valid_sum: Option<AlgebraicExpression<P>> = None;
-
-        for (idx, pcp) in extensions.iter_mut().enumerate() {
-            // is_valid columns cannot be shared between precompiles. Here we do
-            // their remapping into exclusive columns.
-            let is_valid_new_id = is_valid_start + idx as u64;
-
-            // remap is_valid column in constraints and interactions
-            let mut remapped = pcp.machine.clone();
-            remapped.pre_visit_expressions_mut(&mut |expr| {
-                if let AlgebraicExpression::Reference(r) = expr {
-                    assert!(r.id <= is_valid_start);
-                    if r.id == pcp.is_valid_column.id {
-                        // we assume each pcp to have a specific column named "is_valid"
-                        assert!(*r.name == "is_valid");
-                        r.id = is_valid_new_id;
-                        r.name = format!("is_valid_{}", pcp.opcode.global_opcode().as_usize()).into();
-                    } else {
-                        r.name = format!("col_{}", r.id).into();
-                    }
+    tracing::debug!("Stacking {} precompiles", group.len());
+
+    // take the max id in all pcps and add 1.
+    let is_valid_start = 1 + group
+        .iter()
+        .flat_map(|pcp| {
+            pcp.original_instructions
+                .iter()
+                .flat_map(|instr| instr.subs.iter())
+        })
+        .max()
+        .unwrap();
+
+    let mut is_valid_sum: Option<AlgebraicExpression<P>> = None;
+    let mut stacked_constraints = vec![];
+    let mut interactions_by_machine = vec![];
+
+    for (idx, pcp) in group.iter_mut().enumerate() {
+        // is_valid columns cannot be shared between precompiles. Here we do
+        // their remapping into exclusive columns.
+        let is_valid_new_id = is_valid_start + idx as u64;
+
+        // remap is_valid column in constraints and interactions
+        let mut remapped = pcp.machine.clone();
+        remapped.pre_visit_expressions_mut(&mut |expr| {
+            if let AlgebraicExpression::Reference(r) = expr {
+                assert!(r.id <= is_valid_start);
+                if r.id == pcp.is_valid_column.id {
+                    // we assume each pcp to have a specific column named "is_valid"
+                    assert!(*r.name == "is_valid");
+                    r.id = is_valid_new_id;
+                    r.name = format!("is_valid_{}", pcp.opcode.global_opcode().as_usize()).into();
+                } else {
+                    r.name = format!("col_{}", r.id).into();
                 }
-            });
+            }
+        });
 
-            // set the is valid column in the original precompile
-            pcp.is_valid_column.id = is_valid_new_id;
+        // set the is valid column in the original precompile
+        pcp.is_valid_column.id = is_valid_new_id;
+
+        let is_valid = AlgebraicExpression::Reference(AlgebraicReference {
+            name: format!("is_valid_{}", pcp.opcode.global_opcode().as_usize()).into(),
+            id: is_valid_new_id,
+        });
 
-            let is_valid = AlgebraicExpression::Reference(AlgebraicReference {
-                name: format!("is_valid_{}", pcp.opcode.global_opcode().as_usize()).into(),
-                id: is_valid_new_id,
+        // guard interaction payloads so they can be merged later
+        remapped
+            .bus_interactions
+            .iter_mut()
+            .for_each(|interaction| {
+                interaction.args.iter_mut().for_each(|arg| { *arg = is_valid.clone() * arg.clone().normalize(); });
             });
 
-            // guard interaction payloads so they can be merged later
-            remapped
-                .bus_interactions
-                .iter_mut()
-                .for_each(|interaction| {
-                    interaction.args.iter_mut().for_each(|arg| { *arg = is_valid.clone() * arg.clone().normalize(); });
-                });
+        is_valid_sum = is_valid_sum
+            .map(|sum| sum + is_valid.clone())
+            .or_else(|| Some(is_valid.clone()));
 
-            is_valid_sum = is_valid_sum
-                .map(|sum| sum + is_valid.clone())
-                .or_else(|| Some(is_valid.clone()));
+        stacked_constraints.extend(remapped.constraints);
+        interactions_by_machine.push(remapped.bus_interactions);
+    }
 
-            stacked_constraints.extend(remapped.constraints);
-            interactions_by_machine.push(remapped.bus_interactions);
-        }
+    tracing::debug!("Stacked chip has {} constraints", stacked_constraints.len());
+    let mut stacked_constraints = join_constraints(stacked_constraints);
+    tracing::debug!("After joining constraints: {}", stacked_constraints.len());
 
-        tracing::debug!("Stacked chip has {} constraints", stacked_constraints.len());
-        let mut stacked_constraints = join_constraints(stacked_constraints);
-        tracing::debug!("After joining constraints: {}", stacked_constraints.len());
+    // enforce only one is_valid is active
+    stacked_constraints.push(make_bool(is_valid_sum.unwrap()).into());
 
-        // enforce only one is_valid is active
-        stacked_constraints.push(make_bool(is_valid_sum.unwrap()).into());
+    tracing::debug!(
+        "Stacked chip has {} bus interactions",
+        interactions_by_machine.iter().flatten().count()
+    );
+    let stacked_interactions = join_bus_interactions(interactions_by_machine);
+    tracing::debug!("After merging interactions: {}", stacked_interactions.len());
 
-        tracing::debug!(
-            "Stacked chip has {} bus interactions",
-            interactions_by_machine.iter().flatten().count()
-        );
-        let stacked_interactions = join_bus_interactions(interactions_by_machine);
-        tracing::debug!("After merging interactions: {}", stacked_interactions.len());
-
-        let machine = SymbolicMachine {
-            constraints: stacked_constraints,
-            bus_interactions: stacked_interactions,
-        };
-
-        result.push(PowdrStackedPrecompile {
-            precompiles: extensions
-                .into_iter()
-                .map(|p| (p.opcode.clone(), p))
-                .collect(),
-            machine,
-        });
-    }
+    let machine = SymbolicMachine {
+        constraints: stacked_constraints,
+        bus_interactions: stacked_interactions,
+    };
 
-    result
+    PowdrStackedPrecompile {
+        precompiles: group
+            .into_iter()
+            .map(|p| (p.opcode.clone(), p))
+            .collect(),
+        machine,
+    }
 }
 
 /// Merge bus interactions, taking into account args that have the same expression.
-/// When only compatible args are merged, the expression doesn't need an is_valid guard.
 fn join_bus_interactions<P: IntoOpenVm>(
     interactions_by_machine: Vec<Vec<SymbolicBusInteraction<P>>>,
 ) -> Vec<SymbolicBusInteraction<P>> {
@@ -276,14 +271,15 @@ fn join_bus_interactions<P: IntoOpenVm>(
         }
 
         // do the actual merging of the grouped interactions
-        result.extend(to_merge.into_iter().map(merge_bus_interactions));
+        result.extend(to_merge.into_iter().map(combine_bus_interactions));
     }
 
     result
 }
 
-/// helper to combine a set of bus interactions into a single interaction
-fn merge_bus_interactions<P: IntoOpenVm>(
+/// helper to combine a set of bus interactions into a single interaction.
+/// When only compatible args are combined, the expression doesn't need an is_valid guard.
+fn combine_bus_interactions<P: IntoOpenVm>(
     interactions: Vec<SymbolicBusInteraction<P>>,
 ) -> SymbolicBusInteraction<P> {
     assert!(interactions.iter().map(|i| i.id).unique().count() == 1, "grouped interactions should have the same bus");

From 3b8dd68f4a8647fb687d157b9abc52c301bf1e90 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 18 Jun 2025 19:01:10 -0300
Subject: [PATCH 35/39] fix

---
 openvm/src/customize_exe/air_stacking.rs | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index f84b7307b5..18d1f57958 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -30,10 +30,19 @@ pub fn air_stacking<P: IntoOpenVm>(
 ) -> Vec<PowdrStackedPrecompile<P>> {
     assert!(!extensions.is_empty());
 
+    // we normalize expressions in constraints and bus arguments so that they
+    // can later be compared for having the same structure
     extensions.iter_mut().for_each(|ext| {
         ext.machine.constraints.iter_mut().for_each(|c| {
             normalize_guarded(&mut c.expr);
-        })
+        });
+
+        ext.machine
+            .bus_interactions
+            .iter_mut()
+            .for_each(|i| {
+                i.args.iter_mut().for_each(|arg| { *arg = arg.clone().normalize(); });
+            });
     });
 
     // group precompiles by number of columns
@@ -42,8 +51,7 @@ pub fn air_stacking<P: IntoOpenVm>(
         let idx = f32::log(
             pcp.machine.unique_references().count() as f32,
             chip_stacking_log,
-        )
-        .floor() as usize;
+        ).floor() as usize;
         groups.entry(idx).or_default().push(pcp);
     }
 
@@ -127,7 +135,7 @@ fn group_into_stacked<P: IntoOpenVm>(mut group: Vec<PowdrPrecompile<P>>) -> Powd
             .bus_interactions
             .iter_mut()
             .for_each(|interaction| {
-                interaction.args.iter_mut().for_each(|arg| { *arg = is_valid.clone() * arg.clone().normalize(); });
+                interaction.args.iter_mut().for_each(|arg| { *arg = is_valid.clone() * arg.clone(); });
             });
 
         is_valid_sum = is_valid_sum

From 3ada3899a02dfbac90637fff2300b27f7e49fe7a Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Wed, 18 Jun 2025 19:06:53 -0300
Subject: [PATCH 36/39] fmt

---
 autoprecompiles/src/lib.rs               | 22 ++++----
 expression/src/lib.rs                    | 60 ++++++++++++++-------
 openvm/src/customize_exe.rs              |  2 +-
 openvm/src/customize_exe/air_stacking.rs | 69 +++++++++++++++---------
 4 files changed, 99 insertions(+), 54 deletions(-)

diff --git a/autoprecompiles/src/lib.rs b/autoprecompiles/src/lib.rs
index c56a69043a..aa282dc631 100644
--- a/autoprecompiles/src/lib.rs
+++ b/autoprecompiles/src/lib.rs
@@ -315,7 +315,7 @@ fn add_guards_constraint<T: FieldElement>(
 fn add_guards<T: FieldElement>(
     mut machine: SymbolicMachine<T>,
     bus_map: BusMap,
-    strict: bool
+    strict: bool,
 ) -> SymbolicMachine<T> {
     let pre_degree = machine.degree();
     let exec_bus_id = bus_map.get_bus_id(&BusType::ExecutionBridge).unwrap();
@@ -329,16 +329,20 @@ fn add_guards<T: FieldElement>(
     });
 
     machine.constraints = if strict {
-        machine.constraints.into_iter().map(|mut c| {
-            c.expr = is_valid.clone() * c.expr.clone();
-            c
-        }).collect()
+        machine
+            .constraints
+            .into_iter()
+            .map(|mut c| {
+                c.expr = is_valid.clone() * c.expr.clone();
+                c
+            })
+            .collect()
     } else {
         machine
-        .constraints
-        .into_iter()
-        .map(|c| add_guards_constraint(c.expr, &is_valid).into())
-        .collect()
+            .constraints
+            .into_iter()
+            .map(|c| add_guards_constraint(c.expr, &is_valid).into())
+            .collect()
     };
 
     let [execution_bus_receive, execution_bus_send] = machine
diff --git a/expression/src/lib.rs b/expression/src/lib.rs
index 413f246143..b4aefe451c 100644
--- a/expression/src/lib.rs
+++ b/expression/src/lib.rs
@@ -173,7 +173,10 @@ impl<T, R> From<T> for AlgebraicExpression<T, R> {
 }
 
 // Helper to flatten Add trees
-fn flatten_add<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>, acc: &mut Vec<AlgebraicExpression<T, R>>) {
+fn flatten_add<T: Ord + Clone, R: Ord + Clone>(
+    expr: AlgebraicExpression<T, R>,
+    acc: &mut Vec<AlgebraicExpression<T, R>>,
+) {
     match expr {
         AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
             left,
@@ -188,7 +191,10 @@ fn flatten_add<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>,
 }
 
 // Helper to flatten Mul trees
-fn flatten_mul<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>, acc: &mut Vec<AlgebraicExpression<T, R>>) {
+fn flatten_mul<T: Ord + Clone, R: Ord + Clone>(
+    expr: AlgebraicExpression<T, R>,
+    acc: &mut Vec<AlgebraicExpression<T, R>>,
+) {
     match expr {
         AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
             left,
@@ -202,7 +208,9 @@ fn flatten_mul<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>,
     }
 }
 
-fn normalize<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>) -> AlgebraicExpression<T, R> {
+fn normalize<T: Ord + Clone, R: Ord + Clone>(
+    expr: AlgebraicExpression<T, R>,
+) -> AlgebraicExpression<T, R> {
     match expr {
         AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
             left,
@@ -230,14 +238,14 @@ fn normalize<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>) ->
             left,
             op: AlgebraicBinaryOperator::Sub,
             right,
-        }) => {
-            normalize(*left) - normalize(*right)
-        }
+        }) => normalize(*left) - normalize(*right),
         _ => expr,
     }
 }
 
-fn expand<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>) -> AlgebraicExpression<T, R> {
+fn expand<T: Ord + Clone, R: Ord + Clone>(
+    expr: AlgebraicExpression<T, R>,
+) -> AlgebraicExpression<T, R> {
     match expr {
         AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
             left,
@@ -248,17 +256,23 @@ fn expand<T: Ord + Clone, R: Ord + Clone>(expr: AlgebraicExpression<T, R>) -> Al
             let right = expand(*right);
             match (left, right) {
                 // (a + b)*c ==> a*c + b*c
-                (AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                    left: a,
-                    op: AlgebraicBinaryOperator::Add,
-                    right: b,
-                }), c) => expand(*a * c.clone()) + expand(*b * c),
+                (
+                    AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                        left: a,
+                        op: AlgebraicBinaryOperator::Add,
+                        right: b,
+                    }),
+                    c,
+                ) => expand(*a * c.clone()) + expand(*b * c),
                 // a*(b + c) ==> a*b + a*c
-                (a, AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
-                    left: b,
-                    op: AlgebraicBinaryOperator::Add,
-                    right: c,
-                })) => expand(a.clone() * *b) + expand(a.clone() * *c),
+                (
+                    a,
+                    AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
+                        left: b,
+                        op: AlgebraicBinaryOperator::Add,
+                        right: c,
+                    }),
+                ) => expand(a.clone() * *b) + expand(a.clone() * *c),
                 (left, right) => left * right,
             }
         }
@@ -328,11 +342,17 @@ mod tests {
         assert_eq!(normalize(e1), normalize(e2));
 
         // a*(b + c) + d + b*(a + b)
-        let expr1 = normalize(expand(v("a") * (v("b") + v("c")) + v("d") + v("b") * (v("a") + v("b"))));
+        let expr1 = normalize(expand(
+            v("a") * (v("b") + v("c")) + v("d") + v("b") * (v("a") + v("b")),
+        ));
         // a*b + a*c + (d + b*a + b*b)
-        let expr2 = normalize(expand(v("a") * v("b") + v("a") * v("c") + (v("d") + v("b") * v("a") + v("b") * v("b"))));
+        let expr2 = normalize(expand(
+            v("a") * v("b") + v("a") * v("c") + (v("d") + v("b") * v("a") + v("b") * v("b")),
+        ));
         // (d + b*a) + b*b + (a*b + a*c)
-        let expr3 = normalize(expand((v("d") + v("b") * v("a")) + v("b") * v("b") + (v("a") * v("b") + v("a") * v("c"))));
+        let expr3 = normalize(expand(
+            (v("d") + v("b") * v("a")) + v("b") * v("b") + (v("a") * v("b") + v("a") * v("c")),
+        ));
         assert_eq!(expr1, expr2);
         assert_eq!(expr1, expr3);
     }
diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index c1ab802334..83e6119711 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -2,9 +2,9 @@ use std::collections::HashSet;
 use std::collections::{BTreeMap, BTreeSet, HashMap};
 use std::sync::Arc;
 
-use crate::powdr_extension::PowdrStackedPrecompile;
 use crate::extraction_utils::OriginalVmConfig;
 use crate::opcode::ALL_OPCODES;
+use crate::powdr_extension::PowdrStackedPrecompile;
 use crate::utils::UnsupportedOpenVmReferenceError;
 use crate::IntoOpenVm;
 use crate::OpenVmField;
diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index 18d1f57958..f5fbce9a0c 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -1,7 +1,9 @@
 use std::collections::{BTreeSet, HashMap};
 
 use powdr_autoprecompiles::{
-    expression::{AlgebraicExpression, AlgebraicReference}, powdr::make_bool, simplify_expression, SymbolicBusInteraction, SymbolicConstraint, SymbolicMachine
+    expression::{AlgebraicExpression, AlgebraicReference},
+    powdr::make_bool,
+    simplify_expression, SymbolicBusInteraction, SymbolicConstraint, SymbolicMachine,
 };
 
 use powdr_autoprecompiles::powdr::UniqueReferences;
@@ -37,12 +39,11 @@ pub fn air_stacking<P: IntoOpenVm>(
             normalize_guarded(&mut c.expr);
         });
 
-        ext.machine
-            .bus_interactions
-            .iter_mut()
-            .for_each(|i| {
-                i.args.iter_mut().for_each(|arg| { *arg = arg.clone().normalize(); });
+        ext.machine.bus_interactions.iter_mut().for_each(|i| {
+            i.args.iter_mut().for_each(|arg| {
+                *arg = arg.clone().normalize();
             });
+        });
     });
 
     // group precompiles by number of columns
@@ -51,7 +52,8 @@ pub fn air_stacking<P: IntoOpenVm>(
         let idx = f32::log(
             pcp.machine.unique_references().count() as f32,
             chip_stacking_log,
-        ).floor() as usize;
+        )
+        .floor() as usize;
         groups.entry(idx).or_default().push(pcp);
     }
 
@@ -59,7 +61,9 @@ pub fn air_stacking<P: IntoOpenVm>(
 }
 
 /// Takes a group of precompiles and stacks them into a single stacked precompile.
-fn group_into_stacked<P: IntoOpenVm>(mut group: Vec<PowdrPrecompile<P>>) -> PowdrStackedPrecompile<P>{
+fn group_into_stacked<P: IntoOpenVm>(
+    mut group: Vec<PowdrPrecompile<P>>,
+) -> PowdrStackedPrecompile<P> {
     // group has a single precompile, no stacking
     if group.len() == 1 {
         let precompile = group.pop().unwrap();
@@ -135,7 +139,9 @@ fn group_into_stacked<P: IntoOpenVm>(mut group: Vec<PowdrPrecompile<P>>) -> Powd
             .bus_interactions
             .iter_mut()
             .for_each(|interaction| {
-                interaction.args.iter_mut().for_each(|arg| { *arg = is_valid.clone() * arg.clone(); });
+                interaction.args.iter_mut().for_each(|arg| {
+                    *arg = is_valid.clone() * arg.clone();
+                });
             });
 
         is_valid_sum = is_valid_sum
@@ -166,10 +172,7 @@ fn group_into_stacked<P: IntoOpenVm>(mut group: Vec<PowdrPrecompile<P>>) -> Powd
     };
 
     PowdrStackedPrecompile {
-        precompiles: group
-            .into_iter()
-            .map(|p| (p.opcode.clone(), p))
-            .collect(),
+        precompiles: group.into_iter().map(|p| (p.opcode.clone(), p)).collect(),
         machine,
     }
 }
@@ -290,11 +293,15 @@ fn join_bus_interactions<P: IntoOpenVm>(
 fn combine_bus_interactions<P: IntoOpenVm>(
     interactions: Vec<SymbolicBusInteraction<P>>,
 ) -> SymbolicBusInteraction<P> {
-    assert!(interactions.iter().map(|i| i.id).unique().count() == 1, "grouped interactions should have the same bus");
+    assert!(
+        interactions.iter().map(|i| i.id).unique().count() == 1,
+        "grouped interactions should have the same bus"
+    );
     let id = interactions[0].id;
     // multiplicites are just added
     let mult = simplify_expression(
-        interactions.iter()
+        interactions
+            .iter()
             .map(|i| i.mult.clone())
             .reduce(|a, b| a + b)
             .unwrap(),
@@ -386,8 +393,11 @@ fn is_valid_guarded<P: IntoOpenVm>(expr: &AlgebraicExpression<P>) -> bool {
     }
 }
 
-fn normalize_guarded<P: IntoOpenVm>(expr:&mut AlgebraicExpression<P>) {
-    assert!(is_valid_guarded(expr), "not left guarded by is_valid col: {expr}");
+fn normalize_guarded<P: IntoOpenVm>(expr: &mut AlgebraicExpression<P>) {
+    assert!(
+        is_valid_guarded(expr),
+        "not left guarded by is_valid col: {expr}"
+    );
     match expr {
         AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation {
             left: _is_valid,
@@ -395,7 +405,7 @@ fn normalize_guarded<P: IntoOpenVm>(expr:&mut AlgebraicExpression<P>) {
             right,
         }) => {
             *right.as_mut() = (*right.clone()).normalize();
-        },
+        }
         _ => unreachable!(),
     }
 }
@@ -466,9 +476,7 @@ fn has_same_structure<P: IntoOpenVm>(
                 expr: expr2,
             }),
         ) => op1 == op2 && has_same_structure(expr1, expr2),
-        (AlgebraicExpression::Reference(r1), AlgebraicExpression::Reference(r2)) => {
-            r1.id == r2.id
-        }
+        (AlgebraicExpression::Reference(r1), AlgebraicExpression::Reference(r2)) => r1.id == r2.id,
         (AlgebraicExpression::Number(n1), AlgebraicExpression::Number(n2)) => n1 == n2,
         _ => false,
     }
@@ -604,7 +612,12 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
         // There may be multiple such constraints, so we try them all (as some
         // of the ids in the current constraint may already be assigned)
         for c in &pcp.machine.constraints {
-            for c2 in self.pcps.iter().map(|pcp| pcp.machine.constraints.iter()).flatten() {
+            for c2 in self
+                .pcps
+                .iter()
+                .map(|pcp| pcp.machine.constraints.iter())
+                .flatten()
+            {
                 if has_same_structure(
                     &expr_poly_id_by_order(c.expr.clone()),
                     &expr_poly_id_by_order(c2.expr.clone()),
@@ -621,7 +634,12 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
 
         // do the same for bus interactions
         for b in &pcp.machine.bus_interactions {
-            for b2 in self.pcps.iter().map(|pcp| pcp.machine.bus_interactions.iter()).flatten() {
+            for b2 in self
+                .pcps
+                .iter()
+                .map(|pcp| pcp.machine.bus_interactions.iter())
+                .flatten()
+            {
                 if b.id == b2.id && b.args.len() == b2.args.len() {
                     let all_args_same_structure =
                         b.args.iter().zip_eq(b2.args.iter()).all(|(a1, a2)| {
@@ -648,7 +666,10 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
 
 /// Reasign precompile columns taking into acount the given mapping.
 /// When not present, use an unused column id (starting from 0).
-fn assign_columns_from_mapping<P: IntoOpenVm>(pcp: &mut PowdrPrecompile<P>, mut mapping: BiMap<u64, u64>) {
+fn assign_columns_from_mapping<P: IntoOpenVm>(
+    pcp: &mut PowdrPrecompile<P>,
+    mut mapping: BiMap<u64, u64>,
+) {
     let mut curr_id = 0;
     let mut new_poly_id = |old_id: u64| {
         if let Some(id) = mapping.get_by_left(&old_id) {

From f176a0f29a7c5f691b8a5338e0b43968b3a09012 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Thu, 19 Jun 2025 09:35:32 -0300
Subject: [PATCH 37/39] clippy

---
 openvm/src/customize_exe.rs              |  2 +-
 openvm/src/customize_exe/air_stacking.rs | 22 +++++++++-------------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 83e6119711..2329ccb87b 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -151,7 +151,7 @@ pub fn customize(
     let n_skip = config.skip_autoprecompiles as usize;
     tracing::info!("Generating {n_acc} autoprecompiles in parallel");
 
-    let mut degree_bound = config.degree_bound.clone();
+    let mut degree_bound = config.degree_bound;
     // chip stacking needs to guard bus arguments, so we lower the inliner bound on bus interactions
     if config.chip_stacking_log.is_some() {
         degree_bound.bus_interactions -= 1;
diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index f5fbce9a0c..fb535adf8b 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -227,12 +227,12 @@ fn join_bus_interactions<P: IntoOpenVm>(
                     if used.contains(&idx) {
                         continue;
                     }
-                    let i2 = to_merge_set.get(0).unwrap();
+                    let i2 = to_merge_set.first().unwrap();
                     let all_args_same_structure = i
                         .args
                         .iter()
                         .zip_eq(i2.args.iter())
-                        .all(|(a1, a2)| has_same_structure(strip_guard(&a1), strip_guard(&a2)));
+                        .all(|(a1, a2)| has_same_structure(strip_guard(a1), strip_guard(a2)));
                     if all_args_same_structure {
                         // found an exact match
                         to_merge_set.push(i);
@@ -251,12 +251,12 @@ fn join_bus_interactions<P: IntoOpenVm>(
                     if used.contains(&idx) {
                         continue;
                     }
-                    let i2 = to_merge_set.get(0).unwrap();
+                    let i2 = to_merge_set.first().unwrap();
                     let some_args_same_structure = i
                         .args
                         .iter()
                         .zip_eq(i2.args.iter())
-                        .any(|(a1, a2)| has_same_structure(strip_guard(&a1), strip_guard(&a2)));
+                        .any(|(a1, a2)| has_same_structure(strip_guard(a1), strip_guard(a2)));
                     if some_args_same_structure {
                         // found a partial match on some args
                         to_merge_set.push(i);
@@ -539,7 +539,7 @@ fn create_mapping<P: IntoOpenVm>(
                 }),
             ) => {
                 assert_eq!(op1, op2);
-                create_mapping_inner(&expr1, &expr2)
+                create_mapping_inner(expr1, expr2)
             }
             (AlgebraicExpression::Reference(from), AlgebraicExpression::Reference(to)) => {
                 let mut mappings = BiMap::new();
@@ -615,8 +615,7 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
             for c2 in self
                 .pcps
                 .iter()
-                .map(|pcp| pcp.machine.constraints.iter())
-                .flatten()
+                .flat_map(|pcp| pcp.machine.constraints.iter())
             {
                 if has_same_structure(
                     &expr_poly_id_by_order(c.expr.clone()),
@@ -637,8 +636,7 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
             for b2 in self
                 .pcps
                 .iter()
-                .map(|pcp| pcp.machine.bus_interactions.iter())
-                .flatten()
+                .flat_map(|pcp| pcp.machine.bus_interactions.iter())
             {
                 if b.id == b2.id && b.args.len() == b2.args.len() {
                     let all_args_same_structure =
@@ -650,7 +648,7 @@ impl<P: IntoOpenVm> ColumnAssigner<P> {
                         });
                     if all_args_same_structure {
                         for (arg1, arg2) in b.args.iter().zip_eq(b2.args.iter()) {
-                            let new_mappings = create_mapping(&arg1, &arg2);
+                            let new_mappings = create_mapping(arg1, arg2);
                             extend_if_no_conflicts(&mut mappings, new_mappings);
                         }
                         break;
@@ -681,9 +679,7 @@ fn assign_columns_from_mapping<P: IntoOpenVm>(
         }
         assert!(
             mapping.get_by_right(&curr_id).is_none(),
-            "New id {} already exists in mapping: {:?}",
-            curr_id,
-            mapping
+            "New id {curr_id} already exists in mapping: {mapping:?}",
         );
         assert!(!mapping.contains_left(&old_id));
         let id = curr_id;

From 857997a6ffd59eb6bcf4c66d27efaa92c8edd011 Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Mon, 23 Jun 2025 10:29:42 -0300
Subject: [PATCH 38/39] comments

---
 autoprecompiles/src/lib.rs               | 14 ++++++++------
 openvm/src/customize_exe.rs              |  2 +-
 openvm/src/customize_exe/air_stacking.rs |  2 +-
 openvm/src/powdr_extension/vm.rs         |  8 ++++----
 4 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/autoprecompiles/src/lib.rs b/autoprecompiles/src/lib.rs
index aa282dc631..c722ee5454 100644
--- a/autoprecompiles/src/lib.rs
+++ b/autoprecompiles/src/lib.rs
@@ -65,11 +65,11 @@ impl<T: Clone + Ord + std::fmt::Display> Children<AlgebraicExpression<T>>
     }
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash, Ord, PartialOrd)]
+#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
 pub struct SymbolicBusInteraction<T> {
     pub id: u64,
-    pub args: Vec<AlgebraicExpression<T>>,
     pub mult: AlgebraicExpression<T>,
+    pub args: Vec<AlgebraicExpression<T>>,
 }
 
 impl<T: Display> Display for SymbolicBusInteraction<T> {
@@ -246,6 +246,8 @@ pub fn build<T: FieldElement, B: BusInteractionHandler<T> + IsBusStateful<T> + C
     vm_config: VmConfig<T, B>,
     degree_bound: DegreeBound,
     opcode: u32,
+    // If true, disables the optimization of multipliying only constants when guarding with is_valid.
+    // Chip stacking needs variables to be guarded too.
     strict_is_valid_guards: bool,
 ) -> Result<(SymbolicMachine<T>, Vec<Vec<u64>>), crate::constraint_optimizer::Error> {
     let (machine, subs) = statements_to_symbolic_machine(
@@ -315,7 +317,7 @@ fn add_guards_constraint<T: FieldElement>(
 fn add_guards<T: FieldElement>(
     mut machine: SymbolicMachine<T>,
     bus_map: BusMap,
-    strict: bool,
+    is_strict: bool,
 ) -> SymbolicMachine<T> {
     let pre_degree = machine.degree();
     let exec_bus_id = bus_map.get_bus_id(&BusType::ExecutionBridge).unwrap();
@@ -328,7 +330,7 @@ fn add_guards<T: FieldElement>(
         id: max_id,
     });
 
-    machine.constraints = if strict {
+    machine.constraints = if is_strict {
         machine
             .constraints
             .into_iter()
@@ -370,7 +372,7 @@ fn add_guards<T: FieldElement>(
     for b in &mut machine.bus_interactions {
         // already handled exec and pc lookup bus types
         if b.id != exec_bus_id && b.id != pc_lookup_bus_id {
-            if strict || !satisfies_zero_witness(&b.mult) {
+            if is_strict || !satisfies_zero_witness(&b.mult) {
                 // guard the multiplicity by `is_valid`
                 b.mult = is_valid.clone() * b.mult.clone();
                 // TODO this would not have to be cloned if we had *=
@@ -386,7 +388,7 @@ fn add_guards<T: FieldElement>(
 
     machine.constraints.extend(is_valid_mults);
 
-    if !strict {
+    if !is_strict {
         assert_eq!(
             pre_degree,
             machine.degree(),
diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 2329ccb87b..c8d44d13c0 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -263,7 +263,7 @@ pub fn customize(
         tracing::debug!("Chip stacking disabled");
         extensions
             .into_iter()
-            .map(PowdrStackedPrecompile::new_single)
+            .map(PowdrStackedPrecompile::from)
             .collect()
     };
 
diff --git a/openvm/src/customize_exe/air_stacking.rs b/openvm/src/customize_exe/air_stacking.rs
index fb535adf8b..0eb250a398 100644
--- a/openvm/src/customize_exe/air_stacking.rs
+++ b/openvm/src/customize_exe/air_stacking.rs
@@ -71,7 +71,7 @@ fn group_into_stacked<P: IntoOpenVm>(
             "Precompile {} not stacked",
             &precompile.opcode.global_opcode().as_usize()
         );
-        return PowdrStackedPrecompile::new_single(precompile);
+        return precompile.into();
     }
 
     // handle largest precompile first
diff --git a/openvm/src/powdr_extension/vm.rs b/openvm/src/powdr_extension/vm.rs
index 918749ebc4..a1d413e837 100644
--- a/openvm/src/powdr_extension/vm.rs
+++ b/openvm/src/powdr_extension/vm.rs
@@ -104,11 +104,11 @@ pub struct PowdrStackedPrecompile<P: IntoOpenVm> {
     pub machine: SymbolicMachine<P>,
 }
 
-impl<P: IntoOpenVm> PowdrStackedPrecompile<P> {
-    pub fn new_single(precompile: PowdrPrecompile<P>) -> Self {
+impl<P: IntoOpenVm> From<PowdrPrecompile<P>> for PowdrStackedPrecompile<P> {
+    fn from(pcp: PowdrPrecompile<P>) -> Self {
         Self {
-            machine: precompile.machine.clone(),
-            precompiles: BTreeMap::from([(precompile.opcode.clone(), precompile)]),
+            machine: pcp.machine.clone(),
+            precompiles: BTreeMap::from([(pcp.opcode.clone(), pcp)]),
         }
     }
 }

From be0c6b631011a32c5ed0fde6f000498d58a5457f Mon Sep 17 00:00:00 2001
From: Leandro Pacheco <contact@leandropacheco.com>
Date: Thu, 3 Jul 2025 10:44:38 -0300
Subject: [PATCH 39/39] fmt

---
 openvm/src/customize_exe.rs | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
index 1dcd2fc560..6528aa6ff4 100644
--- a/openvm/src/customize_exe.rs
+++ b/openvm/src/customize_exe.rs
@@ -2,9 +2,9 @@ use std::collections::{BTreeSet, HashMap};
 
 use std::sync::Arc;
 
-use crate::powdr_extension::PowdrStackedPrecompile;
 use crate::extraction_utils::{OriginalAirs, OriginalVmConfig};
 use crate::opcode::{branch_opcodes_bigint_set, branch_opcodes_set};
+use crate::powdr_extension::PowdrStackedPrecompile;
 use crate::utils::{fractional_knapsack, KnapsackItem, UnsupportedOpenVmReferenceError};
 use crate::IntoOpenVm;
 use crate::OpenVmField;
@@ -88,10 +88,17 @@ fn generate_apcs_with_pgo<P: IntoOpenVm>(
             bus_map,
             strict_is_valid_guards,
         ),
-        PgoConfig::Instruction(pgo_program_idx_count) => {
-            create_apcs_with_instruction_pgo(blocks, pgo_program_idx_count, airs, config, bus_map, strict_is_valid_guards)
+        PgoConfig::Instruction(pgo_program_idx_count) => create_apcs_with_instruction_pgo(
+            blocks,
+            pgo_program_idx_count,
+            airs,
+            config,
+            bus_map,
+            strict_is_valid_guards,
+        ),
+        PgoConfig::None => {
+            create_apcs_with_no_pgo(blocks, airs, config, bus_map, strict_is_valid_guards)
         }
-        PgoConfig::None => create_apcs_with_no_pgo(blocks, airs, config, bus_map, strict_is_valid_guards),
     };
 
     assert!(res.len() <= config.autoprecompiles as usize);
@@ -446,7 +453,13 @@ fn generate_autoprecompile<P: IntoOpenVm>(
         bus_map: bus_map.clone(),
     };
 
-    let apc = powdr_autoprecompiles::build(program, vm_config, degree_bound, apc_opcode as u32, strict_is_valid_guards)?;
+    let apc = powdr_autoprecompiles::build(
+        program,
+        vm_config,
+        degree_bound,
+        apc_opcode as u32,
+        strict_is_valid_guards,
+    )?;
 
     // Check that substitution values are unique over all instructions
     assert!(apc.subs().iter().flatten().all_unique());