1use alloc::vec::Vec;
40
41pub use alloc::vec::Vec as SchemaVec;
42pub use hekate_core::trace::ColumnType;
43
44#[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 #[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 (@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 (@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 (@push $v:ident, $ty:ident) => {
98 $v.push($crate::schema::ColumnType::$ty);
99 };
100
101 (@push $v:ident, [$ty:ident; $n:expr]) => {
103 $v.extend(::core::iter::repeat($crate::schema::ColumnType::$ty).take($n));
104 };
105}
106
107pub 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}