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}