miden_ace_codegen/layout/
plan.rs1use core::num::NonZeroUsize;
2
3use super::{InputKey, SELECTORS_PER_AIR};
4use crate::EXT_DEGREE;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub(crate) struct InputRegion {
9 pub offset: usize,
10 pub width: usize,
11}
12
13impl InputRegion {
14 pub fn index(&self, local: usize) -> Option<usize> {
16 (local < self.width).then(|| self.offset + local)
17 }
18}
19
20#[derive(Debug, Clone, Copy)]
22pub struct InputCounts {
23 pub preprocessed_width: usize,
25 pub width: usize,
27 pub aux_width: usize,
29 pub num_aux_boundary: usize,
31 pub num_public: usize,
33 pub num_randomness: usize,
35 pub num_quotient_chunks: usize,
37}
38
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub(crate) struct LayoutRegions {
42 pub public_values: InputRegion,
44 pub randomness: InputRegion,
46 pub preprocessed_curr: InputRegion,
48 pub main_curr: InputRegion,
50 pub aux_curr: InputRegion,
52 pub quotient_curr: InputRegion,
54 pub preprocessed_next: InputRegion,
56 pub main_next: InputRegion,
58 pub aux_next: InputRegion,
60 pub quotient_next: InputRegion,
62 pub aux_bus_boundary: InputRegion,
64 pub stark_vars: InputRegion,
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
70pub(crate) struct StarkVarIndices {
71 pub alpha: usize,
73 pub z_pow_n: usize,
75 pub z_k: usize,
77 pub is_first: usize,
79 pub is_last: usize,
81 pub is_transition: usize,
83 pub reserved: usize,
85 pub weight0: usize,
87 pub f: usize,
89 pub s0: usize,
91 pub multi_air: Option<MultiAirIndices>,
93}
94
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97pub(crate) struct MultiAirIndices {
98 pub air_count: NonZeroUsize,
100 pub fold_beta: usize,
102 pub selector_start: usize,
104}
105
106#[derive(Debug, Clone)]
108pub struct InputLayout {
109 pub(crate) regions: LayoutRegions,
111 pub(crate) aux_rand_alpha: usize,
113 pub(crate) aux_rand_beta: usize,
115 pub(crate) stark: StarkVarIndices,
117 pub total_inputs: usize,
119 pub counts: InputCounts,
121}
122
123impl InputLayout {
124 pub(crate) fn mapper(&self) -> super::InputKeyMapper<'_> {
125 super::InputKeyMapper { layout: self }
126 }
127
128 pub fn index(&self, key: InputKey) -> Option<usize> {
130 self.mapper().index_of(key)
131 }
132
133 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}