miden_ace_codegen/layout/
policy.rs

1use core::num::NonZeroUsize;
2
3use super::{
4    InputCounts, InputLayout, InputRegion, LayoutRegions, MultiAirIndices, SELECTORS_PER_AIR,
5    StarkVarIndices,
6};
7use crate::{EXT_DEGREE, randomness};
8
9#[derive(Clone, Copy)]
10enum Alignment {
11    Unaligned = 1,
12    Word = 2,
13    DoubleWord = 4,
14    QuadWord = 8,
15}
16
17#[derive(Clone, Copy)]
18struct LayoutPolicy {
19    public_values: Alignment,
20    randomness: Alignment,
21    preprocessed: Alignment,
22    main: Alignment,
23    aux: Alignment,
24    quotient: Alignment,
25    aux_bus_boundary: Alignment,
26    stark_vars: Alignment,
27    end_align: Option<Alignment>,
28}
29
30/// Whether the layout includes the slots needed to combine multiple AIR instances.
31#[derive(Clone, Copy)]
32enum AirComposition {
33    Single,
34    Multi { air_count: NonZeroUsize },
35}
36
37impl AirComposition {
38    fn multi_air(air_count: usize) -> Self {
39        let air_count =
40            NonZeroUsize::new(air_count).expect("multi-AIR layout requires at least one AIR");
41        Self::Multi { air_count }
42    }
43
44    fn extra_stark_slots(self) -> usize {
45        match self {
46            Self::Single => 0,
47            Self::Multi { air_count } => 1 + air_count.get() * SELECTORS_PER_AIR,
48        }
49    }
50
51    fn multi_air_indices(self, stark_start: usize, base_slots: usize) -> Option<MultiAirIndices> {
52        match self {
53            Self::Single => None,
54            Self::Multi { air_count } => Some(MultiAirIndices {
55                air_count,
56                fold_beta: stark_start + base_slots,
57                selector_start: stark_start + base_slots + 1,
58            }),
59        }
60    }
61}
62
63impl LayoutPolicy {
64    fn native() -> Self {
65        Self {
66            public_values: Alignment::Unaligned,
67            randomness: Alignment::Unaligned,
68            preprocessed: Alignment::Unaligned,
69            main: Alignment::Unaligned,
70            aux: Alignment::Unaligned,
71            quotient: Alignment::Unaligned,
72            aux_bus_boundary: Alignment::Unaligned,
73            stark_vars: Alignment::Unaligned,
74            end_align: None,
75        }
76    }
77
78    fn masm() -> Self {
79        Self {
80            public_values: Alignment::QuadWord,
81            randomness: Alignment::Word,
82            preprocessed: Alignment::DoubleWord,
83            main: Alignment::DoubleWord,
84            aux: Alignment::DoubleWord,
85            quotient: Alignment::DoubleWord,
86            aux_bus_boundary: Alignment::Word,
87            stark_vars: Alignment::Word,
88            end_align: Some(Alignment::Word),
89        }
90    }
91}
92
93struct LayoutBuilder {
94    offset: usize,
95}
96
97impl LayoutBuilder {
98    fn new() -> Self {
99        Self { offset: 0 }
100    }
101
102    fn align(&mut self, alignment: Alignment) {
103        self.offset = self.offset.next_multiple_of(alignment as usize);
104    }
105
106    fn alloc(&mut self, width: usize, alignment: Alignment) -> InputRegion {
107        self.align(alignment);
108        let region = InputRegion { offset: self.offset, width };
109        self.offset += width;
110        region
111    }
112}
113
114impl InputLayout {
115    /// Build a native layout with no alignment padding.
116    pub fn new(counts: InputCounts) -> Self {
117        Self::build_with_policy(counts, LayoutPolicy::native(), AirComposition::Single)
118    }
119
120    /// Build the MASM-compatible layout.
121    pub fn new_masm(counts: InputCounts) -> Self {
122        Self::build_with_policy(counts, LayoutPolicy::masm(), AirComposition::Single)
123    }
124
125    /// Build a native layout for a multi-AIR relation.
126    pub fn new_multi_air(counts: InputCounts, num_airs: usize) -> Self {
127        Self::build_with_policy(counts, LayoutPolicy::native(), AirComposition::multi_air(num_airs))
128    }
129
130    /// Build the MASM-compatible layout for a multi-AIR relation.
131    pub fn new_masm_multi_air(counts: InputCounts, num_airs: usize) -> Self {
132        Self::build_with_policy(counts, LayoutPolicy::masm(), AirComposition::multi_air(num_airs))
133    }
134
135    fn build_with_policy(
136        counts: InputCounts,
137        policy: LayoutPolicy,
138        composition: AirComposition,
139    ) -> Self {
140        const NUM_RANDOMNESS_INPUTS: usize = 2;
141        assert_eq!(
142            counts.num_randomness, NUM_RANDOMNESS_INPUTS,
143            "ACE layouts require exactly alpha and beta randomness inputs"
144        );
145
146        // Every ACE input slot is an extension-field element. Slots 7-9 carry base-field values
147        // embedded as `(value, 0)`.
148        const NUM_STARK_VARS_BASE: usize = 10;
149        let num_stark_vars = NUM_STARK_VARS_BASE + composition.extra_stark_slots();
150
151        let mut builder = LayoutBuilder::new();
152
153        let public_values = builder.alloc(counts.num_public, policy.public_values);
154        let randomness = builder.alloc(NUM_RANDOMNESS_INPUTS, policy.randomness);
155        let (aux_rand_alpha, aux_rand_beta) = randomness::aux_rand_indices(randomness);
156        let preprocessed_curr = builder.alloc(counts.preprocessed_width, policy.preprocessed);
157        let main_curr = builder.alloc(counts.width, policy.main);
158        let aux_coord_width = counts.aux_width * EXT_DEGREE;
159        let aux_curr = builder.alloc(aux_coord_width, policy.aux);
160        let quotient_curr = builder.alloc(counts.num_quotient_chunks * EXT_DEGREE, policy.quotient);
161        let preprocessed_next = builder.alloc(counts.preprocessed_width, policy.preprocessed);
162        let main_next = builder.alloc(counts.width, policy.main);
163        let aux_next = builder.alloc(aux_coord_width, policy.aux);
164        let quotient_next = builder.alloc(counts.num_quotient_chunks * EXT_DEGREE, policy.quotient);
165        let aux_bus_boundary = builder.alloc(counts.num_aux_boundary, policy.aux_bus_boundary);
166        let stark_vars = builder.alloc(num_stark_vars, policy.stark_vars);
167
168        let b = stark_vars.offset;
169        let alpha = b;
170        let z_pow_n = b + 1;
171        let z_k = b + 2;
172        let is_first = b + 3;
173        let is_last = b + 4;
174        let is_transition = b + 5;
175        let reserved = b + 6;
176        let weight0 = b + 7;
177        let f = b + 8;
178        let s0 = b + 9;
179        let multi_air = composition.multi_air_indices(b, NUM_STARK_VARS_BASE);
180
181        if let Some(end_align) = policy.end_align {
182            builder.align(end_align);
183        }
184
185        Self {
186            regions: LayoutRegions {
187                public_values,
188                randomness,
189                preprocessed_curr,
190                main_curr,
191                aux_curr,
192                quotient_curr,
193                preprocessed_next,
194                main_next,
195                aux_next,
196                quotient_next,
197                aux_bus_boundary,
198                stark_vars,
199            },
200            aux_rand_alpha,
201            aux_rand_beta,
202            stark: StarkVarIndices {
203                alpha,
204                z_pow_n,
205                z_k,
206                is_first,
207                is_last,
208                is_transition,
209                reserved,
210                weight0,
211                f,
212                s0,
213                multi_air,
214            },
215            total_inputs: builder.offset,
216            counts,
217        }
218    }
219}
220
221#[cfg(test)]
222mod tests {
223    use super::super::{InputCounts, InputKey, InputLayout};
224
225    fn test_counts() -> InputCounts {
226        InputCounts {
227            preprocessed_width: 0,
228            width: 1,
229            aux_width: 1,
230            num_aux_boundary: 3,
231            num_public: 8,
232            num_randomness: 2,
233            num_quotient_chunks: 1,
234        }
235    }
236
237    #[test]
238    fn multi_air_layout_indexes_air_slots() {
239        let layout = InputLayout::new_masm_multi_air(test_counts(), 3);
240
241        let beta = layout.index(InputKey::MultiAirFoldBeta).unwrap();
242
243        let first0 = layout.index(InputKey::IsFirstAir(0)).unwrap();
244        assert_eq!(first0, beta + 1);
245        assert_eq!(layout.index(InputKey::IsLastAir(0)), Some(first0 + 1));
246        assert_eq!(layout.index(InputKey::IsTransitionAir(0)), Some(first0 + 2));
247        assert_eq!(layout.index(InputKey::IsFirstAir(1)), Some(first0 + 3));
248        assert_eq!(layout.index(InputKey::IsFirstAir(2)), Some(first0 + 6));
249        assert_eq!(layout.index(InputKey::IsFirstAir(3)), None);
250    }
251
252    #[test]
253    #[should_panic(expected = "multi-AIR layout requires at least one AIR")]
254    fn multi_air_layout_rejects_zero_airs() {
255        let _ = InputLayout::new_masm_multi_air(test_counts(), 0);
256    }
257}
miden_ace_codegen/layout/policy.rs

miden_ace_codegen/layout/
policy.rs
