Skip to main content

hekate_program/
schema.rs

1// SPDX-FileCopyrightText: 2026 Andrei Kochergin <andrei@oumuamua.dev>
2// SPDX-FileCopyrightText: 2026 Oumuamua Labs <info@oumuamua.dev>
3// SPDX-License-Identifier: AGPL-3.0-only
4
5//! Typed Column Schema.
6//!
7//! `define_columns!` generates index constants,
8//! `NUM_COLUMNS`, and `build_layout()` from a
9//! single definition. One source of truth —
10//! indices and layout can't drift from each other.
11//!
12//! # Scalar columns
13//!
14//! `NAME: Type` → `const NAME: usize = <offset>`.
15//!
16//! # Array columns
17//!
18//! `NAME: [Type; N]` → `const NAME: usize = <base>`.
19//! Offset advances by N.
20//!
21//! # Example
22//!
23//! ```ignore
24//! define_columns! {
25//!     pub FibColumns {
26//!         A: B32,
27//!         B: B32,
28//!         Q: Bit,
29//!     }
30//! }
31//!
32//! // FibColumns::A == 0
33//! // FibColumns::B == 1
34//! // FibColumns::Q == 2
35//! // FibColumns::NUM_COLUMNS == 3
36//! // FibColumns::build_layout() → [B32, B32, Bit]
37//! ```
38
39use alloc::vec::Vec;
40
41pub use alloc::vec::Vec as SchemaVec;
42pub use hekate_core::trace::ColumnType;
43
44/// Define a typed column schema.
45///
46/// Generates index constants, `NUM_COLUMNS`,
47/// and `build_layout() -> Vec<ColumnType>`.
48#[macro_export]
49macro_rules! define_columns {
50    (
51        $(#[$meta:meta])*
52        $vis:vis $name:ident {
53            $( $field:ident : $kind:tt ),* $(,)?
54        }
55    ) => {
56        $(#[$meta])*
57        #[derive(Clone, Copy, Debug)]
58        $vis struct $name;
59
60        impl $name {
61            $crate::define_columns!(@consts 0usize; $( $field : $kind, )*);
62
63            /// Column layout derived
64            /// from the schema definition.
65            #[allow(unused_mut)]
66            pub fn build_layout() -> $crate::schema::SchemaVec<$crate::schema::ColumnType> {
67                let mut v = $crate::schema::SchemaVec::with_capacity(Self::NUM_COLUMNS);
68                $( $crate::define_columns!(@push v, $kind); )*
69
70                v
71            }
72        }
73    };
74
75    (@consts $offset:expr;) => {
76        pub const NUM_COLUMNS: usize = $offset;
77    };
78
79    // Array:
80    // NAME: [Type; N] — base index, offset += N
81    (@consts $offset:expr; $field:ident : [$ty:ident; $n:expr], $( $rest:tt )*) => {
82        #[allow(dead_code)]
83        pub const $field: usize = $offset;
84        $crate::define_columns!(@consts $offset + $n; $( $rest )*);
85    };
86
87    // Scalar:
88    // NAME: Type — index, offset += 1
89    (@consts $offset:expr; $field:ident : $ty:ident, $( $rest:tt )*) => {
90        #[allow(dead_code)]
91        pub const $field: usize = $offset;
92        $crate::define_columns!(@consts $offset + 1usize; $( $rest )*);
93    };
94
95    // Layout push:
96    // scalar
97    (@push $v:ident, $ty:ident) => {
98        $v.push($crate::schema::ColumnType::$ty);
99    };
100
101    // Layout push: array
102    (@push $v:ident, [$ty:ident; $n:expr]) => {
103        $v.extend(::core::iter::repeat($crate::schema::ColumnType::$ty).take($n));
104    };
105}
106
107/// Cumulative column offsets from chiplet widths.
108///
109/// `offsets[i]` = starting column index for chiplet i.
110pub fn chiplet_offsets(column_counts: &[usize]) -> Vec<usize> {
111    let mut offsets = Vec::with_capacity(column_counts.len());
112    let mut acc = 0;
113
114    for &count in column_counts {
115        offsets.push(acc);
116        acc += count;
117    }
118
119    offsets
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    define_columns! {
127        pub TestScalar {
128            A: B32,
129            B: B32,
130            Q: Bit,
131        }
132    }
133
134    define_columns! {
135        pub TestMixed {
136            ADDR_B0: B32,
137            ADDR_B1: B32,
138            ADDR_B2: B32,
139            ADDR_B3: B32,
140            IS_WRITE: Bit,
141            SELECTOR: Bit,
142            AUX_INV: B128,
143            DIFF_BYTE_IDX: [Bit; 8],
144            DIFF_BIT_IDX: [Bit; 8],
145            A_BITS: [Bit; 8],
146        }
147    }
148
149    define_columns! {
150        pub TestLargeArrays {
151            STATE_BITS: [Bit; 1600],
152            RC_BITS: [Bit; 64],
153            LANES: [B64; 25],
154            S_ROUND: Bit,
155            S_IN_OUT: Bit,
156        }
157    }
158
159    define_columns! {
160        pub TestEmpty {}
161    }
162
163    #[test]
164    fn scalar_indices() {
165        assert_eq!(TestScalar::A, 0);
166        assert_eq!(TestScalar::B, 1);
167        assert_eq!(TestScalar::Q, 2);
168        assert_eq!(TestScalar::NUM_COLUMNS, 3);
169    }
170
171    #[test]
172    fn mixed_scalar_and_array() {
173        assert_eq!(TestMixed::ADDR_B0, 0);
174        assert_eq!(TestMixed::ADDR_B1, 1);
175        assert_eq!(TestMixed::ADDR_B2, 2);
176        assert_eq!(TestMixed::ADDR_B3, 3);
177        assert_eq!(TestMixed::IS_WRITE, 4);
178        assert_eq!(TestMixed::SELECTOR, 5);
179        assert_eq!(TestMixed::AUX_INV, 6);
180        assert_eq!(TestMixed::DIFF_BYTE_IDX, 7);
181        assert_eq!(TestMixed::DIFF_BIT_IDX, 15);
182        assert_eq!(TestMixed::A_BITS, 23);
183        assert_eq!(TestMixed::NUM_COLUMNS, 31);
184    }
185
186    #[test]
187    fn large_arrays() {
188        assert_eq!(TestLargeArrays::STATE_BITS, 0);
189        assert_eq!(TestLargeArrays::RC_BITS, 1600);
190        assert_eq!(TestLargeArrays::LANES, 1664);
191        assert_eq!(TestLargeArrays::S_ROUND, 1689);
192        assert_eq!(TestLargeArrays::S_IN_OUT, 1690);
193        assert_eq!(TestLargeArrays::NUM_COLUMNS, 1691);
194    }
195
196    #[test]
197    fn empty_schema() {
198        assert_eq!(TestEmpty::NUM_COLUMNS, 0);
199        assert!(TestEmpty::build_layout().is_empty());
200    }
201
202    #[test]
203    fn scalar_layout() {
204        let layout = TestScalar::build_layout();
205        assert_eq!(layout.len(), 3);
206        assert_eq!(layout[0], ColumnType::B32);
207        assert_eq!(layout[1], ColumnType::B32);
208        assert_eq!(layout[2], ColumnType::Bit);
209    }
210
211    #[test]
212    fn mixed_layout() {
213        let layout = TestMixed::build_layout();
214        assert_eq!(layout.len(), TestMixed::NUM_COLUMNS);
215        assert_eq!(layout[0], ColumnType::B32);
216        assert_eq!(layout[4], ColumnType::Bit);
217        assert_eq!(layout[6], ColumnType::B128);
218
219        for i in 0..8 {
220            assert_eq!(layout[TestMixed::DIFF_BYTE_IDX + i], ColumnType::Bit);
221        }
222    }
223
224    #[test]
225    fn large_array_layout() {
226        let layout = TestLargeArrays::build_layout();
227        assert_eq!(layout.len(), 1691);
228        assert_eq!(layout[0], ColumnType::Bit);
229        assert_eq!(layout[1599], ColumnType::Bit);
230        assert_eq!(layout[1664], ColumnType::B64);
231        assert_eq!(layout[1688], ColumnType::B64);
232        assert_eq!(layout[1689], ColumnType::Bit);
233    }
234
235    #[test]
236    fn offsets_basic() {
237        let o = chiplet_offsets(&[9, 5, 10, 9, 17, 49]);
238        assert_eq!(o, vec![0, 9, 14, 24, 33, 50]);
239    }
240
241    #[test]
242    fn offsets_empty() {
243        assert!(chiplet_offsets(&[]).is_empty());
244    }
245
246    #[test]
247    fn offsets_single() {
248        assert_eq!(chiplet_offsets(&[42]), vec![0]);
249    }
250}