Skip to main content

miden_ace_codegen/layout/
keys.rs

1use super::InputLayout;
2use crate::EXT_DEGREE;
3
4/// Logical inputs required by the ACE circuit.
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
6pub enum InputKey {
7    /// Public input at the given index.
8    Public(usize),
9    /// Aux randomness α supplied as an input.
10    AuxRandAlpha,
11    /// Aux randomness β supplied as an input.
12    AuxRandBeta,
13    /// Multi-AIR β coefficient for the AIR at instance index `air` (caller order).
14    /// The verifier assigns β to the AIR at proof_order position 0 and 1 to the other.
15    /// Only present in layouts built with `num_airs >= 2`.
16    MultiAirBeta(usize),
17    /// Main trace value at (offset, index).
18    Main { offset: usize, index: usize },
19    /// Base-field coordinate for an aux trace column.
20    AuxCoord {
21        offset: usize,
22        index: usize,
23        coord: usize,
24    },
25    /// Aux bus boundary value at the given index.
26    AuxBusBoundary(usize),
27    /// Reserved stark-vars slot, kept zero. Provides word-alignment padding for the base
28    /// stark-vars block; not referenced by the production circuit.
29    Reserved,
30    /// Composition challenge used to fold constraints.
31    Alpha,
32    /// `zeta^N`, where `N` is the trace length.
33    ZPowN,
34    /// `zeta^(N / max_cycle_len)` for periodic columns.
35    ZK,
36    /// Precomputed first-row selector: `(z^N - 1) / (z - 1)`.
37    IsFirst,
38    /// Precomputed last-row selector: `(z^N - 1) / (z - g^{-1})`.
39    IsLast,
40    /// Precomputed transition selector: `z - g^{-1}`.
41    IsTransition,
42    /// Per-AIR lifted first-row selector at `z^{r_air}` (`r_air = n_max / n_air`) for the
43    /// AIR at instance index `air` (caller order). Equal to the canonical `IsFirst` when
44    /// the AIR is at log_max. Only present in layouts built with `num_airs >= 2`.
45    IsFirstAir(usize),
46    /// Per-AIR lifted last-row selector. Mirror of [`InputKey::IsFirstAir`].
47    IsLastAir(usize),
48    /// Per-AIR lifted transition selector. Mirror of [`InputKey::IsFirstAir`].
49    IsTransitionAir(usize),
50    /// First barycentric weight for quotient recomposition.
51    Weight0,
52    /// `f = h^N`, the chunk shift ratio between cosets.
53    F,
54    /// `s0 = offset^N`, the first chunk shift.
55    S0,
56    /// Base-field coordinate for a quotient chunk opening at `offset`
57    /// (0 = zeta, 1 = g * zeta).
58    QuotientChunkCoord {
59        offset: usize,
60        chunk: usize,
61        coord: usize,
62    },
63}
64
65/// Canonical InputKey → index mapping for a given layout.
66#[derive(Debug, Clone, Copy)]
67pub(crate) struct InputKeyMapper<'a> {
68    pub(super) layout: &'a InputLayout,
69}
70
71impl InputKeyMapper<'_> {
72    /// Return the input index for a key, if it exists in the layout.
73    pub(crate) fn index_of(self, key: InputKey) -> Option<usize> {
74        let layout = self.layout;
75        match key {
76            InputKey::Public(i) => layout.regions.public_values.index(i),
77            InputKey::AuxRandAlpha => Some(layout.aux_rand_alpha),
78            InputKey::AuxRandBeta => Some(layout.aux_rand_beta),
79            InputKey::MultiAirBeta(air) => {
80                layout.stark.multi_air.as_ref().and_then(|m| m.beta(air))
81            },
82            InputKey::Main { offset, index } => match offset {
83                0 => layout.regions.main_curr.index(index),
84                1 => layout.regions.main_next.index(index),
85                _ => None,
86            },
87            InputKey::AuxCoord { offset, index, coord } => {
88                if index >= layout.counts.aux_width || coord >= EXT_DEGREE {
89                    return None;
90                }
91                let local = index * EXT_DEGREE + coord;
92                match offset {
93                    0 => layout.regions.aux_curr.index(local),
94                    1 => layout.regions.aux_next.index(local),
95                    _ => None,
96                }
97            },
98            InputKey::AuxBusBoundary(i) => layout.regions.aux_bus_boundary.index(i),
99            // Extension-field stark vars.
100            InputKey::Alpha => Some(layout.stark.alpha),
101            InputKey::ZPowN => Some(layout.stark.z_pow_n),
102            InputKey::ZK => Some(layout.stark.z_k),
103            InputKey::IsFirst => Some(layout.stark.is_first),
104            InputKey::IsLast => Some(layout.stark.is_last),
105            InputKey::IsTransition => Some(layout.stark.is_transition),
106            InputKey::IsFirstAir(air) => {
107                layout.stark.multi_air.as_ref().and_then(|m| m.selector(air, 0))
108            },
109            InputKey::IsLastAir(air) => {
110                layout.stark.multi_air.as_ref().and_then(|m| m.selector(air, 1))
111            },
112            InputKey::IsTransitionAir(air) => {
113                layout.stark.multi_air.as_ref().and_then(|m| m.selector(air, 2))
114            },
115            InputKey::Reserved => Some(layout.stark.reserved),
116            // Base-field stark vars (stored as (val, 0) in the EF slot).
117            InputKey::Weight0 => Some(layout.stark.weight0),
118            InputKey::F => Some(layout.stark.f),
119            InputKey::S0 => Some(layout.stark.s0),
120            InputKey::QuotientChunkCoord { offset, chunk, coord } => {
121                if chunk >= layout.counts.num_quotient_chunks || coord >= EXT_DEGREE {
122                    return None;
123                }
124                let idx = chunk * EXT_DEGREE + coord;
125                match offset {
126                    0 => layout.regions.quotient_curr.index(idx),
127                    1 => layout.regions.quotient_next.index(idx),
128                    _ => None,
129                }
130            },
131        }
132    }
133}