Skip to main content

miden_ace_codegen/layout/
keys.rs

1use super::InputLayout;
2use crate::EXT_DEGREE;
3
4const AIR_SELECTOR_FIRST_OFFSET: usize = 0;
5const AIR_SELECTOR_LAST_OFFSET: usize = 1;
6const AIR_SELECTOR_TRANSITION_OFFSET: usize = 2;
7
8/// Logical inputs required by the ACE circuit.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub enum InputKey {
11    /// Public input at the given index.
12    Public(usize),
13    /// Aux randomness alpha supplied as an input.
14    AuxRandAlpha,
15    /// Aux randomness beta supplied as an input.
16    AuxRandBeta,
17    /// Challenge used to fold per-AIR constraint roots in proof order.
18    MultiAirFoldBeta,
19    /// Fold coefficient for the AIR instance at the given index.
20    ///
21    /// Factored-emission-only: this key has no READ-layout slot. The factored emitter binds it
22    /// to a shuffle-section gate holding a power of [`InputKey::MultiAirFoldBeta`]; which power
23    /// is chosen per assembled circuit by the caller of
24    /// `FactoredAceCircuit::assemble`.
25    MultiAirFoldCoeff(usize),
26    /// Preprocessed trace value at (offset, index).
27    Preprocessed { offset: usize, index: usize },
28    /// Main trace value at (offset, index).
29    Main { offset: usize, index: usize },
30    /// Base-field coordinate for an aux trace column.
31    AuxCoord {
32        offset: usize,
33        index: usize,
34        coord: usize,
35    },
36    /// Aux bus boundary value at the given index.
37    AuxBusBoundary(usize),
38    /// Reserved stark-vars slot, kept zero.
39    Reserved,
40    /// Composition challenge used to fold constraints.
41    Alpha,
42    /// `zeta^N_max`, where `N_max` is the maximum trace length represented by the circuit.
43    ZPowN,
44    /// Periodic-column evaluation basis `zeta^(N_max / shared_period)`.
45    /// A period-`p` column is evaluated at `ZK^(shared_period / p)`.
46    ZK,
47    /// Precomputed first-row selector: `(z^N - 1) / (z - 1)`.
48    IsFirst,
49    /// Precomputed last-row selector: `(z^N - 1) / (z - g^-1)`.
50    IsLast,
51    /// Precomputed transition selector: `z - g^-1`.
52    IsTransition,
53    /// Per-AIR lifted first-row selector.
54    IsFirstAir(usize),
55    /// Per-AIR lifted last-row selector.
56    IsLastAir(usize),
57    /// Per-AIR lifted transition selector.
58    IsTransitionAir(usize),
59    /// First barycentric weight for quotient recomposition.
60    Weight0,
61    /// Primitive `D`-th root of unity, where `D` is the quotient chunk count; the chunk shift
62    /// ratio between cosets after taking the trace-height power.
63    F,
64    /// `s0 = offset^N`, the first chunk shift.
65    S0,
66    /// Base-field coordinate for a quotient chunk opening at `offset`
67    /// (0 = zeta, 1 = g * zeta).
68    QuotientChunkCoord {
69        offset: usize,
70        chunk: usize,
71        coord: usize,
72    },
73}
74
75/// Canonical InputKey -> index mapping for a given layout.
76#[derive(Debug, Clone, Copy)]
77pub(crate) struct InputKeyMapper<'a> {
78    pub(super) layout: &'a InputLayout,
79}
80
81impl InputKeyMapper<'_> {
82    /// Return the input index for a key, if it exists in the layout.
83    pub(crate) fn index_of(self, key: InputKey) -> Option<usize> {
84        let layout = self.layout;
85        match key {
86            InputKey::Public(i) => layout.regions.public_values.index(i),
87            InputKey::AuxRandAlpha => Some(layout.aux_rand_alpha),
88            InputKey::AuxRandBeta => Some(layout.aux_rand_beta),
89            InputKey::MultiAirFoldBeta => layout.stark.multi_air_fold_beta_index(),
90            // Factored-emission-only key: never present in the READ layout.
91            InputKey::MultiAirFoldCoeff(_) => None,
92            InputKey::Preprocessed { offset, index } => match offset {
93                0 => layout.regions.preprocessed_curr.index(index),
94                1 => layout.regions.preprocessed_next.index(index),
95                _ => None,
96            },
97            InputKey::Main { offset, index } => match offset {
98                0 => layout.regions.main_curr.index(index),
99                1 => layout.regions.main_next.index(index),
100                _ => None,
101            },
102            InputKey::AuxCoord { offset, index, coord } => {
103                if index >= layout.counts.aux_width || coord >= EXT_DEGREE {
104                    return None;
105                }
106                let local = index * EXT_DEGREE + coord;
107                match offset {
108                    0 => layout.regions.aux_curr.index(local),
109                    1 => layout.regions.aux_next.index(local),
110                    _ => None,
111                }
112            },
113            InputKey::AuxBusBoundary(i) => layout.regions.aux_bus_boundary.index(i),
114            InputKey::Reserved => Some(layout.stark.reserved),
115            InputKey::Alpha => Some(layout.stark.alpha),
116            InputKey::ZPowN => Some(layout.stark.z_pow_n),
117            InputKey::ZK => Some(layout.stark.z_k),
118            InputKey::IsFirst => Some(layout.stark.is_first),
119            InputKey::IsLast => Some(layout.stark.is_last),
120            InputKey::IsTransition => Some(layout.stark.is_transition),
121            InputKey::IsFirstAir(i) => {
122                layout.stark.air_selector_index(i, AIR_SELECTOR_FIRST_OFFSET)
123            },
124            InputKey::IsLastAir(i) => layout.stark.air_selector_index(i, AIR_SELECTOR_LAST_OFFSET),
125            InputKey::IsTransitionAir(i) => {
126                layout.stark.air_selector_index(i, AIR_SELECTOR_TRANSITION_OFFSET)
127            },
128            InputKey::Weight0 => Some(layout.stark.weight0),
129            InputKey::F => Some(layout.stark.f),
130            InputKey::S0 => Some(layout.stark.s0),
131            InputKey::QuotientChunkCoord { offset, chunk, coord } => {
132                if chunk >= layout.counts.num_quotient_chunks || coord >= EXT_DEGREE {
133                    return None;
134                }
135                let idx = chunk * EXT_DEGREE + coord;
136                match offset {
137                    0 => layout.regions.quotient_curr.index(idx),
138                    1 => layout.regions.quotient_next.index(idx),
139                    _ => None,
140                }
141            },
142        }
143    }
144}