Skip to main content

hekate_core/
trace.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
5use crate::errors;
6use crate::poly::variant::PolyVariant;
7use alloc::vec;
8use alloc::vec::Vec;
9use core::any::TypeId;
10use core::fmt;
11use core::mem::transmute;
12use hekate_math::{
13    BinaryFieldExtras, Bit, Block8, Block16, Block32, Block64, Block128, CanonicalSerialize, Flat,
14    FlatPromote, HardwareField, PackableField, TowerField,
15};
16use zeroize::Zeroize;
17#[cfg(feature = "secure-memory")]
18use zeroize::ZeroizeOnDrop;
19
20#[derive(Clone, Copy, Debug, Eq, PartialEq)]
21pub enum Error {
22    InvalidParameters {
23        message: &'static str,
24    },
25    ColumnLengthMismatch {
26        expected_len: usize,
27        got_len: usize,
28    },
29    ColumnIndexOutOfBounds {
30        col_idx: usize,
31        num_cols: usize,
32    },
33    RowIndexOutOfBounds {
34        row_idx: usize,
35        num_rows: usize,
36    },
37    PointDimensionMismatch {
38        expected_len: usize,
39        got_len: usize,
40    },
41    ColumnTypeMismatch {
42        col_idx: usize,
43        expected: &'static str,
44        got: &'static str,
45    },
46}
47
48impl fmt::Display for Error {
49    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50        match self {
51            Self::InvalidParameters { message } => {
52                write!(f, "Trace invalid parameters: {message}")
53            }
54            Self::ColumnLengthMismatch {
55                expected_len,
56                got_len,
57            } => write!(
58                f,
59                "Trace column length mismatch: expected {expected_len}, got {got_len}",
60            ),
61            Self::ColumnIndexOutOfBounds { col_idx, num_cols } => write!(
62                f,
63                "Trace column index out of bounds: col_idx={col_idx}, num_cols={num_cols}",
64            ),
65            Self::RowIndexOutOfBounds { row_idx, num_rows } => write!(
66                f,
67                "Trace row index out of bounds: row_idx={row_idx}, num_rows={num_rows}",
68            ),
69            Self::PointDimensionMismatch {
70                expected_len,
71                got_len,
72            } => write!(
73                f,
74                "Trace evaluation point dimension mismatch: expected {expected_len}, got {got_len}",
75            ),
76            Self::ColumnTypeMismatch {
77                col_idx,
78                expected,
79                got,
80            } => write!(
81                f,
82                "Trace column type mismatch at col_idx={col_idx}: expected {expected}, got {got}",
83            ),
84        }
85    }
86}
87
88/// Bound for the proving field `F`:
89/// it must losslessly represent every `ColumnType`
90/// used by the trace (Bit, B8, B16, B32, B64, B128)
91/// and carry the additive-FFT substrate the
92/// Brakedown RS row code encodes over.
93pub trait TraceCompatibleField:
94    TowerField
95    + HardwareField
96    + PackableField
97    + BinaryFieldExtras
98    + FlatPromote<Block8>
99    + FlatPromote<Block16>
100    + FlatPromote<Block32>
101    + FlatPromote<Block64>
102    + FlatPromote<Block128>
103    + From<Bit>
104    + From<Block8>
105    + From<Block16>
106    + From<Block32>
107    + From<Block64>
108    + From<Block128>
109    + Send
110    + Sync
111{
112}
113
114impl<T> TraceCompatibleField for T where
115    T: TowerField
116        + HardwareField
117        + PackableField
118        + BinaryFieldExtras
119        + FlatPromote<Block8>
120        + FlatPromote<Block16>
121        + FlatPromote<Block32>
122        + FlatPromote<Block64>
123        + FlatPromote<Block128>
124        + From<Bit>
125        + From<Block8>
126        + From<Block16>
127        + From<Block32>
128        + From<Block64>
129        + From<Block128>
130        + Send
131        + Sync
132{
133}
134
135// =========================================================
136// TRACE TRAIT DEFINITION
137// =========================================================
138
139/// Execution-trace interface. Separates physical
140/// storage (`TraceColumn`) from the virtual
141/// polynomial view consumed by Sumcheck.
142pub trait Trace: Send + Sync {
143    /// `log2` of the trace height.
144    fn num_vars(&self) -> usize;
145
146    fn columns(&self) -> &[TraceColumn];
147
148    /// `2^num_vars`.
149    fn num_rows(&self) -> errors::Result<usize> {
150        num_rows_from_num_vars(self.num_vars())
151    }
152
153    fn num_cols(&self) -> usize {
154        self.columns().len()
155    }
156
157    fn column_layout(&self) -> Vec<ColumnType> {
158        self.columns().iter().map(|col| col.column_type()).collect()
159    }
160
161    /// Read a single trace cell and lift
162    /// it into `F` in the flat/hardware basis.
163    fn get_element<F: TraceCompatibleField>(
164        &self,
165        col_idx: usize,
166        row_idx: usize,
167    ) -> errors::Result<Flat<F>> {
168        let cols = self.columns();
169        let num_cols = self.num_cols();
170
171        if col_idx >= num_cols {
172            return Err(Error::ColumnIndexOutOfBounds { col_idx, num_cols }.into());
173        }
174
175        let num_rows = self.num_rows()?;
176        if row_idx >= num_rows {
177            return Err(Error::RowIndexOutOfBounds { row_idx, num_rows }.into());
178        }
179
180        match &cols[col_idx] {
181            TraceColumn::Bit(v) => Ok(Flat::from_raw(F::from(v[row_idx]))),
182            TraceColumn::B8(v) => Ok(F::promote_flat(v[row_idx])),
183            TraceColumn::B16(v) => Ok(F::promote_flat(v[row_idx])),
184            TraceColumn::B32(v) => Ok(F::promote_flat(v[row_idx])),
185            TraceColumn::B64(v) => Ok(F::promote_flat(v[row_idx])),
186            TraceColumn::B128(v) => Ok(F::promote_flat(v[row_idx])),
187        }
188    }
189
190    /// Zero-copy typed slice over a column.
191    /// Fails if `F` does not match
192    /// the column's storage type.
193    fn get_column_slice<F: 'static>(&self, col_idx: usize) -> errors::Result<&[F]> {
194        let cols = self.columns();
195        let num_cols = self.num_cols();
196
197        if col_idx >= num_cols {
198            return Err(Error::ColumnIndexOutOfBounds { col_idx, num_cols }.into());
199        }
200
201        let got = core::any::type_name::<F>();
202
203        match &cols[col_idx] {
204            TraceColumn::Bit(vec) => {
205                if TypeId::of::<F>() != TypeId::of::<Bit>() {
206                    return Err(Error::ColumnTypeMismatch {
207                        col_idx,
208                        expected: "Bit",
209                        got,
210                    }
211                    .into());
212                }
213
214                // SAFETY:
215                // The TypeId check guarantees F == Bit.
216                Ok(unsafe { transmute::<&[Bit], &[F]>(vec.as_slice()) })
217            }
218            TraceColumn::B8(vec) => {
219                if TypeId::of::<F>() != TypeId::of::<Flat<Block8>>() {
220                    return Err(Error::ColumnTypeMismatch {
221                        col_idx,
222                        expected: "Flat<Block8>",
223                        got,
224                    }
225                    .into());
226                }
227                Ok(unsafe { transmute::<&[Flat<Block8>], &[F]>(vec.as_slice()) })
228            }
229            TraceColumn::B16(vec) => {
230                if TypeId::of::<F>() != TypeId::of::<Flat<Block16>>() {
231                    return Err(Error::ColumnTypeMismatch {
232                        col_idx,
233                        expected: "Flat<Block16>",
234                        got,
235                    }
236                    .into());
237                }
238                Ok(unsafe { transmute::<&[Flat<Block16>], &[F]>(vec.as_slice()) })
239            }
240            TraceColumn::B32(vec) => {
241                if TypeId::of::<F>() != TypeId::of::<Flat<Block32>>() {
242                    return Err(Error::ColumnTypeMismatch {
243                        col_idx,
244                        expected: "Flat<Block32>",
245                        got,
246                    }
247                    .into());
248                }
249                Ok(unsafe { transmute::<&[Flat<Block32>], &[F]>(vec.as_slice()) })
250            }
251            TraceColumn::B64(vec) => {
252                if TypeId::of::<F>() != TypeId::of::<Flat<Block64>>() {
253                    return Err(Error::ColumnTypeMismatch {
254                        col_idx,
255                        expected: "Flat<Block64>",
256                        got,
257                    }
258                    .into());
259                }
260                Ok(unsafe { transmute::<&[Flat<Block64>], &[F]>(vec.as_slice()) })
261            }
262            TraceColumn::B128(vec) => {
263                if TypeId::of::<F>() != TypeId::of::<Flat<Block128>>() {
264                    return Err(Error::ColumnTypeMismatch {
265                        col_idx,
266                        expected: "Flat<Block128>",
267                        got,
268                    }
269                    .into());
270                }
271                Ok(unsafe { transmute::<&[Flat<Block128>], &[F]>(vec.as_slice()) })
272            }
273        }
274    }
275
276    /// Map physical columns to the `PolyVariant`s
277    /// consumed by Sumcheck. Default is a 1:1
278    /// `BitSlice` / `B{N}Slice` mapping; chiplets
279    /// that pack data (e.g. Keccak) override this
280    /// to expose virtual bit-columns.
281    fn get_poly_variants<F>(&'_ self) -> errors::Result<Vec<PolyVariant<'_, F>>>
282    where
283        F: TraceCompatibleField + 'static,
284    {
285        let cols = self.columns();
286        let mut variants = Vec::with_capacity(cols.len());
287
288        for (i, col) in cols.iter().enumerate() {
289            if i >= self.num_cols() {
290                return Err(errors::Error::Protocol {
291                    protocol: "air",
292                    message: "trace has fewer columns than required by AIR",
293                });
294            }
295
296            let variant = if let Some(s) = col.as_bit_slice() {
297                PolyVariant::BitSlice(s)
298            } else if let Some(s) = col.as_b8_slice() {
299                PolyVariant::B8Slice(s)
300            } else if let Some(s) = col.as_b16_slice() {
301                PolyVariant::B16Slice(s)
302            } else if let Some(s) = col.as_b32_slice() {
303                PolyVariant::B32Slice(s)
304            } else if let Some(s) = col.as_b64_slice() {
305                PolyVariant::B64Slice(s)
306            } else if let Some(s) = col.as_b128_slice() {
307                PolyVariant::B128Slice(s)
308            } else {
309                return Err(errors::Error::Protocol {
310                    protocol: "air",
311                    message: "unsupported trace column variant",
312                });
313            };
314
315            variants.push(variant);
316        }
317
318        Ok(variants)
319    }
320}
321
322// =========================================================
323// COLUMN LAYOUT METADATA
324// =========================================================
325
326/// Column storage type, without the data itself.
327/// Mixed-field traces require the verifier to know
328/// the exact byte width of every opened column in
329/// order to parse the raw LDT bytes.
330#[derive(Clone, Copy, Debug, PartialEq, Eq)]
331pub enum ColumnType {
332    /// AIR authors MUST `cs.assert_boolean(cs.col(idx))`
333    /// on every `Bit` used as selector, constraint operand,
334    /// or LogUp source, parse is byte-preserving, the
335    /// verifier accepts any byte and lifts it to `F`.
336    Bit,
337    B8,
338    B16,
339    B32,
340    B64,
341    B128,
342}
343
344impl ColumnType {
345    #[inline]
346    pub const fn byte_size(&self) -> usize {
347        match self {
348            Self::Bit => 1,
349            Self::B8 => 1,
350            Self::B16 => 2,
351            Self::B32 => 4,
352            Self::B64 => 8,
353            Self::B128 => 16,
354        }
355    }
356
357    /// The tower field this column's cells are committed
358    /// in under the MDS Reed-Solomon row code. GF(2),
359    /// GF(2^8), and GF(2^16) admit no rate-1/2 MDS RS.
360    #[inline]
361    pub const fn rs_field(&self) -> Self {
362        match self {
363            Self::Bit | Self::B8 | Self::B16 => Self::B32,
364            other => *other,
365        }
366    }
367
368    /// Parse a field element from its on-wire
369    /// bytes (little-endian, hardware basis)
370    /// without intermediate allocation.
371    pub fn parse_from_bytes<F>(&self, bytes: &[u8]) -> Flat<F>
372    where
373        F: TraceCompatibleField,
374    {
375        match self {
376            Self::Bit => Flat::from_raw(F::from(Bit::new(bytes[0]))),
377            Self::B8 => F::promote_flat(Flat::from_raw(Block8(bytes[0]))),
378            Self::B16 => {
379                let mut buf = [0u8; 2];
380                buf.copy_from_slice(&bytes[0..2]);
381
382                F::promote_flat(Flat::from_raw(Block16(u16::from_le_bytes(buf))))
383            }
384            Self::B32 => {
385                let mut buf = [0u8; 4];
386                buf.copy_from_slice(&bytes[0..4]);
387
388                F::promote_flat(Flat::from_raw(Block32(u32::from_le_bytes(buf))))
389            }
390            Self::B64 => {
391                let mut buf = [0u8; 8];
392                buf.copy_from_slice(&bytes[0..8]);
393
394                F::promote_flat(Flat::from_raw(Block64(u64::from_le_bytes(buf))))
395            }
396            Self::B128 => {
397                let mut buf = [0u8; 16];
398                buf.copy_from_slice(&bytes[0..16]);
399
400                F::promote_flat(Flat::from_raw(Block128(u128::from_le_bytes(buf))))
401            }
402        }
403    }
404}
405
406/// One typed column of the execution trace.
407/// Stored in hardware (flat) basis for all
408/// non-`Bit` variants.
409#[derive(Clone, Debug, Zeroize)]
410#[cfg_attr(feature = "secure-memory", derive(ZeroizeOnDrop))]
411pub enum TraceColumn {
412    Bit(Vec<Bit>),
413    B8(Vec<Flat<Block8>>),
414    B16(Vec<Flat<Block16>>),
415    B32(Vec<Flat<Block32>>),
416    B64(Vec<Flat<Block64>>),
417    B128(Vec<Flat<Block128>>),
418}
419
420impl TraceColumn {
421    /// Narrow a vector of `Block128` into a column
422    /// of `target_type` by truncating to the low bytes.
423    pub fn from_data(data: Vec<Block128>, target_type: ColumnType) -> Self {
424        match target_type {
425            ColumnType::Bit => {
426                let converted: Vec<Bit> = data
427                    .iter()
428                    .map(|val| {
429                        let bytes = val.to_bytes();
430                        Bit::from(bytes[0] & 1)
431                    })
432                    .collect();
433                TraceColumn::Bit(converted)
434            }
435            ColumnType::B8 => {
436                let converted: Vec<Flat<Block8>> = data
437                    .iter()
438                    .map(|val| {
439                        let bytes = val.to_bytes();
440                        Block8::from(bytes[0]).to_hardware()
441                    })
442                    .collect();
443                TraceColumn::B8(converted)
444            }
445            ColumnType::B16 => {
446                let converted: Vec<Flat<Block16>> = data
447                    .iter()
448                    .map(|val| {
449                        let bytes = val.to_bytes();
450                        let mut chunk = [0u8; 2];
451                        chunk.copy_from_slice(&bytes[0..2]);
452
453                        Block16::from(u16::from_le_bytes(chunk)).to_hardware()
454                    })
455                    .collect();
456                TraceColumn::B16(converted)
457            }
458            ColumnType::B32 => {
459                let converted: Vec<Flat<Block32>> = data
460                    .iter()
461                    .map(|val| {
462                        let bytes = val.to_bytes();
463                        let mut chunk = [0u8; 4];
464                        chunk.copy_from_slice(&bytes[0..4]);
465
466                        Block32::from(u32::from_le_bytes(chunk)).to_hardware()
467                    })
468                    .collect();
469                TraceColumn::B32(converted)
470            }
471            ColumnType::B64 => {
472                let converted: Vec<Flat<Block64>> = data
473                    .iter()
474                    .map(|val| {
475                        let bytes = val.to_bytes();
476                        let mut chunk = [0u8; 8];
477                        chunk.copy_from_slice(&bytes[0..8]);
478
479                        Block64::from(u64::from_le_bytes(chunk)).to_hardware()
480                    })
481                    .collect();
482                TraceColumn::B64(converted)
483            }
484            ColumnType::B128 => {
485                TraceColumn::B128(data.into_iter().map(|value| value.to_hardware()).collect())
486            }
487        }
488    }
489
490    pub fn len(&self) -> usize {
491        match self {
492            Self::Bit(v) => v.len(),
493            Self::B8(v) => v.len(),
494            Self::B16(v) => v.len(),
495            Self::B32(v) => v.len(),
496            Self::B64(v) => v.len(),
497            Self::B128(v) => v.len(),
498        }
499    }
500
501    pub fn is_empty(&self) -> bool {
502        self.len() == 0
503    }
504
505    pub fn is_all_zeros(&self) -> bool {
506        match self {
507            Self::Bit(v) => v.iter().all(|x| x.get() == 0),
508            Self::B8(v) => v.iter().all(|x| x.into_raw().0 == 0),
509            Self::B16(v) => v.iter().all(|x| x.into_raw().0 == 0),
510            Self::B32(v) => v.iter().all(|x| x.into_raw().0 == 0),
511            Self::B64(v) => v.iter().all(|x| x.into_raw().0 == 0),
512            Self::B128(v) => v.iter().all(|x| x.into_raw() == Block128::ZERO),
513        }
514    }
515
516    /// The `ColumnType` tag matching this variant.
517    pub fn column_type(&self) -> ColumnType {
518        match self {
519            Self::Bit(_) => ColumnType::Bit,
520            Self::B8(_) => ColumnType::B8,
521            Self::B16(_) => ColumnType::B16,
522            Self::B32(_) => ColumnType::B32,
523            Self::B64(_) => ColumnType::B64,
524            Self::B128(_) => ColumnType::B128,
525        }
526    }
527
528    /// Append this row's little-endian
529    /// serialization to `buf`. Used for
530    /// Merkle leaf hashing and LDT opening bytes.
531    pub fn append_bytes_at(&self, row_idx: usize, buf: &mut Vec<u8>) {
532        match self {
533            Self::Bit(v) => {
534                buf.push(v[row_idx].get());
535            }
536            Self::B8(v) => {
537                buf.push(v[row_idx].into_raw().0);
538            }
539            Self::B16(v) => {
540                buf.extend_from_slice(&v[row_idx].into_raw().0.to_le_bytes());
541            }
542            Self::B32(v) => {
543                buf.extend_from_slice(&v[row_idx].into_raw().0.to_le_bytes());
544            }
545            Self::B64(v) => {
546                buf.extend_from_slice(&v[row_idx].into_raw().0.to_le_bytes());
547            }
548            Self::B128(v) => {
549                buf.extend_from_slice(&v[row_idx].into_raw().0.to_le_bytes());
550            }
551        }
552    }
553
554    // ===========================================
555    // Typed slice accessors
556    // ===========================================
557
558    pub fn as_bit_slice(&self) -> Option<&[Bit]> {
559        if let Self::Bit(v) = self {
560            Some(v)
561        } else {
562            None
563        }
564    }
565
566    pub fn as_b8_slice(&self) -> Option<&[Flat<Block8>]> {
567        if let Self::B8(v) = self {
568            Some(v)
569        } else {
570            None
571        }
572    }
573    pub fn as_b16_slice(&self) -> Option<&[Flat<Block16>]> {
574        if let Self::B16(v) = self {
575            Some(v)
576        } else {
577            None
578        }
579    }
580    pub fn as_b32_slice(&self) -> Option<&[Flat<Block32>]> {
581        if let Self::B32(v) = self {
582            Some(v)
583        } else {
584            None
585        }
586    }
587    pub fn as_b64_slice(&self) -> Option<&[Flat<Block64>]> {
588        if let Self::B64(v) = self {
589            Some(v)
590        } else {
591            None
592        }
593    }
594    pub fn as_b128_slice(&self) -> Option<&[Flat<Block128>]> {
595        if let Self::B128(v) = self {
596            Some(v)
597        } else {
598            None
599        }
600    }
601}
602
603/// A concrete implementation of
604/// Trace using column-major storage.
605#[derive(Clone, Debug, Zeroize)]
606#[cfg_attr(feature = "secure-memory", derive(ZeroizeOnDrop))]
607pub struct ColumnTrace {
608    pub columns: Vec<TraceColumn>,
609    pub num_vars: usize,
610}
611
612impl Trace for ColumnTrace {
613    fn num_vars(&self) -> usize {
614        self.num_vars
615    }
616
617    fn columns(&self) -> &[TraceColumn] {
618        &self.columns
619    }
620}
621
622impl ColumnTrace {
623    pub fn new(num_vars: usize) -> errors::Result<Self> {
624        let num_rows = num_rows_from_num_vars(num_vars)?;
625        if num_rows == 0 {
626            return Err(Error::InvalidParameters {
627                message: "trace height is zero",
628            }
629            .into());
630        }
631
632        Ok(Self {
633            num_vars,
634            columns: Vec::new(),
635        })
636    }
637
638    /// Consume the trace and return
639    /// its owned column storage.
640    pub fn into_columns(mut self) -> Vec<TraceColumn> {
641        core::mem::take(&mut self.columns)
642    }
643
644    pub fn add_column(&mut self, col: TraceColumn) -> errors::Result<()> {
645        let expected_len = self.num_rows()?;
646        let got_len = col.len();
647
648        if got_len != expected_len {
649            return Err(Error::ColumnLengthMismatch {
650                expected_len,
651                got_len,
652            }
653            .into());
654        }
655
656        self.columns.push(col);
657
658        Ok(())
659    }
660}
661
662pub trait IntoTraceColumn {
663    fn into_trace_column(self) -> TraceColumn;
664}
665
666impl IntoTraceColumn for Vec<Bit> {
667    fn into_trace_column(self) -> TraceColumn {
668        TraceColumn::Bit(self)
669    }
670}
671
672impl IntoTraceColumn for Vec<Block8> {
673    fn into_trace_column(self) -> TraceColumn {
674        TraceColumn::B8(self.into_iter().map(|value| value.to_hardware()).collect())
675    }
676}
677
678impl IntoTraceColumn for Vec<Block16> {
679    fn into_trace_column(self) -> TraceColumn {
680        TraceColumn::B16(self.into_iter().map(|value| value.to_hardware()).collect())
681    }
682}
683
684impl IntoTraceColumn for Vec<Block32> {
685    fn into_trace_column(self) -> TraceColumn {
686        TraceColumn::B32(self.into_iter().map(|value| value.to_hardware()).collect())
687    }
688}
689
690impl IntoTraceColumn for Vec<Block64> {
691    fn into_trace_column(self) -> TraceColumn {
692        TraceColumn::B64(self.into_iter().map(|value| value.to_hardware()).collect())
693    }
694}
695
696impl IntoTraceColumn for Vec<Block128> {
697    fn into_trace_column(self) -> TraceColumn {
698        TraceColumn::B128(self.into_iter().map(|value| value.to_hardware()).collect())
699    }
700}
701
702impl IntoTraceColumn for Vec<Flat<Block8>> {
703    fn into_trace_column(self) -> TraceColumn {
704        TraceColumn::B8(self)
705    }
706}
707
708impl IntoTraceColumn for Vec<Flat<Block16>> {
709    fn into_trace_column(self) -> TraceColumn {
710        TraceColumn::B16(self)
711    }
712}
713
714impl IntoTraceColumn for Vec<Flat<Block32>> {
715    fn into_trace_column(self) -> TraceColumn {
716        TraceColumn::B32(self)
717    }
718}
719
720impl IntoTraceColumn for Vec<Flat<Block64>> {
721    fn into_trace_column(self) -> TraceColumn {
722        TraceColumn::B64(self)
723    }
724}
725
726impl IntoTraceColumn for Vec<Flat<Block128>> {
727    fn into_trace_column(self) -> TraceColumn {
728        TraceColumn::B128(self)
729    }
730}
731
732/// Zero-copy byte views `(ptr, elem_width)` for
733/// every column in a trace. Centralizes the
734/// `#[repr(transparent)]`-dependent pointer casts.
735pub fn get_col_views(columns: &[TraceColumn]) -> Vec<(&[u8], usize)> {
736    columns
737        .iter()
738        .map(|col| match col {
739            TraceColumn::Bit(v) => (
740                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len()) },
741                1,
742            ),
743            TraceColumn::B8(v) => (
744                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len()) },
745                1,
746            ),
747            TraceColumn::B16(v) => (
748                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 2) },
749                2,
750            ),
751            TraceColumn::B32(v) => (
752                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4) },
753                4,
754            ),
755            TraceColumn::B64(v) => (
756                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 8) },
757                8,
758            ),
759            TraceColumn::B128(v) => (
760                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 16) },
761                16,
762            ),
763        })
764        .collect()
765}
766
767// =========================================================
768// TRACE BUILDER
769// =========================================================
770
771macro_rules! impl_width_writes {
772    (
773        $(
774            $block:ty {
775                set: $set:ident,
776                set_array: $set_array:ident,
777                push: $push:ident,
778                set_flat: $set_flat:ident,
779                set_array_flat: $set_array_flat:ident,
780                column: $expect:ident,
781            }
782        )*
783    ) => {
784        $(
785            #[inline]
786            pub fn $set(&mut self, col: usize, row: usize, val: $block) -> errors::Result<()> {
787                self.$set_flat(col, row, val.to_hardware())
788            }
789
790            #[inline]
791            pub fn $set_flat(
792                &mut self,
793                col: usize,
794                row: usize,
795                val: Flat<$block>,
796            ) -> errors::Result<()> {
797                let num_rows = self.num_rows;
798                let data = self.$expect(col)?;
799                let slot = data.get_mut(row).ok_or(Error::RowIndexOutOfBounds {
800                    row_idx: row,
801                    num_rows,
802                })?;
803
804                *slot = val;
805
806                Ok(())
807            }
808
809            #[inline]
810            pub fn $push(&mut self, col: usize, val: $block) -> errors::Result<()> {
811                let row = self.cursor(col)?;
812                self.$set(col, row, val)?;
813
814                self.cursors[col] = row + 1;
815
816                Ok(())
817            }
818
819            pub fn $set_array(
820                &mut self,
821                base: usize,
822                row: usize,
823                values: &[$block],
824            ) -> errors::Result<()> {
825                for (i, &val) in values.iter().enumerate() {
826                    self.$set(base + i, row, val)?;
827                }
828
829                Ok(())
830            }
831
832            pub fn $set_array_flat(
833                &mut self,
834                base: usize,
835                row: usize,
836                values: &[Flat<$block>],
837            ) -> errors::Result<()> {
838                for (i, &val) in values.iter().enumerate() {
839                    self.$set_flat(base + i, row, val)?;
840                }
841
842                Ok(())
843            }
844        )*
845    };
846}
847
848/// Schema-driven builder. Every column
849/// is allocated zero-filled from the layout;
850/// unfilled rows stay zero, so padding is implicit.
851pub struct TraceBuilder {
852    columns: Vec<TraceColumn>,
853    num_vars: usize,
854    num_rows: usize,
855    cursors: Vec<usize>,
856}
857
858impl TraceBuilder {
859    pub fn new(layout: &[ColumnType], num_vars: usize) -> errors::Result<Self> {
860        let num_rows = num_rows_from_num_vars(num_vars)?;
861        let columns = layout
862            .iter()
863            .map(|ct| match ct {
864                ColumnType::Bit => TraceColumn::Bit(vec![Bit::ZERO; num_rows]),
865                ColumnType::B8 => TraceColumn::B8(vec![Block8::ZERO.to_hardware(); num_rows]),
866                ColumnType::B16 => TraceColumn::B16(vec![Block16::ZERO.to_hardware(); num_rows]),
867                ColumnType::B32 => TraceColumn::B32(vec![Block32::ZERO.to_hardware(); num_rows]),
868                ColumnType::B64 => TraceColumn::B64(vec![Block64::ZERO.to_hardware(); num_rows]),
869                ColumnType::B128 => TraceColumn::B128(vec![Block128::ZERO.to_hardware(); num_rows]),
870            })
871            .collect();
872
873        Ok(Self {
874            columns,
875            num_vars,
876            num_rows,
877            cursors: vec![0; layout.len()],
878        })
879    }
880
881    /// `2^num_vars`.
882    #[inline]
883    pub fn num_rows(&self) -> usize {
884        self.num_rows
885    }
886
887    // =========================================================
888    // Column writes: indexed, cursor-push, array
889    // =========================================================
890
891    #[inline]
892    pub fn set_bit(&mut self, col: usize, row: usize, val: Bit) -> errors::Result<()> {
893        let num_rows = self.num_rows;
894        let data = self.expect_bit_col(col)?;
895        let slot = data.get_mut(row).ok_or(Error::RowIndexOutOfBounds {
896            row_idx: row,
897            num_rows,
898        })?;
899
900        *slot = val;
901
902        Ok(())
903    }
904
905    pub fn set_bit_array(&mut self, base: usize, row: usize, values: &[Bit]) -> errors::Result<()> {
906        for (i, &val) in values.iter().enumerate() {
907            self.set_bit(base + i, row, val)?;
908        }
909
910        Ok(())
911    }
912
913    #[inline]
914    pub fn push_bit(&mut self, col: usize, val: Bit) -> errors::Result<()> {
915        let row = self.cursor(col)?;
916        self.set_bit(col, row, val)?;
917
918        self.cursors[col] = row + 1;
919
920        Ok(())
921    }
922
923    impl_width_writes! {
924        Block8 {
925            set: set_b8,
926            set_array: set_b8_array,
927            push: push_b8,
928            set_flat: set_b8_flat,
929            set_array_flat: set_b8_array_flat,
930            column: expect_b8_col,
931        }
932
933        Block16 {
934            set: set_b16,
935            set_array: set_b16_array,
936            push: push_b16,
937            set_flat: set_b16_flat,
938            set_array_flat: set_b16_array_flat,
939            column: expect_b16_col,
940        }
941
942        Block32 {
943            set: set_b32,
944            set_array: set_b32_array,
945            push: push_b32,
946            set_flat: set_b32_flat,
947            set_array_flat: set_b32_array_flat,
948            column: expect_b32_col,
949        }
950
951        Block64 {
952            set: set_b64,
953            set_array: set_b64_array,
954            push: push_b64,
955            set_flat: set_b64_flat,
956            set_array_flat: set_b64_array_flat,
957            column: expect_b64_col,
958        }
959
960        Block128 {
961            set: set_b128,
962            set_array: set_b128_array,
963            push: push_b128,
964            set_flat: set_b128_flat,
965            set_array_flat: set_b128_array_flat,
966            column: expect_b128_col,
967        }
968    }
969
970    // =========================================================
971    // Selector helpers
972    // =========================================================
973
974    /// Write `ONE` into rows `[0, active_rows)`
975    /// of a `Bit` column. The tail stays zero from allocation.
976    pub fn fill_selector(&mut self, col: usize, active_rows: usize) -> errors::Result<()> {
977        let limit = active_rows.min(self.num_rows);
978        let data = self.expect_bit_col(col)?;
979
980        for slot in data.iter_mut().take(limit) {
981            *slot = Bit::ONE;
982        }
983
984        Ok(())
985    }
986
987    // =========================================================
988    // Finalization
989    // =========================================================
990
991    /// Consume the builder and return a `ColumnTrace`.
992    /// Column order matches the schema passed to `new`.
993    pub fn build(self) -> ColumnTrace {
994        ColumnTrace {
995            columns: self.columns,
996            num_vars: self.num_vars,
997        }
998    }
999
1000    // =========================================================
1001    // Internal helpers
1002    // =========================================================
1003
1004    #[inline]
1005    fn cursor(&self, col: usize) -> errors::Result<usize> {
1006        self.cursors.get(col).copied().ok_or_else(|| {
1007            Error::ColumnIndexOutOfBounds {
1008                col_idx: col,
1009                num_cols: self.columns.len(),
1010            }
1011            .into()
1012        })
1013    }
1014
1015    #[inline]
1016    fn expect_bit_col(&mut self, col: usize) -> errors::Result<&mut Vec<Bit>> {
1017        let num_cols = self.columns.len();
1018        let tc = self
1019            .columns
1020            .get_mut(col)
1021            .ok_or(Error::ColumnIndexOutOfBounds {
1022                col_idx: col,
1023                num_cols,
1024            })?;
1025
1026        match tc {
1027            TraceColumn::Bit(data) => Ok(data),
1028            other => Err(Error::ColumnTypeMismatch {
1029                col_idx: col,
1030                expected: "Bit",
1031                got: other.column_type_name(),
1032            }
1033            .into()),
1034        }
1035    }
1036
1037    #[inline]
1038    fn expect_b8_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block8>>> {
1039        let num_cols = self.columns.len();
1040        let tc = self
1041            .columns
1042            .get_mut(col)
1043            .ok_or(Error::ColumnIndexOutOfBounds {
1044                col_idx: col,
1045                num_cols,
1046            })?;
1047
1048        match tc {
1049            TraceColumn::B8(data) => Ok(data),
1050            other => Err(Error::ColumnTypeMismatch {
1051                col_idx: col,
1052                expected: "B8",
1053                got: other.column_type_name(),
1054            }
1055            .into()),
1056        }
1057    }
1058
1059    #[inline]
1060    fn expect_b16_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block16>>> {
1061        let num_cols = self.columns.len();
1062        let tc = self
1063            .columns
1064            .get_mut(col)
1065            .ok_or(Error::ColumnIndexOutOfBounds {
1066                col_idx: col,
1067                num_cols,
1068            })?;
1069
1070        match tc {
1071            TraceColumn::B16(data) => Ok(data),
1072            other => Err(Error::ColumnTypeMismatch {
1073                col_idx: col,
1074                expected: "B16",
1075                got: other.column_type_name(),
1076            }
1077            .into()),
1078        }
1079    }
1080
1081    #[inline]
1082    fn expect_b32_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block32>>> {
1083        let num_cols = self.columns.len();
1084        let tc = self
1085            .columns
1086            .get_mut(col)
1087            .ok_or(Error::ColumnIndexOutOfBounds {
1088                col_idx: col,
1089                num_cols,
1090            })?;
1091
1092        match tc {
1093            TraceColumn::B32(data) => Ok(data),
1094            other => Err(Error::ColumnTypeMismatch {
1095                col_idx: col,
1096                expected: "B32",
1097                got: other.column_type_name(),
1098            }
1099            .into()),
1100        }
1101    }
1102
1103    #[inline]
1104    fn expect_b64_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block64>>> {
1105        let num_cols = self.columns.len();
1106        let tc = self
1107            .columns
1108            .get_mut(col)
1109            .ok_or(Error::ColumnIndexOutOfBounds {
1110                col_idx: col,
1111                num_cols,
1112            })?;
1113
1114        match tc {
1115            TraceColumn::B64(data) => Ok(data),
1116            other => Err(Error::ColumnTypeMismatch {
1117                col_idx: col,
1118                expected: "B64",
1119                got: other.column_type_name(),
1120            }
1121            .into()),
1122        }
1123    }
1124
1125    #[inline]
1126    fn expect_b128_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block128>>> {
1127        let num_cols = self.columns.len();
1128        let tc = self
1129            .columns
1130            .get_mut(col)
1131            .ok_or(Error::ColumnIndexOutOfBounds {
1132                col_idx: col,
1133                num_cols,
1134            })?;
1135
1136        match tc {
1137            TraceColumn::B128(data) => Ok(data),
1138            other => Err(Error::ColumnTypeMismatch {
1139                col_idx: col,
1140                expected: "B128",
1141                got: other.column_type_name(),
1142            }
1143            .into()),
1144        }
1145    }
1146}
1147
1148impl TraceColumn {
1149    fn column_type_name(&self) -> &'static str {
1150        match self {
1151            Self::Bit(_) => "Bit",
1152            Self::B8(_) => "B8",
1153            Self::B16(_) => "B16",
1154            Self::B32(_) => "B32",
1155            Self::B64(_) => "B64",
1156            Self::B128(_) => "B128",
1157        }
1158    }
1159}
1160
1161fn num_rows_from_num_vars(num_vars: usize) -> errors::Result<usize> {
1162    let num_vars_u32 = match u32::try_from(num_vars) {
1163        Ok(v) => v,
1164        Err(_) => {
1165            return Err(Error::InvalidParameters {
1166                message: "num_vars too large",
1167            }
1168            .into());
1169        }
1170    };
1171
1172    let Some(num_rows) = 1usize.checked_shl(num_vars_u32) else {
1173        return Err(Error::InvalidParameters {
1174            message: "num_rows overflow",
1175        }
1176        .into());
1177    };
1178
1179    if num_rows == 0 {
1180        return Err(Error::InvalidParameters {
1181            message: "num_rows is zero",
1182        }
1183        .into());
1184    }
1185
1186    Ok(num_rows)
1187}
1188
1189// =================================================================
1190// UNIT TESTS
1191// =================================================================
1192
1193#[cfg(test)]
1194mod tests {
1195    use super::*;
1196    use crate::errors;
1197    use hekate_math::HardwareField;
1198
1199    fn create_mock_trace(num_vars: usize) -> ColumnTrace {
1200        ColumnTrace::new(num_vars).unwrap()
1201    }
1202
1203    #[test]
1204    fn trace_construction_basic() {
1205        let num_vars = 3;
1206        let mut trace = create_mock_trace(num_vars);
1207
1208        let col_data = vec![Block128::from(1u8); 8];
1209        trace.add_column(col_data.into_trace_column()).unwrap();
1210
1211        assert_eq!(trace.num_rows().unwrap(), 8);
1212        assert_eq!(trace.num_cols(), 1);
1213        assert_eq!(trace.num_vars, 3);
1214    }
1215
1216    #[test]
1217    fn trace_add_column_wrong_len() {
1218        let num_vars = 2;
1219        let mut trace = create_mock_trace(num_vars);
1220
1221        let col_data = vec![Block128::ZERO; 5];
1222        let err = trace
1223            .add_column(col_data.into_trace_column())
1224            .expect_err("Expected length mismatch error");
1225
1226        assert!(matches!(
1227            err,
1228            errors::Error::Trace(Error::ColumnLengthMismatch { .. })
1229        ));
1230    }
1231
1232    #[test]
1233    fn trace_get_element_mixed_types() {
1234        let num_vars = 1;
1235        let mut trace = create_mock_trace(num_vars);
1236
1237        trace
1238            .add_column(TraceColumn::Bit(vec![Bit::new(0), Bit::new(1)]))
1239            .unwrap();
1240        trace
1241            .add_column(vec![Block32::from(10u32), Block32::from(20u32)].into_trace_column())
1242            .unwrap();
1243
1244        let val0_r0: Flat<Block128> = trace.get_element(0, 0).unwrap();
1245        let val0_r1: Flat<Block128> = trace.get_element(0, 1).unwrap();
1246        let val1_r0: Flat<Block128> = trace.get_element(1, 0).unwrap();
1247        let val1_r1: Flat<Block128> = trace.get_element(1, 1).unwrap();
1248
1249        assert_eq!(val0_r0.into_raw(), Block128::ZERO);
1250        assert_eq!(val0_r1.into_raw(), Block128::ONE);
1251
1252        let expected_10 = Block128::promote_flat(Block32::from(10u32).to_hardware()).into_raw();
1253        let expected_20 = Block128::promote_flat(Block32::from(20u32).to_hardware()).into_raw();
1254
1255        assert_eq!(val1_r0.into_raw(), expected_10);
1256        assert_eq!(val1_r1.into_raw(), expected_20);
1257    }
1258
1259    #[test]
1260    fn get_element_oob_row() {
1261        let mut trace = create_mock_trace(1);
1262        trace
1263            .add_column(TraceColumn::Bit(vec![Bit::ZERO; 2]))
1264            .unwrap();
1265        trace
1266            .get_element::<Block128>(0, 2)
1267            .expect_err("Expected out-of-bounds row error");
1268    }
1269
1270    #[test]
1271    fn get_element_oob_col() {
1272        let trace = create_mock_trace(1);
1273        trace
1274            .get_element::<Block128>(0, 0)
1275            .expect_err("Expected out-of-bounds column error");
1276    }
1277
1278    // ===========================================
1279    // SAFETY & SLICE TESTS
1280    // ===========================================
1281
1282    #[test]
1283    fn get_column_slice_correct_type() {
1284        let mut trace = create_mock_trace(2);
1285        let data = vec![
1286            Block32::from(1u32),
1287            Block32::from(2u32),
1288            Block32::from(3u32),
1289            Block32::from(4u32),
1290        ];
1291        trace.add_column(data.clone().into_trace_column()).unwrap();
1292
1293        let expected_hw: Vec<Flat<Block32>> = data.into_iter().map(|x| x.to_hardware()).collect();
1294
1295        let slice: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1296        assert_eq!(slice, expected_hw.as_slice());
1297    }
1298
1299    #[test]
1300    fn get_column_slice_wrong_type() {
1301        let mut trace = create_mock_trace(1);
1302        trace
1303            .add_column(vec![Block128::ZERO; 2].into_trace_column())
1304            .unwrap();
1305
1306        trace
1307            .get_column_slice::<Flat<Block32>>(0)
1308            .expect_err("Expected column type mismatch error");
1309    }
1310
1311    #[test]
1312    fn trace_stores_hardware_basis() {
1313        let mut trace = create_mock_trace(2);
1314
1315        let tower_data = vec![
1316            Block32::from(42u32),
1317            Block32::from(13u32),
1318            Block32::from(255u32),
1319            Block32::from(1u32),
1320        ];
1321
1322        let expected_hardware: Vec<Block32> = tower_data
1323            .iter()
1324            .map(|x| x.to_hardware().into_raw())
1325            .collect();
1326
1327        trace.add_column(tower_data.into_trace_column()).unwrap();
1328
1329        let stored: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1330
1331        for (i, (&stored_val, &expected_val)) in
1332            stored.iter().zip(expected_hardware.iter()).enumerate()
1333        {
1334            assert_eq!(
1335                stored_val.into_raw(),
1336                expected_val,
1337                "Row {}: stored value {:?} != expected hardware {:?}",
1338                i,
1339                stored_val,
1340                expected_val
1341            );
1342        }
1343    }
1344
1345    #[test]
1346    fn trace_hardware_basis_homomorphism() {
1347        let mut trace = create_mock_trace(3);
1348
1349        let a_tower = vec![Block32::from(5u32); 8];
1350        let b_tower = vec![Block32::from(7u32); 8];
1351
1352        trace
1353            .add_column(a_tower.clone().into_trace_column())
1354            .unwrap();
1355        trace
1356            .add_column(b_tower.clone().into_trace_column())
1357            .unwrap();
1358
1359        let a_stored: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1360        let b_stored: &[Flat<Block32>] = trace.get_column_slice(1).unwrap();
1361
1362        let a_hw_expected = a_tower[0].to_hardware().into_raw();
1363        let b_hw_expected = b_tower[0].to_hardware().into_raw();
1364
1365        assert_eq!(a_stored[0].into_raw(), a_hw_expected);
1366        assert_eq!(b_stored[0].into_raw(), b_hw_expected);
1367
1368        let product_hw = a_stored[0] * b_stored[0];
1369        let product_expected = (a_tower[0] * b_tower[0]).to_hardware().into_raw();
1370
1371        assert_eq!(product_hw.into_raw(), product_expected);
1372    }
1373
1374    // =========================================================
1375    // TRACE BUILDER TESTS
1376    // =========================================================
1377
1378    #[test]
1379    fn trace_builder_construction_and_auto_padding() {
1380        let layout = &[ColumnType::B32, ColumnType::Bit];
1381        let tb = TraceBuilder::new(layout, 2).unwrap(); // 4 rows
1382        assert_eq!(tb.num_rows(), 4);
1383
1384        let trace = tb.build();
1385        assert_eq!(trace.num_cols(), 2);
1386        assert_eq!(trace.num_rows().unwrap(), 4);
1387
1388        // All values should be zero (auto-padding)
1389        assert!(trace.columns[0].is_all_zeros());
1390        assert!(trace.columns[1].is_all_zeros());
1391    }
1392
1393    #[test]
1394    fn trace_builder_set_b32_stores_hardware_basis() {
1395        let layout = &[ColumnType::B32];
1396        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1397
1398        tb.set_b32(0, 0, Block32::from(42u32)).unwrap();
1399        tb.set_b32(0, 1, Block32::from(13u32)).unwrap();
1400
1401        let trace = tb.build();
1402        let stored: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1403
1404        assert_eq!(stored[0], Block32::from(42u32).to_hardware());
1405        assert_eq!(stored[1], Block32::from(13u32).to_hardware());
1406    }
1407
1408    #[test]
1409    fn trace_builder_column_ordering_matches_schema() {
1410        let layout = &[ColumnType::Bit, ColumnType::B32, ColumnType::B128];
1411        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1412
1413        tb.set_bit(0, 0, Bit::ONE).unwrap();
1414        tb.set_b32(1, 0, Block32::from(99u32)).unwrap();
1415        tb.set_b128(2, 0, Block128::from(7u8)).unwrap();
1416
1417        let trace = tb.build();
1418        assert_eq!(trace.columns[0].column_type(), ColumnType::Bit);
1419        assert_eq!(trace.columns[1].column_type(), ColumnType::B32);
1420        assert_eq!(trace.columns[2].column_type(), ColumnType::B128);
1421    }
1422
1423    #[test]
1424    fn trace_builder_type_mismatch_returns_error() {
1425        let layout = &[ColumnType::Bit];
1426        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1427
1428        let err = tb.set_b32(0, 0, Block32::ZERO);
1429        assert!(err.is_err());
1430    }
1431
1432    #[test]
1433    fn trace_builder_row_out_of_bounds_returns_error() {
1434        let layout = &[ColumnType::B32];
1435        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1436
1437        let err = tb.set_b32(0, 2, Block32::ZERO);
1438        assert!(err.is_err());
1439    }
1440
1441    #[test]
1442    fn trace_builder_col_out_of_bounds_returns_error() {
1443        let layout = &[ColumnType::B32];
1444        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1445
1446        let err = tb.set_b32(1, 0, Block32::ZERO);
1447        assert!(err.is_err());
1448    }
1449
1450    #[test]
1451    fn trace_builder_fill_selector() {
1452        let layout = &[ColumnType::Bit];
1453        let mut tb = TraceBuilder::new(layout, 2).unwrap(); // 4 rows
1454
1455        tb.fill_selector(0, 3).unwrap(); // rows 0,1,2 = ONE
1456
1457        let trace = tb.build();
1458        let bits = trace.columns[0].as_bit_slice().unwrap();
1459        assert_eq!(bits[0], Bit::ONE);
1460        assert_eq!(bits[1], Bit::ONE);
1461        assert_eq!(bits[2], Bit::ONE);
1462        assert_eq!(bits[3], Bit::ZERO); // padding
1463    }
1464
1465    #[test]
1466    fn trace_builder_push_mode() {
1467        let layout = &[ColumnType::B32, ColumnType::Bit];
1468        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1469
1470        tb.push_b32(0, Block32::from(10u32)).unwrap();
1471        tb.push_b32(0, Block32::from(20u32)).unwrap();
1472        tb.push_bit(1, Bit::ONE).unwrap();
1473        // row 1 of col 1 stays zero (auto-pad)
1474
1475        let trace = tb.build();
1476        let b32s: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1477        assert_eq!(b32s[0], Block32::from(10u32).to_hardware());
1478        assert_eq!(b32s[1], Block32::from(20u32).to_hardware());
1479
1480        let bits = trace.columns[1].as_bit_slice().unwrap();
1481        assert_eq!(bits[0], Bit::ONE);
1482        assert_eq!(bits[1], Bit::ZERO);
1483    }
1484
1485    #[test]
1486    fn trace_builder_set_b32_array() {
1487        let layout = &[ColumnType::B32, ColumnType::B32, ColumnType::B32];
1488        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1489
1490        let vals = [
1491            Block32::from(1u32),
1492            Block32::from(2u32),
1493            Block32::from(3u32),
1494        ];
1495        tb.set_b32_array(0, 0, &vals).unwrap();
1496
1497        let trace = tb.build();
1498        for (i, &expected) in vals.iter().enumerate() {
1499            let stored: &[Flat<Block32>] = trace.get_column_slice(i).unwrap();
1500            assert_eq!(stored[0], expected.to_hardware());
1501        }
1502    }
1503
1504    #[test]
1505    fn trace_builder_set_b8() {
1506        let layout = &[ColumnType::B8];
1507        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1508
1509        tb.set_b8(0, 0, Block8(0xAB)).unwrap();
1510        tb.set_b8(0, 1, Block8(0xCD)).unwrap();
1511
1512        let trace = tb.build();
1513        let stored: &[Flat<Block8>] = trace.get_column_slice(0).unwrap();
1514        assert_eq!(stored[0], Block8(0xAB).to_hardware());
1515        assert_eq!(stored[1], Block8(0xCD).to_hardware());
1516    }
1517
1518    #[test]
1519    fn trace_builder_set_b16() {
1520        let layout = &[ColumnType::B16];
1521        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1522
1523        tb.set_b16(0, 0, Block16(1000)).unwrap();
1524        tb.set_b16(0, 1, Block16(2000)).unwrap();
1525
1526        let trace = tb.build();
1527        let stored: &[Flat<Block16>] = trace.get_column_slice(0).unwrap();
1528        assert_eq!(stored[0], Block16(1000).to_hardware());
1529        assert_eq!(stored[1], Block16(2000).to_hardware());
1530    }
1531
1532    #[test]
1533    fn trace_builder_set_b64() {
1534        let layout = &[ColumnType::B64];
1535        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1536
1537        tb.set_b64(0, 0, Block64(0xDEADBEEF_CAFEBABE)).unwrap();
1538
1539        let trace = tb.build();
1540        let stored: &[Flat<Block64>] = trace.get_column_slice(0).unwrap();
1541        assert_eq!(stored[0], Block64(0xDEADBEEF_CAFEBABE).to_hardware());
1542        assert_eq!(stored[1], Block64::ZERO.to_hardware()); // auto-padding
1543    }
1544
1545    #[test]
1546    fn trace_builder_set_b128() {
1547        let layout = &[ColumnType::B128];
1548        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1549
1550        let val = Block128::from(0xFFu8);
1551        tb.set_b128(0, 0, val).unwrap();
1552
1553        let trace = tb.build();
1554        let stored: &[Flat<Block128>] = trace.get_column_slice(0).unwrap();
1555        assert_eq!(stored[0], val.to_hardware());
1556    }
1557
1558    #[test]
1559    fn trace_builder_push_all_types() {
1560        let layout = &[
1561            ColumnType::Bit,
1562            ColumnType::B8,
1563            ColumnType::B16,
1564            ColumnType::B32,
1565            ColumnType::B64,
1566            ColumnType::B128,
1567        ];
1568        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1569
1570        tb.push_bit(0, Bit::ONE).unwrap();
1571        tb.push_b8(1, Block8(0x42)).unwrap();
1572        tb.push_b16(2, Block16(1234)).unwrap();
1573        tb.push_b32(3, Block32::from(5678u32)).unwrap();
1574        tb.push_b64(4, Block64(9999)).unwrap();
1575        tb.push_b128(5, Block128::from(77u8)).unwrap();
1576
1577        let trace = tb.build();
1578        assert_eq!(trace.num_cols(), 6);
1579
1580        let bits = trace.columns[0].as_bit_slice().unwrap();
1581        assert_eq!(bits[0], Bit::ONE);
1582        assert_eq!(bits[1], Bit::ZERO); // auto-pad
1583
1584        let b8s: &[Flat<Block8>] = trace.get_column_slice(1).unwrap();
1585        assert_eq!(b8s[0], Block8(0x42).to_hardware());
1586
1587        let b16s: &[Flat<Block16>] = trace.get_column_slice(2).unwrap();
1588        assert_eq!(b16s[0], Block16(1234).to_hardware());
1589
1590        let b32s: &[Flat<Block32>] = trace.get_column_slice(3).unwrap();
1591        assert_eq!(b32s[0], Block32::from(5678u32).to_hardware());
1592
1593        let b64s: &[Flat<Block64>] = trace.get_column_slice(4).unwrap();
1594        assert_eq!(b64s[0], Block64(9999).to_hardware());
1595
1596        let b128s: &[Flat<Block128>] = trace.get_column_slice(5).unwrap();
1597        assert_eq!(b128s[0], Block128::from(77u8).to_hardware());
1598    }
1599
1600    #[test]
1601    fn trace_builder_type_mismatch_all_setters() {
1602        // B32 column, every non-B32 setter should fail
1603        let layout = &[ColumnType::B32];
1604        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1605
1606        assert!(tb.set_bit(0, 0, Bit::ONE).is_err());
1607        assert!(tb.set_b8(0, 0, Block8(1)).is_err());
1608        assert!(tb.set_b16(0, 0, Block16(1)).is_err());
1609        assert!(tb.set_b64(0, 0, Block64(1)).is_err());
1610        assert!(tb.set_b128(0, 0, Block128::ONE).is_err());
1611
1612        // Correct type succeeds
1613        assert!(tb.set_b32(0, 0, Block32::ONE).is_ok());
1614    }
1615
1616    #[test]
1617    fn trace_builder_array_setters_all_types() {
1618        let layout = &[
1619            ColumnType::Bit,
1620            ColumnType::Bit,
1621            ColumnType::B8,
1622            ColumnType::B8,
1623            ColumnType::B64,
1624            ColumnType::B64,
1625            ColumnType::B128,
1626            ColumnType::B128,
1627        ];
1628        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1629
1630        tb.set_bit_array(0, 0, &[Bit::ONE, Bit::ZERO]).unwrap();
1631        tb.set_b8_array(2, 0, &[Block8(10), Block8(20)]).unwrap();
1632        tb.set_b64_array(4, 0, &[Block64(100), Block64(200)])
1633            .unwrap();
1634        tb.set_b128_array(6, 0, &[Block128::ONE, Block128::from(2u8)])
1635            .unwrap();
1636
1637        let trace = tb.build();
1638
1639        let bits = trace.columns[0].as_bit_slice().unwrap();
1640        assert_eq!(bits[0], Bit::ONE);
1641
1642        let bits1 = trace.columns[1].as_bit_slice().unwrap();
1643        assert_eq!(bits1[0], Bit::ZERO);
1644
1645        let b8s: &[Flat<Block8>] = trace.get_column_slice(2).unwrap();
1646        assert_eq!(b8s[0], Block8(10).to_hardware());
1647
1648        let b8s1: &[Flat<Block8>] = trace.get_column_slice(3).unwrap();
1649        assert_eq!(b8s1[0], Block8(20).to_hardware());
1650    }
1651
1652    #[test]
1653    fn trace_builder_invalid_num_vars() {
1654        let layout = &[ColumnType::B32];
1655        // num_vars too large to shift
1656        assert!(TraceBuilder::new(layout, 128).is_err());
1657    }
1658}