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    /// Preprocessed trace value at (offset, index).
20    Preprocessed { offset: usize, index: usize },
21    /// Main trace value at (offset, index).
22    Main { offset: usize, index: usize },
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    /// Reserved stark-vars slot, kept zero.
32    Reserved,
33    /// Composition challenge used to fold constraints.
34    Alpha,
35    /// `zeta^N_max`, where `N_max` is the maximum trace length represented by the circuit.
36    ZPowN,
37    /// Periodic-column evaluation basis `zeta^(N_max / shared_period)`.
38    /// A period-`p` column is evaluated at `ZK^(shared_period / p)`.
39    ZK,
40    /// Precomputed first-row selector: `(z^N - 1) / (z - 1)`.
41    IsFirst,
42    /// Precomputed last-row selector: `(z^N - 1) / (z - g^-1)`.
43    IsLast,
44    /// Precomputed transition selector: `z - g^-1`.
45    IsTransition,
46    /// Per-AIR lifted first-row selector.
47    IsFirstAir(usize),
48    /// Per-AIR lifted last-row selector.
49    IsLastAir(usize),
50    /// Per-AIR lifted transition selector.
51    IsTransitionAir(usize),
52    /// First barycentric weight for quotient recomposition.
53    Weight0,
54    /// `f = h^N`, the chunk shift ratio between cosets.
55    F,
56    /// `s0 = offset^N`, the first chunk shift.
57    S0,
58    /// Base-field coordinate for a quotient chunk opening at `offset`
59    /// (0 = zeta, 1 = g * zeta).
60    QuotientChunkCoord {
61        offset: usize,
62        chunk: usize,
63        coord: usize,
64    },
65}
66
67/// Canonical InputKey -> index mapping for a given layout.
68#[derive(Debug, Clone, Copy)]
69pub(crate) struct InputKeyMapper<'a> {
70    pub(super) layout: &'a InputLayout,
71}
72
73impl InputKeyMapper<'_> {
74    /// Return the input index for a key, if it exists in the layout.
75    pub(crate) fn index_of(self, key: InputKey) -> Option<usize> {
76        let layout = self.layout;
77        match key {
78            InputKey::Public(i) => layout.regions.public_values.index(i),
79            InputKey::AuxRandAlpha => Some(layout.aux_rand_alpha),
80            InputKey::AuxRandBeta => Some(layout.aux_rand_beta),
81            InputKey::MultiAirFoldBeta => layout.stark.multi_air_fold_beta_index(),
82            InputKey::Preprocessed { offset, index } => match offset {
83                0 => layout.regions.preprocessed_curr.index(index),
84                1 => layout.regions.preprocessed_next.index(index),
85                _ => None,
86            },
87            InputKey::Main { offset, index } => match offset {
88                0 => layout.regions.main_curr.index(index),
89                1 => layout.regions.main_next.index(index),
90                _ => None,
91            },
92            InputKey::AuxCoord { offset, index, coord } => {
93                if index >= layout.counts.aux_width || coord >= EXT_DEGREE {
94                    return None;
95                }
96                let local = index * EXT_DEGREE + coord;
97                match offset {
98                    0 => layout.regions.aux_curr.index(local),
99                    1 => layout.regions.aux_next.index(local),
100                    _ => None,
101                }
102            },
103            InputKey::AuxBusBoundary(i) => layout.regions.aux_bus_boundary.index(i),
104            InputKey::Reserved => Some(layout.stark.reserved),
105            InputKey::Alpha => Some(layout.stark.alpha),
106            InputKey::ZPowN => Some(layout.stark.z_pow_n),
107            InputKey::ZK => Some(layout.stark.z_k),
108            InputKey::IsFirst => Some(layout.stark.is_first),
109            InputKey::IsLast => Some(layout.stark.is_last),
110            InputKey::IsTransition => Some(layout.stark.is_transition),
111            InputKey::IsFirstAir(i) => {
112                layout.stark.air_selector_index(i, AIR_SELECTOR_FIRST_OFFSET)
113            },
114            InputKey::IsLastAir(i) => layout.stark.air_selector_index(i, AIR_SELECTOR_LAST_OFFSET),
115            InputKey::IsTransitionAir(i) => {
116                layout.stark.air_selector_index(i, AIR_SELECTOR_TRANSITION_OFFSET)
117            },
118            InputKey::Weight0 => Some(layout.stark.weight0),
119            InputKey::F => Some(layout.stark.f),
120            InputKey::S0 => Some(layout.stark.s0),
121            InputKey::QuotientChunkCoord { offset, chunk, coord } => {
122                if chunk >= layout.counts.num_quotient_chunks || coord >= EXT_DEGREE {
123                    return None;
124                }
125                let idx = chunk * EXT_DEGREE + coord;
126                match offset {
127                    0 => layout.regions.quotient_curr.index(idx),
128                    1 => layout.regions.quotient_next.index(idx),
129                    _ => None,
130                }
131            },
132        }
133    }
134}