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 variable-length public input (VLPI) reduction slots (in EF elements).
28    /// This is derived from `AceConfig::num_vlpi_groups` by the layout policy:
29    /// MASM expands each group to 2 EF slots (word-aligned); Native uses 1 per group.
30    pub num_vlpi: usize,
31    /// Number of randomness challenges used by the AIR.
32    pub num_randomness: usize,
33    /// Number of periodic columns.
34    pub num_periodic: 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-length public values.
43    pub public_values: InputRegion,
44    /// Region containing variable-length public input reductions.
45    pub vlpi_reductions: InputRegion,
46    /// Region containing randomness inputs (alpha, beta).
47    pub randomness: 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    /// Main trace OOD values at `g * zeta`.
55    pub main_next: InputRegion,
56    /// Aux trace OOD coordinates at `g * zeta`.
57    pub aux_next: InputRegion,
58    /// Quotient chunk OOD coordinates at `g * zeta`.
59    pub quotient_next: InputRegion,
60    /// Aux bus boundary values.
61    pub aux_bus_boundary: InputRegion,
62    /// Stark variables (selectors, powers, weights).
63    pub stark_vars: InputRegion,
64}
65
66/// Indexes of canonical verifier scalars inside the stark-vars block.
67///
68/// Every slot in the ACE input array is an extension-field (EF) element --
69/// the circuit operates entirely in the extension field. However, some of
70/// these scalars are inherently base-field values that the MASM verifier
71/// stores as `(val, 0)` in the EF slot.
72///
73/// See the module documentation on [`super::super::dag::lower`] for how each
74/// variable enters the verifier expression.
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
76pub(crate) struct StarkVarIndices {
77    // -- Extension-field values (slots 0-5) --
78    /// Composition challenge `alpha` for folding constraints.
79    pub alpha: usize,
80    /// `zeta^N` where N is the trace length.
81    pub z_pow_n: usize,
82    /// `zeta^(N / max_cycle_len)` for periodic column evaluation.
83    pub z_k: usize,
84    /// Precomputed first-row selector: `(z^N - 1) / (z - 1)`.
85    pub is_first: usize,
86    /// Precomputed last-row selector: `(z^N - 1) / (z - g^{-1})`.
87    pub is_last: usize,
88    /// Precomputed transition selector: `z - g^{-1}`.
89    pub is_transition: usize,
90    /// Batching challenge `gamma` for reduced_aux_values.
91    pub gamma: usize,
92
93    // -- Base-field values stored as (val, 0) in EF slots --
94    /// First barycentric weight `1 / (k * s0^{k-1})`.
95    pub weight0: usize,
96    /// `f = h^N` (chunk shift ratio between cosets).
97    pub f: usize,
98    /// `s0 = offset^N` (first chunk shift).
99    pub s0: usize,
100}
101
102/// ACE input layout for Plonky3-based verifier logic.
103///
104/// This describes the exact ordering and alignment of inputs consumed by the
105/// ACE chiplet (READ section).
106#[derive(Debug, Clone)]
107pub struct InputLayout {
108    /// Grouped regions for the ACE input layout.
109    pub(crate) regions: LayoutRegions,
110    /// Input index for aux randomness alpha.
111    pub(crate) aux_rand_alpha: usize,
112    /// Input index for aux randomness beta.
113    pub(crate) aux_rand_beta: usize,
114    /// Stride between logical VLPI groups (2 for MASM word-aligned, 1 for native).
115    pub(crate) vlpi_stride: usize,
116    /// Indexes into the stark-vars region.
117    pub(crate) stark: StarkVarIndices,
118    /// Total number of inputs (length of the READ section).
119    pub total_inputs: usize,
120    /// Counts used to derive the layout.
121    pub counts: InputCounts,
122}
123
124impl InputLayout {
125    pub(crate) fn mapper(&self) -> super::InputKeyMapper<'_> {
126        super::InputKeyMapper { layout: self }
127    }
128
129    /// Map a logical `InputKey` into the flat input index, if present.
130    pub fn index(&self, key: InputKey) -> Option<usize> {
131        self.mapper().index_of(key)
132    }
133
134    /// Validate internal invariants for this layout (region sizes, key ranges, randomness inputs).
135    pub(crate) fn validate(&self) {
136        let mut max_end = 0usize;
137        for region in [
138            self.regions.public_values,
139            self.regions.vlpi_reductions,
140            self.regions.randomness,
141            self.regions.main_curr,
142            self.regions.aux_curr,
143            self.regions.quotient_curr,
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        let aux_coord_width = self.counts.aux_width * EXT_DEGREE;
156        assert_eq!(self.regions.aux_curr.width, aux_coord_width, "aux_curr width mismatch");
157        assert_eq!(self.regions.aux_next.width, aux_coord_width, "aux_next width mismatch");
158
159        let quotient_width = self.counts.num_quotient_chunks * EXT_DEGREE;
160        assert_eq!(
161            self.regions.quotient_curr.width, quotient_width,
162            "quotient_curr width mismatch"
163        );
164        assert_eq!(
165            self.regions.quotient_next.width, quotient_width,
166            "quotient_next width mismatch"
167        );
168        assert_eq!(
169            self.regions.aux_bus_boundary.width, self.counts.aux_width,
170            "aux bus boundary width mismatch"
171        );
172
173        let stark_start = self.regions.stark_vars.offset;
174        let stark_end = stark_start + self.regions.stark_vars.width;
175        let check = |name: &str, idx: usize| {
176            assert!(idx >= stark_start && idx < stark_end, "stark var {name} out of range");
177        };
178        // Extension-field slots.
179        check("alpha", self.stark.alpha);
180        check("z_pow_n", self.stark.z_pow_n);
181        check("z_k", self.stark.z_k);
182        check("is_first", self.stark.is_first);
183        check("is_last", self.stark.is_last);
184        check("is_transition", self.stark.is_transition);
185        check("gamma", self.stark.gamma);
186        // Base-field slots (stored as (val, 0) in the EF slot).
187        check("weight0", self.stark.weight0);
188        check("f", self.stark.f);
189        check("s0", self.stark.s0);
190
191        let rand_start = self.regions.randomness.offset;
192        let rand_end = rand_start + self.regions.randomness.width;
193        assert!(
194            self.aux_rand_alpha >= rand_start && self.aux_rand_alpha < rand_end,
195            "aux_rand_alpha out of randomness region"
196        );
197        assert!(
198            self.aux_rand_beta >= rand_start && self.aux_rand_beta < rand_end,
199            "aux_rand_beta out of randomness region"
200        );
201    }
202}