Skip to main content

miden_ace_codegen/layout/
plan.rs

1use super::InputKey;
2use crate::EXT_DEGREE;
3
4/// A contiguous region of inputs within the ACE READ layout.
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub(crate) struct InputRegion {
7    pub offset: usize,
8    pub width: usize,
9}
10
11impl InputRegion {
12    /// Map a region-local index to a global input index.
13    pub fn index(&self, local: usize) -> Option<usize> {
14        (local < self.width).then(|| self.offset + local)
15    }
16}
17
18/// Counts needed to build the ACE input layout.
19#[derive(Debug, Clone, Copy)]
20pub struct InputCounts {
21    /// Width of the main trace.
22    pub width: usize,
23    /// Width of the aux trace.
24    pub aux_width: usize,
25    /// Number of committed boundary values (accumulator column finals).
26    pub num_aux_boundary: usize,
27    /// Number of public inputs.
28    pub num_public: usize,
29    /// Number of randomness challenges used by the AIR.
30    pub num_randomness: usize,
31    /// Number of periodic columns.
32    pub num_periodic: usize,
33    /// Number of quotient chunks.
34    pub num_quotient_chunks: usize,
35}
36
37/// Grouped regions for the ACE input layout.
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub(crate) struct LayoutRegions {
40    /// Region containing fixed-length public values.
41    pub public_values: InputRegion,
42    /// Region containing randomness inputs (alpha, beta).
43    pub randomness: InputRegion,
44    /// Main trace OOD values at `zeta`.
45    pub main_curr: InputRegion,
46    /// Aux trace OOD coordinates at `zeta`.
47    pub aux_curr: InputRegion,
48    /// Quotient chunk OOD coordinates at `zeta`.
49    pub quotient_curr: InputRegion,
50    /// Main trace OOD values at `g * zeta`.
51    pub main_next: InputRegion,
52    /// Aux trace OOD coordinates at `g * zeta`.
53    pub aux_next: InputRegion,
54    /// Quotient chunk OOD coordinates at `g * zeta`.
55    pub quotient_next: InputRegion,
56    /// Aux bus boundary values.
57    pub aux_bus_boundary: InputRegion,
58    /// Stark variables (selectors, powers, weights).
59    pub stark_vars: InputRegion,
60}
61
62/// Indexes of canonical verifier scalars inside the stark-vars block.
63///
64/// Every slot in the ACE input array is an extension-field (EF) element --
65/// the circuit operates entirely in the extension field. However, some of
66/// these scalars are inherently base-field values that the MASM verifier
67/// stores as `(val, 0)` in the EF slot.
68///
69/// See the module documentation on [`super::super::dag::lower`] for how each
70/// variable enters the verifier expression.
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
72pub(crate) struct StarkVarIndices {
73    // -- Extension-field values (slots 0-5) --
74    /// Composition challenge `alpha` for folding constraints.
75    pub alpha: usize,
76    /// `zeta^N` where N is the trace length.
77    pub z_pow_n: usize,
78    /// `zeta^(N / max_cycle_len)` for periodic column evaluation.
79    pub z_k: usize,
80    /// Precomputed first-row selector: `(z^N - 1) / (z - 1)`.
81    pub is_first: usize,
82    /// Precomputed last-row selector: `(z^N - 1) / (z - g^{-1})`.
83    pub is_last: usize,
84    /// Precomputed transition selector: `z - g^{-1}`.
85    pub is_transition: usize,
86    /// Reserved word-alignment slot for the base stark-vars block (kept zero).
87    pub reserved: usize,
88
89    // -- Base-field values stored as (val, 0) in EF slots --
90    /// First barycentric weight `1 / (k * s0^{k-1})`.
91    pub weight0: usize,
92    /// `f = h^N` (chunk shift ratio between cosets).
93    pub f: usize,
94    /// `s0 = offset^N` (first chunk shift).
95    pub s0: usize,
96
97    // -- Multi-AIR additions (only present when the layout was built with `num_airs >= 2`) --
98    /// Slot block for the per-AIR β coefficients and lifted selectors.
99    pub multi_air: Option<MultiAirVarIndices>,
100}
101
102/// Indexes of the per-AIR stark-vars slots inside a multi-AIR layout.
103///
104/// The block occupies `4 * num_airs` consecutive EF slots starting at `base`:
105/// first one β coefficient per AIR (instance order), then one
106/// `(is_first, is_last, is_transition)` lifted-selector triple per AIR.
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub(crate) struct MultiAirVarIndices {
109    /// First slot of the block (the β coefficient for AIR 0).
110    pub base: usize,
111    /// Number of AIRs the layout was built for.
112    pub num_airs: usize,
113}
114
115impl MultiAirVarIndices {
116    /// Slot of the β coefficient for the AIR at instance index `air`.
117    pub fn beta(&self, air: usize) -> Option<usize> {
118        (air < self.num_airs).then(|| self.base + air)
119    }
120
121    /// Slot of the lifted selector `sel` (0 = is_first, 1 = is_last, 2 = is_transition)
122    /// for the AIR at instance index `air`.
123    pub fn selector(&self, air: usize, sel: usize) -> Option<usize> {
124        (air < self.num_airs && sel < 3).then(|| self.base + self.num_airs + 3 * air + sel)
125    }
126
127    /// Total number of slots in the block.
128    pub fn width(&self) -> usize {
129        4 * self.num_airs
130    }
131}
132
133/// ACE input layout for Plonky3-based verifier logic.
134///
135/// This describes the exact ordering and alignment of inputs consumed by the
136/// ACE chiplet (READ section).
137#[derive(Debug, Clone)]
138pub struct InputLayout {
139    /// Grouped regions for the ACE input layout.
140    pub(crate) regions: LayoutRegions,
141    /// Input index for aux randomness alpha.
142    pub(crate) aux_rand_alpha: usize,
143    /// Input index for aux randomness beta.
144    pub(crate) aux_rand_beta: usize,
145    /// Indexes into the stark-vars region.
146    pub(crate) stark: StarkVarIndices,
147    /// Total number of inputs (length of the READ section).
148    pub total_inputs: usize,
149    /// Counts used to derive the layout.
150    pub counts: InputCounts,
151}
152
153impl InputLayout {
154    pub(crate) fn mapper(&self) -> super::InputKeyMapper<'_> {
155        super::InputKeyMapper { layout: self }
156    }
157
158    /// Map a logical `InputKey` into the flat input index, if present.
159    pub fn index(&self, key: InputKey) -> Option<usize> {
160        self.mapper().index_of(key)
161    }
162
163    /// Validate internal invariants for this layout (region sizes, key ranges, randomness inputs).
164    pub(crate) fn validate(&self) {
165        let mut max_end = 0usize;
166        for region in [
167            self.regions.public_values,
168            self.regions.randomness,
169            self.regions.main_curr,
170            self.regions.aux_curr,
171            self.regions.quotient_curr,
172            self.regions.main_next,
173            self.regions.aux_next,
174            self.regions.quotient_next,
175            self.regions.aux_bus_boundary,
176            self.regions.stark_vars,
177        ] {
178            max_end = max_end.max(region.offset.saturating_add(region.width));
179        }
180
181        assert!(max_end <= self.total_inputs, "regions exceed total_inputs");
182
183        let aux_coord_width = self.counts.aux_width * EXT_DEGREE;
184        assert_eq!(self.regions.aux_curr.width, aux_coord_width, "aux_curr width mismatch");
185        assert_eq!(self.regions.aux_next.width, aux_coord_width, "aux_next width mismatch");
186
187        let quotient_width = self.counts.num_quotient_chunks * EXT_DEGREE;
188        assert_eq!(
189            self.regions.quotient_curr.width, quotient_width,
190            "quotient_curr width mismatch"
191        );
192        assert_eq!(
193            self.regions.quotient_next.width, quotient_width,
194            "quotient_next width mismatch"
195        );
196        assert_eq!(
197            self.regions.aux_bus_boundary.width, self.counts.num_aux_boundary,
198            "aux bus boundary width mismatch"
199        );
200
201        let stark_start = self.regions.stark_vars.offset;
202        let stark_end = stark_start + self.regions.stark_vars.width;
203        let check = |name: &str, idx: usize| {
204            assert!(idx >= stark_start && idx < stark_end, "stark var {name} out of range");
205        };
206        // Extension-field slots.
207        check("alpha", self.stark.alpha);
208        check("z_pow_n", self.stark.z_pow_n);
209        check("z_k", self.stark.z_k);
210        check("is_first", self.stark.is_first);
211        check("is_last", self.stark.is_last);
212        check("is_transition", self.stark.is_transition);
213        check("reserved", self.stark.reserved);
214        // Base-field slots (stored as (val, 0) in the EF slot).
215        check("weight0", self.stark.weight0);
216        check("f", self.stark.f);
217        check("s0", self.stark.s0);
218        if let Some(multi_air) = &self.stark.multi_air {
219            check("multi_air block start", multi_air.base);
220            check("multi_air block end", multi_air.base + multi_air.width() - 1);
221        }
222
223        let rand_start = self.regions.randomness.offset;
224        let rand_end = rand_start + self.regions.randomness.width;
225        assert!(
226            self.aux_rand_alpha >= rand_start && self.aux_rand_alpha < rand_end,
227            "aux_rand_alpha out of randomness region"
228        );
229        assert!(
230            self.aux_rand_beta >= rand_start && self.aux_rand_beta < rand_end,
231            "aux_rand_beta out of randomness region"
232        );
233    }
234}