Skip to main content

miden_ace_codegen/layout/
plan.rs

1use core::num::NonZeroUsize;
2
3use super::{InputKey, SELECTORS_PER_AIR};
4use crate::EXT_DEGREE;
5
6/// A contiguous region of inputs within the ACE READ layout.
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub(crate) struct InputRegion {
9    pub offset: usize,
10    pub width: usize,
11}
12
13impl InputRegion {
14    /// Map a region-local index to a global input index.
15    pub fn index(&self, local: usize) -> Option<usize> {
16        (local < self.width).then(|| self.offset + local)
17    }
18}
19
20/// Counts needed to build the ACE input layout.
21#[derive(Debug, Clone, Copy)]
22pub struct InputCounts {
23    /// Width of the preprocessed trace.
24    pub preprocessed_width: usize,
25    /// Width of the main trace.
26    pub width: usize,
27    /// Width of the aux trace.
28    pub aux_width: usize,
29    /// Number of committed boundary values.
30    pub num_aux_boundary: usize,
31    /// Number of public inputs.
32    pub num_public: usize,
33    /// Number of randomness challenges used by the AIR.
34    pub num_randomness: usize,
35    /// Number of quotient chunks.
36    pub num_quotient_chunks: usize,
37}
38
39/// Grouped regions for the ACE input layout.
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub(crate) struct LayoutRegions {
42    /// Region containing fixed public values.
43    pub public_values: InputRegion,
44    /// Region containing randomness inputs (alpha, beta).
45    pub randomness: InputRegion,
46    /// Preprocessed trace OOD values at `zeta`.
47    pub preprocessed_curr: InputRegion,
48    /// Main trace OOD values at `zeta`.
49    pub main_curr: InputRegion,
50    /// Aux trace OOD coordinates at `zeta`.
51    pub aux_curr: InputRegion,
52    /// Quotient chunk OOD coordinates at `zeta`.
53    pub quotient_curr: InputRegion,
54    /// Preprocessed trace OOD values at `g * zeta`.
55    pub preprocessed_next: InputRegion,
56    /// Main trace OOD values at `g * zeta`.
57    pub main_next: InputRegion,
58    /// Aux trace OOD coordinates at `g * zeta`.
59    pub aux_next: InputRegion,
60    /// Quotient chunk OOD coordinates at `g * zeta`.
61    pub quotient_next: InputRegion,
62    /// Aux bus boundary values.
63    pub aux_bus_boundary: InputRegion,
64    /// Stark variables (selectors, powers, weights).
65    pub stark_vars: InputRegion,
66}
67
68/// Indexes of verifier scalars inside the stark-vars block.
69#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub(crate) struct StarkVarIndices {
71    /// Composition challenge `alpha`.
72    pub alpha: usize,
73    /// `zeta^N_max`, where `N_max` is the maximum trace length represented by the circuit.
74    pub z_pow_n: usize,
75    /// Shared periodic-column basis `zeta^(N_max / shared_period)`.
76    pub z_k: usize,
77    /// First-row selector.
78    pub is_first: usize,
79    /// Last-row selector.
80    pub is_last: usize,
81    /// Transition selector.
82    pub is_transition: usize,
83    /// Reserved slot kept zero.
84    pub reserved: usize,
85    /// First barycentric weight.
86    pub weight0: usize,
87    /// Chunk shift ratio.
88    pub f: usize,
89    /// First coset shift.
90    pub s0: usize,
91    /// Extra slots present only for a multi-AIR layout.
92    pub multi_air: Option<MultiAirIndices>,
93}
94
95/// Stark-var slots added by a multi-AIR layout.
96#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97pub(crate) struct MultiAirIndices {
98    /// Number of AIR instances represented in the layout.
99    pub air_count: NonZeroUsize,
100    /// Multi-AIR fold beta slot.
101    pub fold_beta: usize,
102    /// First per-AIR selector slot.
103    pub selector_start: usize,
104}
105
106/// ACE input layout for circuit evaluation.
107#[derive(Debug, Clone)]
108pub struct InputLayout {
109    /// Grouped regions for the ACE input layout.
110    pub(crate) regions: LayoutRegions,
111    /// Input index for aux randomness alpha.
112    pub(crate) aux_rand_alpha: usize,
113    /// Input index for aux randomness beta.
114    pub(crate) aux_rand_beta: usize,
115    /// Indexes into the stark-vars region.
116    pub(crate) stark: StarkVarIndices,
117    /// Total number of ACE READ inputs.
118    pub total_inputs: usize,
119    /// Counts used to derive the layout.
120    pub counts: InputCounts,
121}
122
123impl InputLayout {
124    pub(crate) fn mapper(&self) -> super::InputKeyMapper<'_> {
125        super::InputKeyMapper { layout: self }
126    }
127
128    /// Map a logical `InputKey` into the flat input index, if present.
129    pub fn index(&self, key: InputKey) -> Option<usize> {
130        self.mapper().index_of(key)
131    }
132
133    /// Validate internal layout invariants.
134    pub(crate) fn validate(&self) {
135        let mut max_end = 0usize;
136        for region in [
137            self.regions.public_values,
138            self.regions.randomness,
139            self.regions.preprocessed_curr,
140            self.regions.main_curr,
141            self.regions.aux_curr,
142            self.regions.quotient_curr,
143            self.regions.preprocessed_next,
144            self.regions.main_next,
145            self.regions.aux_next,
146            self.regions.quotient_next,
147            self.regions.aux_bus_boundary,
148            self.regions.stark_vars,
149        ] {
150            max_end = max_end.max(region.offset.saturating_add(region.width));
151        }
152
153        assert!(max_end <= self.total_inputs, "regions exceed total_inputs");
154
155        assert_eq!(
156            self.regions.preprocessed_curr.width, self.counts.preprocessed_width,
157            "preprocessed_curr width mismatch"
158        );
159        assert_eq!(
160            self.regions.preprocessed_next.width, self.counts.preprocessed_width,
161            "preprocessed_next width mismatch"
162        );
163        assert_eq!(self.regions.main_curr.width, self.counts.width, "main_curr width mismatch");
164        assert_eq!(self.regions.main_next.width, self.counts.width, "main_next width mismatch");
165
166        let aux_coord_width = self.counts.aux_width * EXT_DEGREE;
167        assert_eq!(self.regions.aux_curr.width, aux_coord_width, "aux_curr width mismatch");
168        assert_eq!(self.regions.aux_next.width, aux_coord_width, "aux_next width mismatch");
169
170        let quotient_width = self.counts.num_quotient_chunks * EXT_DEGREE;
171        assert_eq!(
172            self.regions.quotient_curr.width, quotient_width,
173            "quotient_curr width mismatch"
174        );
175        assert_eq!(
176            self.regions.quotient_next.width, quotient_width,
177            "quotient_next width mismatch"
178        );
179        assert_eq!(
180            self.regions.aux_bus_boundary.width, self.counts.num_aux_boundary,
181            "aux bus boundary width mismatch"
182        );
183
184        let stark_start = self.regions.stark_vars.offset;
185        let stark_end = stark_start + self.regions.stark_vars.width;
186        let check = |name: &str, idx: usize| {
187            assert!(idx >= stark_start && idx < stark_end, "stark var {name} out of range");
188        };
189        check("alpha", self.stark.alpha);
190        check("z_pow_n", self.stark.z_pow_n);
191        check("z_k", self.stark.z_k);
192        check("is_first", self.stark.is_first);
193        check("is_last", self.stark.is_last);
194        check("is_transition", self.stark.is_transition);
195        check("reserved", self.stark.reserved);
196        check("weight0", self.stark.weight0);
197        check("f", self.stark.f);
198        check("s0", self.stark.s0);
199        if let Some(multi_air) = self.stark.multi_air {
200            check("multi_air_fold_beta", multi_air.fold_beta);
201            for i in 0..(multi_air.air_count.get() * SELECTORS_PER_AIR) {
202                check("air_selector", multi_air.selector_start + i);
203            }
204        }
205
206        let rand_start = self.regions.randomness.offset;
207        let rand_end = rand_start + self.regions.randomness.width;
208        assert!(
209            self.aux_rand_alpha >= rand_start && self.aux_rand_alpha < rand_end,
210            "aux_rand_alpha out of randomness region"
211        );
212        assert!(
213            self.aux_rand_beta >= rand_start && self.aux_rand_beta < rand_end,
214            "aux_rand_beta out of randomness region"
215        );
216    }
217}
218
219impl StarkVarIndices {
220    pub(crate) fn multi_air_fold_beta_index(&self) -> Option<usize> {
221        self.multi_air.map(|multi_air| multi_air.fold_beta)
222    }
223
224    pub(crate) fn air_selector_index(
225        &self,
226        air_index: usize,
227        selector_offset: usize,
228    ) -> Option<usize> {
229        let multi_air = self.multi_air?;
230        (air_index < multi_air.air_count.get() && selector_offset < SELECTORS_PER_AIR)
231            .then_some(multi_air.selector_start + air_index * SELECTORS_PER_AIR + selector_offset)
232    }
233}