Use insertion idx as air id

autoprecompiles/src/adapter.rs

@@ -87,7 +87,6 @@ where
     >;
     type CustomBusTypes: Clone + Display + Sync + Eq + PartialEq;
     type ApcStats: Send + Sync;
-    type AirId: Eq + Hash + Send + Sync;
 
     fn into_field(e: Self::PowdrField) -> Self::Field;
 

openvm/src/customize_exe.rs

@@ -63,7 +63,6 @@ impl<'a> Adapter for BabyBearOpenVmApcAdapter<'a> {
         OpenVmMemoryBusInteraction<Self::PowdrField, V>;
     type CustomBusTypes = OpenVmBusType;
     type ApcStats = OvmApcStats;
-    type AirId = String;
 
     fn into_field(e: Self::PowdrField) -> Self::Field {
         openvm_stark_sdk::p3_baby_bear::BabyBear::from_canonical_u32(

openvm/src/extraction_utils.rs

@@ -1,3 +1,4 @@
+use std::collections::btree_map::Entry;
 use std::collections::{BTreeMap, HashMap};
 use std::sync::{Arc, Mutex};
 
@@ -50,32 +51,35 @@ use crate::utils::symbolic_to_algebraic;
 // TODO: Use `<PackedChallenge<BabyBearSC> as FieldExtensionAlgebra<Val<BabyBearSC>>>::D` instead after fixing p3 dependency
 const EXT_DEGREE: usize = 4;
 
+#[derive(Clone, Serialize, Deserialize)]
+pub struct OriginalMachine<F> {
+    pub machine: SymbolicMachine<F>,
+    pub name: String,
+    pub metrics: AirMetrics,
+}
+
 #[derive(Clone, Serialize, Deserialize, Default)]
 pub struct OriginalAirs<F> {
-    pub(crate) opcode_to_air: HashMap<VmOpcode, String>,
-    pub(crate) air_name_to_machine: BTreeMap<String, (SymbolicMachine<F>, AirMetrics)>,
+    pub(crate) air_id_by_opcode: HashMap<VmOpcode, usize>,
+    pub(crate) machine_by_air_id: BTreeMap<usize, OriginalMachine<F>>,
 }
 
 impl<F> InstructionHandler for OriginalAirs<F> {
     type Field = F;
     type Instruction = Instr<F>;
-    type AirId = String;
+    type AirId = usize;
 
     fn get_instruction_air_and_id(
         &self,
         instruction: &Self::Instruction,
     ) -> (Self::AirId, &SymbolicMachine<Self::Field>) {
-        let id = self
-            .opcode_to_air
-            .get(&instruction.0.opcode)
-            .unwrap()
-            .clone();
-        let air = &self.air_name_to_machine.get(&id).unwrap().0;
+        let id = *self.air_id_by_opcode.get(&instruction.0.opcode).unwrap();
+        let air = &self.machine_by_air_id.get(&id).unwrap().machine;
         (id, air)
     }
 
     fn is_allowed(&self, instruction: &Self::Instruction) -> bool {
-        self.opcode_to_air.contains_key(&instruction.0.opcode)
+        self.air_id_by_opcode.contains_key(&instruction.0.opcode)
     }
 
     fn is_branching(&self, instruction: &Self::Instruction) -> bool {
@@ -95,59 +99,62 @@ impl<F> OriginalAirs<F> {
     pub fn insert_opcode(
         &mut self,
         opcode: VmOpcode,
-        air_name: String,
-        machine: impl Fn() -> Result<(SymbolicMachine<F>, AirMetrics), UnsupportedOpenVmReferenceError>,
+        air_id: usize,
+        machine: impl Fn() -> Result<OriginalMachine<F>, UnsupportedOpenVmReferenceError>,
     ) -> Result<(), UnsupportedOpenVmReferenceError> {
-        if self.opcode_to_air.contains_key(&opcode) {
+        if self.air_id_by_opcode.contains_key(&opcode) {
             panic!("Opcode {opcode} already exists");
         }
         // Insert the machine only if `air_name` isn't already present
-        if !self.air_name_to_machine.contains_key(&air_name) {
+        if let Entry::Vacant(e) = self.machine_by_air_id.entry(air_id) {
             let machine_instance = machine()?;
-            self.air_name_to_machine
-                .insert(air_name.clone(), machine_instance);
+            e.insert(machine_instance);
         }
 
-        self.opcode_to_air.insert(opcode, air_name);
+        self.air_id_by_opcode.insert(opcode, air_id);
+
         Ok(())
     }
 
     pub fn get_instruction_metrics(&self, opcode: VmOpcode) -> Option<&AirMetrics> {
-        self.opcode_to_air.get(&opcode).and_then(|air_name| {
-            self.air_name_to_machine
-                .get(air_name)
-                .map(|(_, metrics)| metrics)
+        self.air_id_by_opcode.get(&opcode).and_then(|air_id| {
+            self.machine_by_air_id
+                .get(air_id)
+                .map(|machine| &machine.metrics)
         })
     }
 
     pub fn allow_list(&self) -> Vec<VmOpcode> {
-        self.opcode_to_air.keys().cloned().collect()
+        self.air_id_by_opcode.keys().cloned().collect()
     }
 
     pub fn airs_by_name(&self) -> impl Iterator<Item = (&String, &SymbolicMachine<F>)> {
-        self.air_name_to_machine
-            .iter()
-            .map(|(name, (machine, _))| (name, machine))
+        self.machine_by_air_id
+            .values()
+            .map(|machine| (&machine.name, &machine.machine))
     }
 }
 
 /// For each air name, the dimension of a record arena needed to store the
 /// records for a single APC call.
-pub fn record_arena_dimension_by_air_name_per_apc_call<F>(
+pub fn record_arena_dimension_by_air_id_per_apc_call<F>(
     apc: &Apc<F, Instr<F>>,
     air_by_opcode_id: &OriginalAirs<F>,
-) -> BTreeMap<String, RecordArenaDimension> {
+) -> BTreeMap<usize, RecordArenaDimension> {
     apc.instructions().iter().map(|instr| &instr.0.opcode).fold(
         BTreeMap::new(),
         |mut acc, opcode| {
             // Get the air name for this opcode
-            let air_name = air_by_opcode_id.opcode_to_air.get(opcode).unwrap();
+            let air_id = air_by_opcode_id.air_id_by_opcode.get(opcode).unwrap();
 
             // Increment the height for this air name, initializing if necessary
-            acc.entry(air_name.clone())
+            acc.entry(*air_id)
                 .or_insert_with(|| {
-                    let (_, air_metrics) =
-                        air_by_opcode_id.air_name_to_machine.get(air_name).unwrap();
+                    let air_metrics = &air_by_opcode_id
+                        .machine_by_air_id
+                        .get(air_id)
+                        .unwrap()
+                        .metrics;
 
                     RecordArenaDimension {
                         height: 0,
@@ -300,11 +307,12 @@ impl OriginalVmConfig {
             })
             .map(|(op, executor_id)| {
                 let insertion_index = chip_inventory.executor_idx_to_insertion_idx[executor_id];
-                let air_ref = &chip_inventory.airs().ext_airs()[insertion_index];
-                (op, air_ref)
+                (op, insertion_index)
             }) // find executor for opcode
-            .try_fold(OriginalAirs::default(), |mut airs, (op, air_ref)| {
-                airs.insert_opcode(op, air_ref.name(), || {
+            .try_fold(OriginalAirs::default(), |mut airs, (op, insertion_idx)| {
+                airs.insert_opcode(op, insertion_idx, || {
+                    let air_ref = &chip_inventory.airs().ext_airs()[insertion_idx];
+                    let name = air_ref.name();
                     let columns = get_columns(air_ref.clone());
                     let constraints = get_constraints(air_ref.clone());
                     let metrics = get_air_metrics(air_ref.clone(), max_degree);
@@ -321,14 +329,17 @@ impl OriginalVmConfig {
                         .map(|expr| openvm_bus_interaction_to_powdr(expr, &columns))
                         .collect::<Result<_, _>>()?;
 
-                    Ok((
-                        SymbolicMachine {
-                            constraints: powdr_exprs.into_iter().map(Into::into).collect(),
-                            bus_interactions: powdr_bus_interactions,
-                            derived_columns: vec![],
-                        },
+                    let machine = SymbolicMachine {
+                        constraints: powdr_exprs.into_iter().map(Into::into).collect(),
+                        bus_interactions: powdr_bus_interactions,
+                        derived_columns: vec![],
+                    };
+
+                    Ok(OriginalMachine {
+                        machine,
+                        name,
                         metrics,
-                    ))
+                    })
                 })?;
 
                 Ok(airs)

openvm/src/powdr_extension/executor/mod.rs

@@ -8,7 +8,7 @@ use std::{
 
 use crate::{
     extraction_utils::{
-        record_arena_dimension_by_air_name_per_apc_call, OriginalAirs, OriginalVmConfig,
+        record_arena_dimension_by_air_id_per_apc_call, OriginalAirs, OriginalVmConfig,
     },
     Instr,
 };
@@ -94,10 +94,10 @@ impl<A: Arena> OriginalArenas<A> {
 
     /// Returns the arena of the given air name.
     /// - Panics if the arenas are not initialized.
-    pub fn take_arena(&mut self, air_name: &str) -> Option<A> {
+    pub fn take_arena(&mut self, air_id: usize) -> Option<A> {
         match self {
             OriginalArenas::Uninitialized => panic!("original arenas are uninitialized"),
-            OriginalArenas::Initialized(initialized) => initialized.take_arena(air_name),
+            OriginalArenas::Initialized(initialized) => initialized.take_arena(air_id),
         }
     }
 
@@ -125,7 +125,7 @@ impl<A: Arena> OriginalArenas<A> {
 #[derive(Default)]
 pub struct InitializedOriginalArenas<A> {
     arenas: Vec<Option<A>>,
-    air_name_to_arena_index: HashMap<String, usize>,
+    air_id_to_arena_index: HashMap<usize, usize>,
     pub number_of_calls: usize,
 }
 
@@ -137,33 +137,32 @@ impl<A: Arena> InitializedOriginalArenas<A> {
         apc: &Arc<Apc<BabyBear, Instr<BabyBear>>>,
     ) -> Self {
         let record_arena_dimensions =
-            record_arena_dimension_by_air_name_per_apc_call(apc, original_airs);
-        let (air_name_to_arena_index, arenas) =
-            record_arena_dimensions.into_iter().enumerate().fold(
-                (HashMap::new(), Vec::new()),
-                |(mut air_name_to_arena_index, mut arenas),
-                 (
-                    idx,
-                    (
-                        air_name,
-                        RecordArenaDimension {
-                            height: num_calls,
-                            width: air_width,
-                        },
-                    ),
-                )| {
-                    air_name_to_arena_index.insert(air_name, idx);
-                    arenas.push(Some(A::with_capacity(
-                        num_calls * apc_call_count_estimate,
-                        air_width,
-                    )));
-                    (air_name_to_arena_index, arenas)
-                },
-            );
+            record_arena_dimension_by_air_id_per_apc_call(apc, original_airs);
+        let (air_id_to_arena_index, arenas) = record_arena_dimensions.into_iter().enumerate().fold(
+            (HashMap::new(), Vec::new()),
+            |(mut air_name_to_arena_index, mut arenas),
+             (
+                idx,
+                (
+                    air_id,
+                    RecordArenaDimension {
+                        height: num_calls,
+                        width: air_width,
+                    },
+                ),
+            )| {
+                air_name_to_arena_index.insert(air_id, idx);
+                arenas.push(Some(A::with_capacity(
+                    num_calls * apc_call_count_estimate,
+                    air_width,
+                )));
+                (air_name_to_arena_index, arenas)
+            },
+        );
 
         Self {
             arenas,
-            air_name_to_arena_index,
+            air_id_to_arena_index,
             // This is the actual number of calls, which we don't know yet. It will be updated during preflight execution.
             number_of_calls: 0,
         }
@@ -177,8 +176,8 @@ impl<A: Arena> InitializedOriginalArenas<A> {
             .expect("arena missing for index")
     }
 
-    fn take_arena(&mut self, air_name: &str) -> Option<A> {
-        let index = *self.air_name_to_arena_index.get(air_name)?;
+    fn take_arena(&mut self, air_id: usize) -> Option<A> {
+        let index = *self.air_id_to_arena_index.get(&air_id)?;
         self.arenas[index].take()
     }
 }
@@ -574,10 +573,10 @@ impl PowdrExecutor {
         let executor_inventory = base_config.sdk_config.sdk.create_executors().unwrap();
 
         let arena_index_by_name =
-            record_arena_dimension_by_air_name_per_apc_call(apc.as_ref(), &air_by_opcode_id)
+            record_arena_dimension_by_air_id_per_apc_call(apc.as_ref(), &air_by_opcode_id)
                 .iter()
                 .enumerate()
-                .map(|(idx, (name, _))| (name.clone(), idx))
+                .map(|(idx, (name, _))| (*name, idx))
                 .collect::<HashMap<_, _>>();
 
         let cached_instructions_meta = apc
@@ -590,7 +589,7 @@ impl PowdrExecutor {
                     .expect("missing executor for opcode")
                     as usize;
                 let air_name = air_by_opcode_id
-                    .opcode_to_air
+                    .air_id_by_opcode
                     .get(&instruction.0.opcode)
                     .expect("missing air for opcode");
                 let arena_index = *arena_index_by_name

openvm/src/powdr_extension/trace_generator/cpu/mod.rs

@@ -119,24 +119,22 @@ impl PowdrTraceGeneratorCpu {
             .inventory
         };
 
-        let dummy_trace_by_air_name: HashMap<String, SharedCpuTrace<BabyBear>> = chip_inventory
+        let dummy_trace_by_air_id: HashMap<usize, SharedCpuTrace<BabyBear>> = chip_inventory
             .chips()
             .iter()
             .enumerate()
             .rev()
-            .filter_map(|(insertion_idx, chip)| {
-                let air_name = chip_inventory.airs().ext_airs()[insertion_idx].name();
-
+            .filter_map(|(air_id, chip)| {
                 let record_arena = {
-                    match original_arenas.take_arena(&air_name) {
+                    match original_arenas.take_arena(air_id) {
                         Some(ra) => ra,
                         None => return None, // skip this iteration, because we only have record arena for chips that are used
                     }
                 };
 
                 let shared_trace = chip.generate_proving_ctx(record_arena).common_main.unwrap();
 
-                Some((air_name, SharedCpuTrace::from(shared_trace)))
+                Some((air_id, SharedCpuTrace::from(shared_trace)))
             })
             .collect();
 
@@ -146,7 +144,7 @@ impl PowdrTraceGeneratorCpu {
             apc_poly_id_to_index,
             columns_to_compute,
         } = generate_trace(
-            &dummy_trace_by_air_name,
+            &dummy_trace_by_air_id,
             &self.original_airs,
             num_apc_calls,
             &self.apc,