miden_air/constraints/
columns.rs

1//! Column layout types for the main and auxiliary execution traces.
2//!
3//! These `#[repr(C)]` structs provide typed, named access to trace columns.
4//! They are borrowed zero-copy from raw `[T; WIDTH]` slices and are used
5//! exclusively by constraint code. They are independent of trace storage
6//! (`MainTrace`, `TraceStorage`, etc.).
7
8use core::{
9    borrow::{Borrow, BorrowMut},
10    mem::size_of,
11};
12
13use super::{
14    chiplets::columns::{
15        AceCols, AceEvalCols, AceReadCols, BitwiseCols, ControllerCols, KernelRomCols, MemoryCols,
16        PermutationCols,
17    },
18    decoder::columns::DecoderCols,
19    range::columns::RangeCols,
20    stack::columns::StackCols,
21    system::columns::SystemCols,
22};
23use crate::trace::{CHIPLETS_WIDTH, TRACE_WIDTH};
24
25// CORE TRACE COLUMN STRUCT
26// ================================================================================================
27
28/// Column layout of the core execution trace.
29///
30/// `CoreCols` covers the system, decoder, stack, and range-check segments — the columns owned
31/// by `CoreAir`. It is also the layout of the leading `NUM_CORE_COLS` columns of the unified
32/// `TRACE_WIDTH`-wide main trace, so it can be borrowed from either a per-AIR
33/// `[T; NUM_CORE_COLS]` slice or the prefix of a `[T; TRACE_WIDTH]` row via
34/// `Borrow<CoreCols<T>>`.
35#[repr(C)]
36#[derive(Debug, Clone, Default)]
37pub struct CoreCols<T> {
38    pub system: SystemCols<T>,
39    pub decoder: DecoderCols<T>,
40    pub stack: StackCols<T>,
41    pub range: RangeCols<T>,
42}
43
44/// Number of columns in the core trace (51), derived from the struct layout.
45pub const NUM_CORE_COLS: usize = size_of::<CoreCols<u8>>();
46
47impl<T> Borrow<CoreCols<T>> for [T] {
48    fn borrow(&self) -> &CoreCols<T> {
49        debug_assert_eq!(self.len(), NUM_CORE_COLS);
50        let (prefix, shorts, suffix) = unsafe { self.align_to::<CoreCols<T>>() };
51        debug_assert!(prefix.is_empty() && suffix.is_empty() && shorts.len() == 1);
52        &shorts[0]
53    }
54}
55
56impl<T> BorrowMut<CoreCols<T>> for [T] {
57    fn borrow_mut(&mut self) -> &mut CoreCols<T> {
58        debug_assert_eq!(self.len(), NUM_CORE_COLS);
59        let (prefix, shorts, suffix) = unsafe { self.align_to_mut::<CoreCols<T>>() };
60        debug_assert!(prefix.is_empty() && suffix.is_empty() && shorts.len() == 1);
61        &mut shorts[0]
62    }
63}
64
65impl<T> CoreCols<T> {
66    /// Returns the column layout as a flat slice of length NUM_CORE_COLS, in column-index
67    /// order. Useful for column-index → field lookups (e.g. a `CoreCols<&str>` name table).
68    pub fn as_slice(&self) -> &[T] {
69        let ptr = self as *const Self as *const T;
70        unsafe { core::slice::from_raw_parts(ptr, NUM_CORE_COLS) }
71    }
72}
73
74// CHIPLETS TRACE COLUMN STRUCT
75// ================================================================================================
76
77/// Column layout of the chiplets execution trace.
78///
79/// `ChipletCols` covers the `s_00` and `s_01` chiplet selectors, `chip_clk`, and the 19 shared
80/// chiplet data columns — the columns owned by `ChipletsAir`. It is also the layout of the
81/// trailing `NUM_CHIPLETS_COLS` columns of the unified main trace, so it can be borrowed from
82/// either a per-AIR `[T; NUM_CHIPLETS_COLS]` slice or the suffix of a `[T; TRACE_WIDTH]` row via
83/// `Borrow<ChipletCols<T>>`.
84#[repr(C)]
85#[derive(Clone, Debug)]
86pub struct ChipletCols<T> {
87    /// Permutation segment selector: consumed by `build_chiplet_selectors`.
88    pub s_00: T,
89    /// Controller segment selector: consumed by `build_chiplet_selectors`.
90    pub s_01: T,
91    /// Chiplet-trace row counter: starts at 1 on the first row, increments by 1 each row.
92    pub chip_clk: T,
93    pub(crate) chiplets: [T; CHIPLETS_WIDTH - 3],
94}
95
96/// Number of columns in the chiplets trace (22), derived from the struct layout.
97pub const NUM_CHIPLETS_COLS: usize = size_of::<ChipletCols<u8>>();
98
99impl<T> ChipletCols<T> {
100    /// Returns the 6 chiplet selector columns `[s_00, s_01, s1, s2, s3, s4]`.
101    ///
102    /// `s_00` and `s_01` are the two physical selectors for the permutation and controller
103    /// sub-chiplets. `s1..s4` subdivide the remaining chiplets under the virtual
104    /// `s0 = 1 - (s_00 + s_01)`.
105    pub fn chiplet_selectors(&self) -> [T; 6]
106    where
107        T: Copy,
108    {
109        [
110            self.s_00,
111            self.s_01,
112            self.chiplets[0],
113            self.chiplets[1],
114            self.chiplets[2],
115            self.chiplets[3],
116        ]
117    }
118
119    /// Returns a typed borrow of the bitwise chiplet columns (chiplets\[1..14\]).
120    pub fn bitwise(&self) -> &BitwiseCols<T> {
121        self.chiplets[1..14].borrow()
122    }
123
124    /// Returns a typed borrow of the memory chiplet columns (chiplets\[2..17\]).
125    pub fn memory(&self) -> &MemoryCols<T> {
126        self.chiplets[2..17].borrow()
127    }
128
129    /// Returns the lower 16-bit limb of the memory word address (chiplets\[17\]).
130    ///
131    /// Range-check auxiliary column populated by the trace builder for the lookup-bus
132    /// emitter; not part of [`MemoryCols`] because the memory AIR's own transition
133    /// constraints don't act on it.
134    pub fn memory_word_addr_lo(&self) -> T
135    where
136        T: Copy,
137    {
138        self.chiplets[17]
139    }
140
141    /// Returns the upper 16-bit limb of the memory word address (chiplets\[18\]).
142    ///
143    /// See [`Self::memory_word_addr_lo`] for the same caveat about the range-check
144    /// auxiliary columns living outside [`MemoryCols`].
145    pub fn memory_word_addr_hi(&self) -> T
146    where
147        T: Copy,
148    {
149        self.chiplets[18]
150    }
151
152    /// Returns a typed borrow of the ACE chiplet columns (chiplets\[3..19\]).
153    pub fn ace(&self) -> &AceCols<T> {
154        self.chiplets[3..].borrow()
155    }
156
157    /// Returns a typed borrow of the kernel ROM chiplet columns (chiplets\[4..9\]).
158    pub fn kernel_rom(&self) -> &KernelRomCols<T> {
159        self.chiplets[4..9].borrow()
160    }
161
162    /// Returns a typed borrow of the permutation sub-chiplet columns (chiplets\[0..19\]).
163    pub fn permutation(&self) -> &PermutationCols<T> {
164        self.chiplets[..].borrow()
165    }
166
167    /// Returns a typed borrow of the controller sub-chiplet columns (chiplets\[0..19\]).
168    pub fn controller(&self) -> &ControllerCols<T> {
169        self.chiplets[..].borrow()
170    }
171}
172
173impl<T> Borrow<ChipletCols<T>> for [T] {
174    fn borrow(&self) -> &ChipletCols<T> {
175        debug_assert_eq!(self.len(), NUM_CHIPLETS_COLS);
176        let (prefix, shorts, suffix) = unsafe { self.align_to::<ChipletCols<T>>() };
177        debug_assert!(prefix.is_empty() && suffix.is_empty() && shorts.len() == 1);
178        &shorts[0]
179    }
180}
181
182impl<T> BorrowMut<ChipletCols<T>> for [T] {
183    fn borrow_mut(&mut self) -> &mut ChipletCols<T> {
184        debug_assert_eq!(self.len(), NUM_CHIPLETS_COLS);
185        let (prefix, shorts, suffix) = unsafe { self.align_to_mut::<ChipletCols<T>>() };
186        debug_assert!(prefix.is_empty() && suffix.is_empty() && shorts.len() == 1);
187        &mut shorts[0]
188    }
189}
190
191// Compile-time invariant: the two halves cover the full main trace exactly.
192const _: () = assert!(NUM_CORE_COLS + NUM_CHIPLETS_COLS == TRACE_WIDTH);
193
194// CONST HELPERS
195// ================================================================================================
196
197/// Generates an array `[0, 1, 2, ..., N-1]` at compile time.
198pub const fn indices_arr<const N: usize>() -> [usize; N] {
199    let mut arr = [0; N];
200    let mut i = 0;
201    while i < N {
202        arr[i] = i;
203        i += 1;
204    }
205    arr
206}
207
208// COLUMN COUNTS
209// ================================================================================================
210//
211// The auxiliary trace is the LogUp lookup-argument segment built per-AIR by `CoreAir`'s
212// and `ChipletsAir`'s `build_aux_trace` (see `air/src/constraints/lookup/`): 4 Core
213// columns + 3 Chiplets columns.
214
215pub const NUM_SYSTEM_COLS: usize = size_of::<SystemCols<u8>>();
216pub const NUM_DECODER_COLS: usize = size_of::<DecoderCols<u8>>();
217pub const NUM_STACK_COLS: usize = size_of::<StackCols<u8>>();
218pub const NUM_RANGE_COLS: usize = size_of::<RangeCols<u8>>();
219pub const NUM_BITWISE_COLS: usize = size_of::<BitwiseCols<u8>>();
220pub const NUM_MEMORY_COLS: usize = size_of::<MemoryCols<u8>>();
221pub const NUM_ACE_COLS: usize = size_of::<AceCols<u8>>();
222pub const NUM_ACE_READ_COLS: usize = size_of::<AceReadCols<u8>>();
223pub const NUM_ACE_EVAL_COLS: usize = size_of::<AceEvalCols<u8>>();
224pub const NUM_KERNEL_ROM_COLS: usize = size_of::<KernelRomCols<u8>>();
225
226const _: () = assert!(NUM_SYSTEM_COLS == 6);
227const _: () = assert!(NUM_DECODER_COLS == 24);
228const _: () = assert!(NUM_STACK_COLS == 19);
229const _: () = assert!(NUM_RANGE_COLS == 2);
230const _: () = assert!(NUM_BITWISE_COLS == 13);
231const _: () = assert!(NUM_MEMORY_COLS == 15);
232const _: () = assert!(NUM_ACE_COLS == 16);
233const _: () = assert!(NUM_ACE_READ_COLS == 4);
234const _: () = assert!(NUM_ACE_EVAL_COLS == 4);
235const _: () = assert!(NUM_KERNEL_ROM_COLS == 5);
236
237// TESTS
238// ================================================================================================
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243    use crate::trace::{DECODER_TRACE_WIDTH, STACK_TRACE_WIDTH, SYS_TRACE_WIDTH};
244
245    /// Per-AIR index maps used only by the column-layout tests below. Each field holds its
246    /// column index inside its own AIR. `CORE_COL_MAP` lines up with unified-trace offsets
247    /// since `CoreCols` sits at offset 0; `CHIPLET_COL_MAP` is 0-based within `ChipletCols`.
248    const CORE_COL_MAP: CoreCols<usize> = unsafe {
249        core::mem::transmute::<[usize; NUM_CORE_COLS], CoreCols<usize>>(
250            indices_arr::<NUM_CORE_COLS>(),
251        )
252    };
253
254    const CHIPLET_COL_MAP: ChipletCols<usize> = unsafe {
255        core::mem::transmute::<[usize; NUM_CHIPLETS_COLS], ChipletCols<usize>>(indices_arr::<
256            NUM_CHIPLETS_COLS,
257        >())
258    };
259
260    /// Column offset of the decoder section within the unified main trace.
261    const DECODER_OFFSET: usize = SYS_TRACE_WIDTH;
262    /// Column offset of the stack section within the unified main trace.
263    const STACK_OFFSET: usize = SYS_TRACE_WIDTH + DECODER_TRACE_WIDTH;
264    /// Column offset of the range-check section within the unified main trace.
265    const RANGE_OFFSET: usize = STACK_OFFSET + STACK_TRACE_WIDTH;
266
267    // --- Core trace column map vs offset constants -------------------------------------------
268    //
269    // `CoreCols` starts at offset 0 of the unified main trace, so its per-AIR indices match
270    // the unified trace offsets one-for-one.
271
272    #[test]
273    fn col_map_system() {
274        assert_eq!(CORE_COL_MAP.system.clk, 0);
275        assert_eq!(CORE_COL_MAP.system.ctx, 1);
276        assert_eq!(CORE_COL_MAP.system.fn_hash[0], 2);
277        assert_eq!(CORE_COL_MAP.system.fn_hash[3], 5);
278    }
279
280    #[test]
281    fn col_map_decoder() {
282        assert_eq!(CORE_COL_MAP.decoder.addr, DECODER_OFFSET);
283        assert_eq!(CORE_COL_MAP.decoder.op_bits[0], DECODER_OFFSET + 1);
284        assert_eq!(CORE_COL_MAP.decoder.op_bits[6], DECODER_OFFSET + 7);
285        assert_eq!(CORE_COL_MAP.decoder.hasher_state[0], DECODER_OFFSET + 8);
286        assert_eq!(CORE_COL_MAP.decoder.in_span, DECODER_OFFSET + 16);
287        assert_eq!(CORE_COL_MAP.decoder.group_count, DECODER_OFFSET + 17);
288        assert_eq!(CORE_COL_MAP.decoder.op_index, DECODER_OFFSET + 18);
289        assert_eq!(CORE_COL_MAP.decoder.batch_flags[0], DECODER_OFFSET + 19);
290        assert_eq!(CORE_COL_MAP.decoder.extra[0], DECODER_OFFSET + 22);
291    }
292
293    #[test]
294    fn col_map_stack() {
295        assert_eq!(CORE_COL_MAP.stack.top[0], STACK_OFFSET);
296        assert_eq!(CORE_COL_MAP.stack.top[15], STACK_OFFSET + 15);
297        assert_eq!(CORE_COL_MAP.stack.b0, STACK_OFFSET + 16);
298        assert_eq!(CORE_COL_MAP.stack.b1, STACK_OFFSET + 17);
299        assert_eq!(CORE_COL_MAP.stack.h0, STACK_OFFSET + 18);
300    }
301
302    #[test]
303    fn col_map_range() {
304        assert_eq!(CORE_COL_MAP.range.multiplicity, RANGE_OFFSET);
305        assert_eq!(CORE_COL_MAP.range.value, RANGE_OFFSET + 1);
306    }
307
308    // --- Chiplet trace column map -------------------------------------------------------------
309    //
310    // `CHIPLET_COL_MAP` is 0-based within `ChipletCols`.
311
312    #[test]
313    fn col_map_chiplets() {
314        assert_eq!(CHIPLET_COL_MAP.s_00, 0);
315        assert_eq!(CHIPLET_COL_MAP.s_01, 1);
316        assert_eq!(CHIPLET_COL_MAP.chip_clk, 2);
317        assert_eq!(CHIPLET_COL_MAP.chiplets[0], 3);
318        assert_eq!(CHIPLET_COL_MAP.chiplets[18], 21);
319    }
320
321    // --- Multi-AIR split: CoreCols + ChipletCols widths ---------------------------------------
322
323    /// `NUM_CORE_COLS` matches the sum of the segment widths it covers.
324    #[test]
325    fn core_cols_width() {
326        assert_eq!(
327            NUM_CORE_COLS,
328            NUM_SYSTEM_COLS + NUM_DECODER_COLS + NUM_STACK_COLS + NUM_RANGE_COLS,
329        );
330    }
331
332    /// `NUM_CHIPLETS_COLS` matches the chiplets segment width.
333    #[test]
334    fn chiplet_cols_width() {
335        assert_eq!(NUM_CHIPLETS_COLS, CHIPLETS_WIDTH);
336    }
337
338    // --- Layout snapshots ---------------------------------------------------------------------
339    //
340    // These pin the resolved column-index maps of each `Cols<T>` view to `.snap` files,
341    // so any layout change surfaces as a snapshot diff in PR review.
342    // Regenerate with `cargo insta review` (or `INSTA_UPDATE=auto cargo test -p miden-air`).
343
344    /// Builds a `$cols<usize>` index map by reinterpreting `[0, 1, …, N-1]` through the
345    /// struct's `#[repr(C)]` layout, where `N = size_of::<$cols<u8>>()`.
346    macro_rules! col_map {
347        ($cols:ident) => {{
348            const N: usize = core::mem::size_of::<$cols<u8>>();
349            const M: $cols<usize> =
350                unsafe { core::mem::transmute::<[usize; N], $cols<usize>>(indices_arr::<N>()) };
351            M
352        }};
353    }
354
355    #[test]
356    fn core_col_map_layout() {
357        insta::assert_debug_snapshot!(CORE_COL_MAP);
358    }
359
360    #[test]
361    fn chiplet_col_map_layout() {
362        insta::assert_debug_snapshot!(CHIPLET_COL_MAP);
363    }
364
365    #[test]
366    fn bitwise_col_map_layout() {
367        insta::assert_debug_snapshot!(col_map!(BitwiseCols));
368    }
369
370    #[test]
371    fn memory_col_map_layout() {
372        insta::assert_debug_snapshot!(col_map!(MemoryCols));
373    }
374
375    #[test]
376    fn ace_col_map_layout() {
377        insta::assert_debug_snapshot!(col_map!(AceCols));
378    }
379
380    #[test]
381    fn ace_read_col_map_layout() {
382        insta::assert_debug_snapshot!(col_map!(AceReadCols));
383    }
384
385    #[test]
386    fn ace_eval_col_map_layout() {
387        insta::assert_debug_snapshot!(col_map!(AceEvalCols));
388    }
389
390    #[test]
391    fn kernel_rom_col_map_layout() {
392        insta::assert_debug_snapshot!(col_map!(KernelRomCols));
393    }
394
395    #[test]
396    fn hasher_controller_col_map_layout() {
397        insta::assert_debug_snapshot!(col_map!(ControllerCols));
398    }
399
400    #[test]
401    fn hasher_permutation_col_map_layout() {
402        insta::assert_debug_snapshot!(col_map!(PermutationCols));
403    }
404}
miden_air/constraints/columns.rs

miden_air/constraints/
columns.rs
