Skip to main content

hekate_program/
schema.rs

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