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 Core. Set to β if Core is at proof_order position 0
14    /// (`core_height ≤ chiplets_height`), else 1. Only present in `is_multi_air = true`.
15    MultiAirBetaCore,
16    /// Multi-AIR β coefficient for Chiplets. Complement of `MultiAirBetaCore`.
17    MultiAirBetaChip,
18    /// Main trace value at (offset, index).
19    Main {
20        offset: usize,
21        index: usize,
22    },
23    /// Base-field coordinate for an aux trace column.
24    AuxCoord {
25        offset: usize,
26        index: usize,
27        coord: usize,
28    },
29    /// Aux bus boundary value at the given index.
30    AuxBusBoundary(usize),
31    /// Variable-length public input reduction at the given group index.
32    VlpiReduction(usize),
33    /// Batching challenge gamma for combining the constraint evaluation with the
34    /// auxiliary trace boundary checks.
35    Gamma,
36    /// Composition challenge used to fold constraints.
37    Alpha,
38    /// `zeta^N`, where `N` is the trace length.
39    ZPowN,
40    /// `zeta^(N / max_cycle_len)` for periodic columns.
41    ZK,
42    /// Precomputed first-row selector: `(z^N - 1) / (z - 1)`.
43    IsFirst,
44    /// Precomputed last-row selector: `(z^N - 1) / (z - g^{-1})`.
45    IsLast,
46    /// Precomputed transition selector: `z - g^{-1}`.
47    IsTransition,
48    /// Per-AIR lifted selectors for Core at `z^{r_core}` (`r_core = n_max / n_core`).
49    /// Equal to the canonical `IsFirst`/`IsLast`/`IsTransition` when Core is at log_max.
50    /// Only present in `is_multi_air = true`.
51    IsFirstCore,
52    IsLastCore,
53    IsTransitionCore,
54    /// Per-AIR lifted selectors for Chiplets at `z^{r_chip}`. Mirror of `*Core`.
55    IsFirstChip,
56    IsLastChip,
57    IsTransitionChip,
58    /// First barycentric weight for quotient recomposition.
59    Weight0,
60    /// `f = h^N`, the chunk shift ratio between cosets.
61    F,
62    /// `s0 = offset^N`, the first chunk shift.
63    S0,
64    /// Base-field coordinate for a quotient chunk opening at `offset`
65    /// (0 = zeta, 1 = g * zeta).
66    QuotientChunkCoord {
67        offset: usize,
68        chunk: usize,
69        coord: usize,
70    },
71}
72
73/// Canonical InputKey → index mapping for a given layout.
74#[derive(Debug, Clone, Copy)]
75pub(crate) struct InputKeyMapper<'a> {
76    pub(super) layout: &'a InputLayout,
77}
78
79impl InputKeyMapper<'_> {
80    /// Return the input index for a key, if it exists in the layout.
81    pub(crate) fn index_of(self, key: InputKey) -> Option<usize> {
82        let layout = self.layout;
83        match key {
84            InputKey::Public(i) => layout.regions.public_values.index(i),
85            InputKey::AuxRandAlpha => Some(layout.aux_rand_alpha),
86            InputKey::AuxRandBeta => Some(layout.aux_rand_beta),
87            InputKey::MultiAirBetaCore => layout.stark.multi_air_beta_core,
88            InputKey::MultiAirBetaChip => layout.stark.multi_air_beta_chip,
89            InputKey::Main { offset, index } => match offset {
90                0 => layout.regions.main_curr.index(index),
91                1 => layout.regions.main_next.index(index),
92                _ => None,
93            },
94            InputKey::AuxCoord { offset, index, coord } => {
95                if index >= layout.counts.aux_width || coord >= EXT_DEGREE {
96                    return None;
97                }
98                let local = index * EXT_DEGREE + coord;
99                match offset {
100                    0 => layout.regions.aux_curr.index(local),
101                    1 => layout.regions.aux_next.index(local),
102                    _ => None,
103                }
104            },
105            InputKey::AuxBusBoundary(i) => layout.regions.aux_bus_boundary.index(i),
106            InputKey::VlpiReduction(i) => {
107                let local = i * layout.vlpi_stride;
108                layout.regions.vlpi_reductions.index(local)
109            },
110            // Extension-field stark vars.
111            InputKey::Alpha => Some(layout.stark.alpha),
112            InputKey::ZPowN => Some(layout.stark.z_pow_n),
113            InputKey::ZK => Some(layout.stark.z_k),
114            InputKey::IsFirst => Some(layout.stark.is_first),
115            InputKey::IsLast => Some(layout.stark.is_last),
116            InputKey::IsTransition => Some(layout.stark.is_transition),
117            InputKey::IsFirstCore => layout.stark.is_first_core,
118            InputKey::IsLastCore => layout.stark.is_last_core,
119            InputKey::IsTransitionCore => layout.stark.is_transition_core,
120            InputKey::IsFirstChip => layout.stark.is_first_chip,
121            InputKey::IsLastChip => layout.stark.is_last_chip,
122            InputKey::IsTransitionChip => layout.stark.is_transition_chip,
123            InputKey::Gamma => Some(layout.stark.gamma),
124            // Base-field stark vars (stored as (val, 0) in the EF slot).
125            InputKey::Weight0 => Some(layout.stark.weight0),
126            InputKey::F => Some(layout.stark.f),
127            InputKey::S0 => Some(layout.stark.s0),
128            InputKey::QuotientChunkCoord { offset, chunk, coord } => {
129                if chunk >= layout.counts.num_quotient_chunks || coord >= EXT_DEGREE {
130                    return None;
131                }
132                let idx = chunk * EXT_DEGREE + coord;
133                match offset {
134                    0 => layout.regions.quotient_curr.index(idx),
135                    1 => layout.regions.quotient_next.index(idx),
136                    _ => None,
137                }
138            },
139        }
140    }
141}