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 public inputs.
26    pub num_public: usize,
27    /// Number of randomness challenges used by the AIR.
28    pub num_randomness: usize,
29    /// Number of periodic columns.
30    pub num_periodic: usize,
31    /// Number of auxiliary (stark var) inputs reserved.
32    pub num_aux_inputs: 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
64pub(crate) struct StarkVarIndices {
65    /// Index of `zeta` in the stark-vars block.
66    pub z: usize,
67    /// Index of the composition challenge `alpha`.
68    pub alpha: usize,
69    /// Index of `g^{-1}`.
70    pub g_inv: usize,
71    /// Index of `zeta^N`.
72    pub z_pow_n: usize,
73    /// Index of `g^{-2}`.
74    pub g_inv2: usize,
75    /// Index of `z_k`.
76    pub z_k: usize,
77    /// Index of `weight0`.
78    pub weight0: usize,
79    /// Index of `g`.
80    pub g: usize,
81    /// Index of `s0`.
82    pub s0: usize,
83    /// Index of `1 / (zeta - g^{-1})`.
84    pub inv_z_minus_g_inv: usize,
85    /// Index of `1 / (zeta - 1)`.
86    pub inv_z_minus_one: usize,
87    /// Index of `1 / (zeta^N - 1)`.
88    pub inv_vanishing: usize,
89}
90
91/// ACE input layout for Plonky3-based verifier logic.
92///
93/// This describes the exact ordering and alignment of inputs consumed by the
94/// ACE chiplet (READ section).
95#[derive(Debug, Clone)]
96pub struct InputLayout {
97    /// Grouped regions for the ACE input layout.
98    pub(crate) regions: LayoutRegions,
99    /// Input index for aux randomness alpha.
100    pub(crate) aux_rand_alpha: usize,
101    /// Input index for aux randomness beta.
102    pub(crate) aux_rand_beta: usize,
103    /// Indexes into the stark-vars region.
104    pub(crate) stark: StarkVarIndices,
105    /// Total number of inputs (length of the READ section).
106    pub total_inputs: usize,
107    /// Counts used to derive the layout.
108    pub counts: InputCounts,
109}
110
111impl InputLayout {
112    pub(crate) fn mapper(&self) -> super::InputKeyMapper<'_> {
113        super::InputKeyMapper { layout: self }
114    }
115
116    /// Map a logical `InputKey` into the flat input index, if present.
117    pub fn index(&self, key: InputKey) -> Option<usize> {
118        self.mapper().index_of(key)
119    }
120
121    /// Validate internal invariants for this layout (region sizes, key ranges, randomness inputs).
122    pub(crate) fn validate(&self) {
123        let mut max_end = 0usize;
124        for region in [
125            self.regions.public_values,
126            self.regions.randomness,
127            self.regions.main_curr,
128            self.regions.aux_curr,
129            self.regions.quotient_curr,
130            self.regions.main_next,
131            self.regions.aux_next,
132            self.regions.quotient_next,
133            self.regions.aux_bus_boundary,
134            self.regions.stark_vars,
135        ] {
136            max_end = max_end.max(region.offset.saturating_add(region.width));
137        }
138
139        assert!(max_end <= self.total_inputs, "regions exceed total_inputs");
140
141        let aux_coord_width = self.counts.aux_width * EXT_DEGREE;
142        assert_eq!(self.regions.aux_curr.width, aux_coord_width, "aux_curr width mismatch");
143        assert_eq!(self.regions.aux_next.width, aux_coord_width, "aux_next width mismatch");
144
145        let quotient_width = self.counts.num_quotient_chunks * EXT_DEGREE;
146        assert_eq!(
147            self.regions.quotient_curr.width, quotient_width,
148            "quotient_curr width mismatch"
149        );
150        assert_eq!(
151            self.regions.quotient_next.width, quotient_width,
152            "quotient_next width mismatch"
153        );
154        assert_eq!(
155            self.regions.aux_bus_boundary.width, self.counts.aux_width,
156            "aux bus boundary width mismatch"
157        );
158
159        let stark_end = self.regions.stark_vars.offset + self.regions.stark_vars.width;
160        for (name, idx) in [
161            ("z", self.stark.z),
162            ("alpha", self.stark.alpha),
163            ("g_inv", self.stark.g_inv),
164            ("z_pow_n", self.stark.z_pow_n),
165            ("g_inv2", self.stark.g_inv2),
166            ("z_k", self.stark.z_k),
167            ("weight0", self.stark.weight0),
168            ("g", self.stark.g),
169            ("s0", self.stark.s0),
170            ("inv_z_minus_g_inv", self.stark.inv_z_minus_g_inv),
171            ("inv_z_minus_one", self.stark.inv_z_minus_one),
172            ("inv_vanishing", self.stark.inv_vanishing),
173        ] {
174            assert!(
175                idx >= self.regions.stark_vars.offset && idx < stark_end,
176                "stark var {name} out of range"
177            );
178        }
179
180        let rand_start = self.regions.randomness.offset;
181        let rand_end = rand_start + self.regions.randomness.width;
182        assert!(
183            self.aux_rand_alpha >= rand_start && self.aux_rand_alpha < rand_end,
184            "aux_rand_alpha out of randomness region"
185        );
186        assert!(
187            self.aux_rand_beta >= rand_start && self.aux_rand_beta < rand_end,
188            "aux_rand_beta out of randomness region"
189        );
190    }
191}