refactor: replace direct indexing for bus_prefix and beta_powers

air/src/ace.rs

@@ -49,7 +49,8 @@ pub enum MessageElement {
 pub enum ProductFactor {
     /// Claimed final value of an auxiliary trace column, by column index.
     BusBoundary(usize),
-    /// A bus message computed from its elements as `bus_prefix[bus] + sum(beta^i * elements[i])`.
+    /// A bus message computed from its elements as the bus-specific prefix plus
+    /// `sum(beta^i * elements[i])` (same convention as [`trace::Challenges::encode`]).
     /// The first field is the bus type index (see `trace::bus_types`).
     Message(usize, Vec<MessageElement>),
     /// Multiset product reduced from variable-length public inputs, by group index.
@@ -157,9 +158,9 @@ where
     sum
 }
 
-/// Encode a bus message as `bus_prefix[bus] + sum(beta^i * elements[i])`.
+/// Encode a bus message as the bus-specific prefix plus `sum(beta^i * elements[i])`.
 ///
-/// The bus prefix provides domain separation: `bus_prefix[bus] = alpha + (bus+1) * gamma`
+/// The bus prefix provides domain separation: `alpha + (bus+1) * gamma` for each bus index.
 /// where `gamma = beta^MAX_MESSAGE_WIDTH`. This matches [`trace::Challenges::encode`].
 fn encode_bus_message<EF>(
     builder: &mut DagBuilder<EF>,

air/src/constraints/range/bus.rs

@@ -18,7 +18,7 @@ use miden_crypto::stark::air::{ExtensionBuilder, WindowAccess};
 use crate::{
     MainCols, MidenAirBuilder,
     constraints::{chiplets::selectors::ChipletSelectors, utils::BoolNot},
-    trace::{Challenges, bus_types::RANGE_CHECK_BUS, range},
+    trace::{Challenges, range},
 };
 
 // ENTRY POINTS
@@ -56,7 +56,7 @@ pub fn enforce_bus<AB>(
     let b_local = aux_local[range::B_RANGE_COL_IDX];
     let b_next = aux_next[range::B_RANGE_COL_IDX];
 
-    let alpha = &challenges.bus_prefix[RANGE_CHECK_BUS];
+    let alpha = &challenges.bus_prefix.range_check_bus;
 
     // Denominators for LogUp
     let mem = local.memory();

air/src/trace/challenges.rs

@@ -1,7 +1,7 @@
 //! Unified bus challenge encoding with per-bus domain separation.
 //!
 //! Provides [`Challenges`], a single struct for encoding multiset/LogUp bus messages
-//! as `bus_prefix[bus] + <beta, message>`. Each bus interaction type gets a unique
+//! as `bus_prefix.prefix_for_bus(bus) + <beta, message>`. Each bus interaction type gets a unique
 //! prefix to ensure domain separation.
 //!
 //! This type is used by:
@@ -17,14 +17,107 @@ use core::ops::{AddAssign, Mul};
 
 use miden_core::field::PrimeCharacteristicRing;
 
-use super::{MAX_MESSAGE_WIDTH, bus_types::NUM_BUS_TYPES};
+use super::{
+    MAX_MESSAGE_WIDTH,
+    bus_message::{
+        ADDR_IDX, CAPACITY_COEFF_END_IDX, CAPACITY_START_IDX, LABEL_IDX, NODE_INDEX_IDX,
+        STATE_COEFF_END_IDX, STATE_START_IDX,
+    },
+    bus_types::{
+        ACE_WIRING_BUS, BLOCK_HASH_TABLE, BLOCK_STACK_TABLE, CHIPLETS_BUS,
+        LOG_PRECOMPILE_TRANSCRIPT, NUM_BUS_TYPES, OP_GROUP_TABLE, RANGE_CHECK_BUS, SIBLING_TABLE,
+        STACK_OVERFLOW_TABLE,
+    },
+};
 
-/// Encodes multiset/LogUp contributions as **bus_prefix\[bus\] + \<beta, message\>**.
+/// Precomputed powers `beta^0 .. beta^(MAX_MESSAGE_WIDTH-1)` for bus message dot products.
+///
+/// Layout matches [`super::bus_message`].
+#[derive(Clone, Debug)]
+pub struct BetaPowers<EF> {
+    pub label: EF,
+    pub addr: EF,
+    pub node_index: EF,
+    pub state: [EF; 8],
+    pub capacity: [EF; 4],
+    pub beta_pow_15: EF,
+}
+
+impl<EF: Clone> BetaPowers<EF> {
+    /// Coefficient `beta^i` for `i < MAX_MESSAGE_WIDTH` (see [`super::bus_message`] indices).
+    #[inline(always)]
+    pub fn at(&self, idx: usize) -> &EF {
+        match idx {
+            LABEL_IDX => &self.label,
+            ADDR_IDX => &self.addr,
+            NODE_INDEX_IDX => &self.node_index,
+            STATE_START_IDX..=STATE_COEFF_END_IDX => &self.state[idx - STATE_START_IDX],
+            CAPACITY_START_IDX..=CAPACITY_COEFF_END_IDX => &self.capacity[idx - CAPACITY_START_IDX],
+            x if x == MAX_MESSAGE_WIDTH - 1 => &self.beta_pow_15,
+            _ => unreachable!("beta power index {idx} out of range (< {MAX_MESSAGE_WIDTH})"),
+        }
+    }
+
+    /// `beta^CAPACITY_DOMAIN_IDX` — second capacity slot (e.g. control-block `op_code`).
+    #[inline(always)]
+    pub fn capacity_domain(&self) -> &EF {
+        &self.capacity[super::bus_message::CAPACITY_DOMAIN_IDX - CAPACITY_START_IDX]
+    }
+
+    /// Dense `beta^3 ..= beta^14`: Hasher rate coefficients then capacity (12 elements).
+    #[inline(always)]
+    pub fn sponge_state(&self) -> [EF; 12] {
+        core::array::from_fn(|i| {
+            if i < 8 {
+                self.state[i].clone()
+            } else {
+                self.capacity[i - 8].clone()
+            }
+        })
+    }
+}
+
+/// Per-bus domain separation constants: each field is `alpha + (bus_const + 1) * gamma`
+/// where `gamma = beta^MAX_MESSAGE_WIDTH` (see [`Challenges::new`]).
+///
+/// Field order matches [`super::bus_types`].
+#[derive(Clone, Debug)]
+pub struct BusPrefix<EF> {
+    pub chiplets_bus: EF,
+    pub block_stack_table: EF,
+    pub block_hash_table: EF,
+    pub op_group_table: EF,
+    pub stack_overflow_table: EF,
+    pub sibling_table: EF,
+    pub log_precompile_transcript: EF,
+    pub range_check_bus: EF,
+    pub ace_wiring_bus: EF,
+}
+
+impl<EF> BusPrefix<EF> {
+    /// Prefix for the bus index used by [`Challenges::encode`] / [`Challenges::encode_sparse`].
+    #[inline(always)]
+    pub(crate) fn prefix_for_bus(&self, bus: usize) -> &EF {
+        match bus {
+            CHIPLETS_BUS => &self.chiplets_bus,
+            BLOCK_STACK_TABLE => &self.block_stack_table,
+            BLOCK_HASH_TABLE => &self.block_hash_table,
+            OP_GROUP_TABLE => &self.op_group_table,
+            STACK_OVERFLOW_TABLE => &self.stack_overflow_table,
+            SIBLING_TABLE => &self.sibling_table,
+            LOG_PRECOMPILE_TRANSCRIPT => &self.log_precompile_transcript,
+            RANGE_CHECK_BUS => &self.range_check_bus,
+            ACE_WIRING_BUS => &self.ace_wiring_bus,
+            _ => unreachable!("bus index {bus} is not a valid bus type (< {NUM_BUS_TYPES})"),
+        }
+    }
+}
+
+/// Encodes multiset/LogUp contributions as **per-bus prefix + \<beta, message\>**.
 ///
 /// - `alpha`: randomness base (kept public for direct access by range checker etc.)
-/// - `beta_powers`: precomputed powers `[beta^0, beta^1, ..., beta^(MAX_MESSAGE_WIDTH-1)]`
-/// - `bus_prefix`: per-bus domain separation constants `bus_prefix[i] = alpha + (i+1) *
-///   beta^MAX_MESSAGE_WIDTH`
+/// - `beta_powers`: precomputed powers `beta^0 .. beta^(MAX_MESSAGE_WIDTH-1)` (see [`BetaPowers`])
+/// - `bus_prefix`: per-bus domain separation (see [`BusPrefix`])
 ///
 /// The challenges are derived from permutation randomness:
 /// - `alpha = challenges[0]`
@@ -33,27 +126,43 @@ use super::{MAX_MESSAGE_WIDTH, bus_types::NUM_BUS_TYPES};
 /// Precomputed once and passed by reference to all bus components.
 pub struct Challenges<EF: PrimeCharacteristicRing> {
     pub alpha: EF,
-    pub beta_powers: [EF; MAX_MESSAGE_WIDTH],
-    /// Per-bus domain separation: `bus_prefix[i] = alpha + (i+1) * gamma`
-    /// where `gamma = beta^MAX_MESSAGE_WIDTH`.
-    pub bus_prefix: [EF; NUM_BUS_TYPES],
+    pub beta_powers: BetaPowers<EF>,
+    pub bus_prefix: BusPrefix<EF>,
 }
 
 impl<EF: PrimeCharacteristicRing> Challenges<EF> {
     /// Builds `alpha`, precomputed `beta` powers, and per-bus prefixes.
     pub fn new(alpha: EF, beta: EF) -> Self {
-        let mut beta_powers = core::array::from_fn(|_| EF::ONE);
+        let mut arr: [EF; MAX_MESSAGE_WIDTH] = core::array::from_fn(|_| EF::ONE);
         for i in 1..MAX_MESSAGE_WIDTH {
-            beta_powers[i] = beta_powers[i - 1].clone() * beta.clone();
+            arr[i] = arr[i - 1].clone() * beta.clone();
         }
         // gamma = beta^MAX_MESSAGE_WIDTH (one power beyond the message range)
-        let gamma = beta_powers[MAX_MESSAGE_WIDTH - 1].clone() * beta;
-        let bus_prefix =
-            core::array::from_fn(|i| alpha.clone() + gamma.clone() * EF::from_u32((i as u32) + 1));
+        let gamma = arr[MAX_MESSAGE_WIDTH - 1].clone() * beta;
+        let beta_powers = BetaPowers {
+            label: arr[LABEL_IDX].clone(),
+            addr: arr[ADDR_IDX].clone(),
+            node_index: arr[NODE_INDEX_IDX].clone(),
+            state: core::array::from_fn(|i| arr[STATE_START_IDX + i].clone()),
+            capacity: core::array::from_fn(|i| arr[CAPACITY_START_IDX + i].clone()),
+            beta_pow_15: arr[MAX_MESSAGE_WIDTH - 1].clone(),
+        };
+        let term = |k: u32| alpha.clone() + gamma.clone() * EF::from_u32(k);
+        let bus_prefix = BusPrefix {
+            chiplets_bus: term(1),
+            block_stack_table: term(2),
+            block_hash_table: term(3),
+            op_group_table: term(4),
+            stack_overflow_table: term(5),
+            sibling_table: term(6),
+            log_precompile_transcript: term(7),
+            range_check_bus: term(8),
+            ace_wiring_bus: term(9),
+        };
         Self { alpha, beta_powers, bus_prefix }
     }
 
-    /// Encodes as **bus_prefix\[bus\] + sum(beta_powers\[i\] * elem\[i\])** with K consecutive
+    /// Encodes as **prefix_for_bus(bus) + sum(beta_powers.at(i) * elem\[i\])** with K consecutive
     /// elements.
     ///
     /// The `bus` parameter selects the bus interaction type for domain separation.
@@ -67,14 +176,14 @@ impl<EF: PrimeCharacteristicRing> Challenges<EF> {
             bus < NUM_BUS_TYPES,
             "Bus index {bus} exceeds NUM_BUS_TYPES ({NUM_BUS_TYPES})"
         );
-        let mut acc = self.bus_prefix[bus].clone();
+        let mut acc = self.bus_prefix.prefix_for_bus(bus).clone();
         for (i, elem) in elems.into_iter().enumerate() {
-            acc += self.beta_powers[i].clone() * elem;
+            acc += self.beta_powers.at(i).clone() * elem;
         }
         acc
     }
 
-    /// Encodes as **bus_prefix\[bus\] + sum(beta_powers\[layout\[i\]\] * values\[i\])** using
+    /// Encodes as **prefix_for_bus(bus) + sum(beta_powers.at(layout\[i\]) * values\[i\])** using
     /// sparse positions.
     ///
     /// The `bus` parameter selects the bus interaction type for domain separation.
@@ -92,15 +201,13 @@ impl<EF: PrimeCharacteristicRing> Challenges<EF> {
             bus < NUM_BUS_TYPES,
             "Bus index {bus} exceeds NUM_BUS_TYPES ({NUM_BUS_TYPES})"
         );
-        let mut acc = self.bus_prefix[bus].clone();
+        let mut acc = self.bus_prefix.prefix_for_bus(bus).clone();
         for (idx, value) in layout.into_iter().zip(values) {
             debug_assert!(
-                idx < self.beta_powers.len(),
-                "encode_sparse index {} exceeds beta_powers length ({})",
-                idx,
-                self.beta_powers.len()
+                idx < MAX_MESSAGE_WIDTH,
+                "encode_sparse index {idx} exceeds beta_powers range ({MAX_MESSAGE_WIDTH})"
             );
-            acc += self.beta_powers[idx].clone() * value;
+            acc += self.beta_powers.at(idx).clone() * value;
         }
         acc
     }

air/src/trace/mod.rs

@@ -4,7 +4,7 @@ use chiplets::hasher::RATE_LEN;
 use miden_core::utils::range;
 
 mod challenges;
-pub use challenges::Challenges;
+pub use challenges::{BetaPowers, BusPrefix, Challenges};
 
 pub mod chiplets;
 pub mod decoder;
@@ -188,45 +188,53 @@ pub const MAX_MESSAGE_WIDTH: usize = 16;
 /// Bus message coefficient indices.
 ///
 /// These define the standard positions for encoding bus messages using the pattern:
-/// `alpha + sum(beta_powers\[i\] * elem\[i\])` where:
+/// `alpha + sum(beta_powers.at(i) * elem\[i\])` where:
 /// - `alpha` is the randomness base (accessed directly as `.alpha`)
-/// - `beta_powers\[i\] = beta^i` are the powers of beta
+/// - `beta_powers.at(i) = beta^i` are the powers of beta (see [`BetaPowers::at`])
 ///
-/// These indices refer to positions in the `beta_powers` array, not including alpha.
+/// These indices refer to positions in [`Challenges::beta_powers`] ([`BetaPowers`]), not including
+/// alpha.
 ///
 /// This layout is shared between:
 /// - AIR constraint builders (symbolic expressions): `Challenges<AB::ExprEF>`
 /// - Processor auxiliary trace builders (concrete field elements): `Challenges<E>`
 pub mod bus_message {
-    /// Label coefficient index: `beta_powers[0] = beta^0`.
+    /// Label coefficient index: `beta_powers.label` (`beta^0`).
     ///
     /// Used for transition type/operation label.
     pub const LABEL_IDX: usize = 0;
 
-    /// Address coefficient index: `beta_powers[1] = beta^1`.
+    /// Address coefficient index: `beta_powers.addr` (`beta^1`).
     ///
     /// Used for chiplet address.
     pub const ADDR_IDX: usize = 1;
 
-    /// Node index coefficient index: `beta_powers[2] = beta^2`.
+    /// Node index coefficient index: `beta_powers.node_index` (`beta^2`).
     ///
     /// Used for Merkle path position. Set to 0 for non-Merkle operations (SPAN, RESPAN, HPERM,
     /// etc.).
     pub const NODE_INDEX_IDX: usize = 2;
 
-    /// State start coefficient index: `beta_powers[3] = beta^3`.
+    /// State start coefficient index: `beta_powers.state[0] = beta^3`.
     ///
     /// Beginning of hasher state. Hasher state occupies 8 consecutive coefficients:
-    /// `beta_powers[3..11]` (beta^3..beta^10) for `state[0..7]` (rate portion: RATE0 || RATE1).
+    /// `beta_powers.state[0..8]` (`beta^3..beta^10`) for `state[0..7]` (rate portion: RATE0 ||
+    /// RATE1).
     pub const STATE_START_IDX: usize = 3;
 
-    /// Capacity start coefficient index: `beta_powers[11] = beta^11`.
+    /// Inclusive end index for the eight `state` coefficients (`beta^3..=beta^10`).
+    pub const STATE_COEFF_END_IDX: usize = STATE_START_IDX + 7;
+
+    /// Capacity start coefficient index: `beta_powers.capacity[0] = beta^11`.
     ///
     /// Beginning of hasher capacity. Hasher capacity occupies 4 consecutive coefficients:
-    /// `beta_powers[11..15]` (beta^11..beta^14) for `capacity[0..3]`.
+    /// `beta_powers.capacity[0..4]` (`beta^11..beta^14`) for `capacity[0..3]`.
     pub const CAPACITY_START_IDX: usize = 11;
 
-    /// Capacity domain coefficient index: `beta_powers[12] = beta^12`.
+    /// Inclusive end index for the four `capacity` coefficients (`beta^11..=beta^14`).
+    pub const CAPACITY_COEFF_END_IDX: usize = CAPACITY_START_IDX + 3;
+
+    /// Capacity domain coefficient index: (`beta^12`).
     ///
     /// Second capacity element. Used for encoding operation-specific data (e.g., op_code in control
     /// block messages).
@@ -237,7 +245,8 @@ pub mod bus_message {
 ///
 /// Each constant identifies a distinct bus interaction type. When encoding a message,
 /// the bus index is passed to [`Challenges::encode`] or [`Challenges::encode_sparse`],
-/// which uses `bus_prefix[bus]` as the additive base instead of bare `alpha`.
+/// which uses `bus_prefix.prefix_for_bus(bus)` (or the corresponding named field) as the additive
+/// base instead of bare `alpha`.
 ///
 /// This ensures messages from different buses are always distinct, even if they share
 /// the same coefficient layout and labels. This is a prerequisite for a future unified bus.