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
771/// Schema-driven builder. Every column
772/// is allocated zero-filled from the layout;
773/// unfilled rows stay zero, so padding is implicit.
774pub struct TraceBuilder {
775    columns: Vec<TraceColumn>,
776    num_vars: usize,
777    num_rows: usize,
778    cursors: Vec<usize>,
779}
780
781impl TraceBuilder {
782    pub fn new(layout: &[ColumnType], num_vars: usize) -> errors::Result<Self> {
783        let num_rows = num_rows_from_num_vars(num_vars)?;
784        let columns = layout
785            .iter()
786            .map(|ct| match ct {
787                ColumnType::Bit => TraceColumn::Bit(vec![Bit::ZERO; num_rows]),
788                ColumnType::B8 => TraceColumn::B8(vec![Block8::ZERO.to_hardware(); num_rows]),
789                ColumnType::B16 => TraceColumn::B16(vec![Block16::ZERO.to_hardware(); num_rows]),
790                ColumnType::B32 => TraceColumn::B32(vec![Block32::ZERO.to_hardware(); num_rows]),
791                ColumnType::B64 => TraceColumn::B64(vec![Block64::ZERO.to_hardware(); num_rows]),
792                ColumnType::B128 => TraceColumn::B128(vec![Block128::ZERO.to_hardware(); num_rows]),
793            })
794            .collect();
795
796        Ok(Self {
797            columns,
798            num_vars,
799            num_rows,
800            cursors: vec![0; layout.len()],
801        })
802    }
803
804    /// `2^num_vars`.
805    #[inline]
806    pub fn num_rows(&self) -> usize {
807        self.num_rows
808    }
809
810    // =========================================================
811    // Indexed write (random access)
812    // =========================================================
813
814    #[inline]
815    pub fn set_bit(&mut self, col: usize, row: usize, val: Bit) -> errors::Result<()> {
816        let num_rows = self.num_rows;
817        let data = self.expect_bit_col(col)?;
818        let slot = data.get_mut(row).ok_or(Error::RowIndexOutOfBounds {
819            row_idx: row,
820            num_rows,
821        })?;
822
823        *slot = val;
824
825        Ok(())
826    }
827
828    #[inline]
829    pub fn set_b8(&mut self, col: usize, row: usize, val: Block8) -> errors::Result<()> {
830        let num_rows = self.num_rows;
831        let data = self.expect_b8_col(col)?;
832        let slot = data.get_mut(row).ok_or(Error::RowIndexOutOfBounds {
833            row_idx: row,
834            num_rows,
835        })?;
836
837        *slot = val.to_hardware();
838
839        Ok(())
840    }
841
842    #[inline]
843    pub fn set_b16(&mut self, col: usize, row: usize, val: Block16) -> errors::Result<()> {
844        let num_rows = self.num_rows;
845        let data = self.expect_b16_col(col)?;
846        let slot = data.get_mut(row).ok_or(Error::RowIndexOutOfBounds {
847            row_idx: row,
848            num_rows,
849        })?;
850
851        *slot = val.to_hardware();
852
853        Ok(())
854    }
855
856    #[inline]
857    pub fn set_b32(&mut self, col: usize, row: usize, val: Block32) -> errors::Result<()> {
858        let num_rows = self.num_rows;
859        let data = self.expect_b32_col(col)?;
860        let slot = data.get_mut(row).ok_or(Error::RowIndexOutOfBounds {
861            row_idx: row,
862            num_rows,
863        })?;
864
865        *slot = val.to_hardware();
866
867        Ok(())
868    }
869
870    #[inline]
871    pub fn set_b64(&mut self, col: usize, row: usize, val: Block64) -> errors::Result<()> {
872        let num_rows = self.num_rows;
873        let data = self.expect_b64_col(col)?;
874        let slot = data.get_mut(row).ok_or(Error::RowIndexOutOfBounds {
875            row_idx: row,
876            num_rows,
877        })?;
878
879        *slot = val.to_hardware();
880
881        Ok(())
882    }
883
884    #[inline]
885    pub fn set_b128(&mut self, col: usize, row: usize, val: Block128) -> errors::Result<()> {
886        let num_rows = self.num_rows;
887        let data = self.expect_b128_col(col)?;
888        let slot = data.get_mut(row).ok_or(Error::RowIndexOutOfBounds {
889            row_idx: row,
890            num_rows,
891        })?;
892
893        *slot = val.to_hardware();
894
895        Ok(())
896    }
897
898    // =========================================================
899    // Push write (sequential overwrite-at-cursor)
900    // =========================================================
901
902    #[inline]
903    pub fn push_bit(&mut self, col: usize, val: Bit) -> errors::Result<()> {
904        let row = self.cursor(col)?;
905        self.set_bit(col, row, val)?;
906
907        self.cursors[col] = row + 1;
908
909        Ok(())
910    }
911
912    #[inline]
913    pub fn push_b8(&mut self, col: usize, val: Block8) -> errors::Result<()> {
914        let row = self.cursor(col)?;
915        self.set_b8(col, row, val)?;
916
917        self.cursors[col] = row + 1;
918
919        Ok(())
920    }
921
922    #[inline]
923    pub fn push_b16(&mut self, col: usize, val: Block16) -> errors::Result<()> {
924        let row = self.cursor(col)?;
925        self.set_b16(col, row, val)?;
926
927        self.cursors[col] = row + 1;
928
929        Ok(())
930    }
931
932    #[inline]
933    pub fn push_b32(&mut self, col: usize, val: Block32) -> errors::Result<()> {
934        let row = self.cursor(col)?;
935        self.set_b32(col, row, val)?;
936
937        self.cursors[col] = row + 1;
938
939        Ok(())
940    }
941
942    #[inline]
943    pub fn push_b64(&mut self, col: usize, val: Block64) -> errors::Result<()> {
944        let row = self.cursor(col)?;
945        self.set_b64(col, row, val)?;
946
947        self.cursors[col] = row + 1;
948
949        Ok(())
950    }
951
952    #[inline]
953    pub fn push_b128(&mut self, col: usize, val: Block128) -> errors::Result<()> {
954        let row = self.cursor(col)?;
955        self.set_b128(col, row, val)?;
956
957        self.cursors[col] = row + 1;
958
959        Ok(())
960    }
961
962    // =========================================================
963    // Array column helpers
964    // =========================================================
965
966    pub fn set_bit_array(&mut self, base: usize, row: usize, values: &[Bit]) -> errors::Result<()> {
967        for (i, &val) in values.iter().enumerate() {
968            self.set_bit(base + i, row, val)?;
969        }
970
971        Ok(())
972    }
973
974    pub fn set_b8_array(
975        &mut self,
976        base: usize,
977        row: usize,
978        values: &[Block8],
979    ) -> errors::Result<()> {
980        for (i, &val) in values.iter().enumerate() {
981            self.set_b8(base + i, row, val)?;
982        }
983
984        Ok(())
985    }
986
987    pub fn set_b16_array(
988        &mut self,
989        base: usize,
990        row: usize,
991        values: &[Block16],
992    ) -> errors::Result<()> {
993        for (i, &val) in values.iter().enumerate() {
994            self.set_b16(base + i, row, val)?;
995        }
996
997        Ok(())
998    }
999
1000    pub fn set_b32_array(
1001        &mut self,
1002        base: usize,
1003        row: usize,
1004        values: &[Block32],
1005    ) -> errors::Result<()> {
1006        for (i, &val) in values.iter().enumerate() {
1007            self.set_b32(base + i, row, val)?;
1008        }
1009
1010        Ok(())
1011    }
1012
1013    pub fn set_b64_array(
1014        &mut self,
1015        base: usize,
1016        row: usize,
1017        values: &[Block64],
1018    ) -> errors::Result<()> {
1019        for (i, &val) in values.iter().enumerate() {
1020            self.set_b64(base + i, row, val)?;
1021        }
1022
1023        Ok(())
1024    }
1025
1026    pub fn set_b128_array(
1027        &mut self,
1028        base: usize,
1029        row: usize,
1030        values: &[Block128],
1031    ) -> errors::Result<()> {
1032        for (i, &val) in values.iter().enumerate() {
1033            self.set_b128(base + i, row, val)?;
1034        }
1035
1036        Ok(())
1037    }
1038
1039    // =========================================================
1040    // Selector helpers
1041    // =========================================================
1042
1043    /// Write `ONE` into rows `[0, active_rows)`
1044    /// of a `Bit` column. The tail stays zero from allocation.
1045    pub fn fill_selector(&mut self, col: usize, active_rows: usize) -> errors::Result<()> {
1046        let limit = active_rows.min(self.num_rows);
1047        let data = self.expect_bit_col(col)?;
1048
1049        for slot in data.iter_mut().take(limit) {
1050            *slot = Bit::ONE;
1051        }
1052
1053        Ok(())
1054    }
1055
1056    // =========================================================
1057    // Finalization
1058    // =========================================================
1059
1060    /// Consume the builder and return a `ColumnTrace`.
1061    /// Column order matches the schema passed to `new`.
1062    pub fn build(self) -> ColumnTrace {
1063        ColumnTrace {
1064            columns: self.columns,
1065            num_vars: self.num_vars,
1066        }
1067    }
1068
1069    // =========================================================
1070    // Internal helpers
1071    // =========================================================
1072
1073    #[inline]
1074    fn cursor(&self, col: usize) -> errors::Result<usize> {
1075        self.cursors.get(col).copied().ok_or_else(|| {
1076            Error::ColumnIndexOutOfBounds {
1077                col_idx: col,
1078                num_cols: self.columns.len(),
1079            }
1080            .into()
1081        })
1082    }
1083
1084    #[inline]
1085    fn expect_bit_col(&mut self, col: usize) -> errors::Result<&mut Vec<Bit>> {
1086        let num_cols = self.columns.len();
1087        let tc = self
1088            .columns
1089            .get_mut(col)
1090            .ok_or(Error::ColumnIndexOutOfBounds {
1091                col_idx: col,
1092                num_cols,
1093            })?;
1094
1095        match tc {
1096            TraceColumn::Bit(data) => Ok(data),
1097            other => Err(Error::ColumnTypeMismatch {
1098                col_idx: col,
1099                expected: "Bit",
1100                got: other.column_type_name(),
1101            }
1102            .into()),
1103        }
1104    }
1105
1106    #[inline]
1107    fn expect_b8_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block8>>> {
1108        let num_cols = self.columns.len();
1109        let tc = self
1110            .columns
1111            .get_mut(col)
1112            .ok_or(Error::ColumnIndexOutOfBounds {
1113                col_idx: col,
1114                num_cols,
1115            })?;
1116
1117        match tc {
1118            TraceColumn::B8(data) => Ok(data),
1119            other => Err(Error::ColumnTypeMismatch {
1120                col_idx: col,
1121                expected: "B8",
1122                got: other.column_type_name(),
1123            }
1124            .into()),
1125        }
1126    }
1127
1128    #[inline]
1129    fn expect_b16_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block16>>> {
1130        let num_cols = self.columns.len();
1131        let tc = self
1132            .columns
1133            .get_mut(col)
1134            .ok_or(Error::ColumnIndexOutOfBounds {
1135                col_idx: col,
1136                num_cols,
1137            })?;
1138
1139        match tc {
1140            TraceColumn::B16(data) => Ok(data),
1141            other => Err(Error::ColumnTypeMismatch {
1142                col_idx: col,
1143                expected: "B16",
1144                got: other.column_type_name(),
1145            }
1146            .into()),
1147        }
1148    }
1149
1150    #[inline]
1151    fn expect_b32_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block32>>> {
1152        let num_cols = self.columns.len();
1153        let tc = self
1154            .columns
1155            .get_mut(col)
1156            .ok_or(Error::ColumnIndexOutOfBounds {
1157                col_idx: col,
1158                num_cols,
1159            })?;
1160
1161        match tc {
1162            TraceColumn::B32(data) => Ok(data),
1163            other => Err(Error::ColumnTypeMismatch {
1164                col_idx: col,
1165                expected: "B32",
1166                got: other.column_type_name(),
1167            }
1168            .into()),
1169        }
1170    }
1171
1172    #[inline]
1173    fn expect_b64_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block64>>> {
1174        let num_cols = self.columns.len();
1175        let tc = self
1176            .columns
1177            .get_mut(col)
1178            .ok_or(Error::ColumnIndexOutOfBounds {
1179                col_idx: col,
1180                num_cols,
1181            })?;
1182
1183        match tc {
1184            TraceColumn::B64(data) => Ok(data),
1185            other => Err(Error::ColumnTypeMismatch {
1186                col_idx: col,
1187                expected: "B64",
1188                got: other.column_type_name(),
1189            }
1190            .into()),
1191        }
1192    }
1193
1194    #[inline]
1195    fn expect_b128_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block128>>> {
1196        let num_cols = self.columns.len();
1197        let tc = self
1198            .columns
1199            .get_mut(col)
1200            .ok_or(Error::ColumnIndexOutOfBounds {
1201                col_idx: col,
1202                num_cols,
1203            })?;
1204
1205        match tc {
1206            TraceColumn::B128(data) => Ok(data),
1207            other => Err(Error::ColumnTypeMismatch {
1208                col_idx: col,
1209                expected: "B128",
1210                got: other.column_type_name(),
1211            }
1212            .into()),
1213        }
1214    }
1215}
1216
1217impl TraceColumn {
1218    fn column_type_name(&self) -> &'static str {
1219        match self {
1220            Self::Bit(_) => "Bit",
1221            Self::B8(_) => "B8",
1222            Self::B16(_) => "B16",
1223            Self::B32(_) => "B32",
1224            Self::B64(_) => "B64",
1225            Self::B128(_) => "B128",
1226        }
1227    }
1228}
1229
1230fn num_rows_from_num_vars(num_vars: usize) -> errors::Result<usize> {
1231    let num_vars_u32 = match u32::try_from(num_vars) {
1232        Ok(v) => v,
1233        Err(_) => {
1234            return Err(Error::InvalidParameters {
1235                message: "num_vars too large",
1236            }
1237            .into());
1238        }
1239    };
1240
1241    let Some(num_rows) = 1usize.checked_shl(num_vars_u32) else {
1242        return Err(Error::InvalidParameters {
1243            message: "num_rows overflow",
1244        }
1245        .into());
1246    };
1247
1248    if num_rows == 0 {
1249        return Err(Error::InvalidParameters {
1250            message: "num_rows is zero",
1251        }
1252        .into());
1253    }
1254
1255    Ok(num_rows)
1256}
1257
1258// =================================================================
1259// UNIT TESTS
1260// =================================================================
1261
1262#[cfg(test)]
1263mod tests {
1264    use super::*;
1265    use crate::errors;
1266    use hekate_math::HardwareField;
1267
1268    fn create_mock_trace(num_vars: usize) -> ColumnTrace {
1269        ColumnTrace::new(num_vars).unwrap()
1270    }
1271
1272    #[test]
1273    fn trace_construction_basic() {
1274        let num_vars = 3;
1275        let mut trace = create_mock_trace(num_vars);
1276
1277        let col_data = vec![Block128::from(1u8); 8];
1278        trace.add_column(col_data.into_trace_column()).unwrap();
1279
1280        assert_eq!(trace.num_rows().unwrap(), 8);
1281        assert_eq!(trace.num_cols(), 1);
1282        assert_eq!(trace.num_vars, 3);
1283    }
1284
1285    #[test]
1286    fn trace_add_column_wrong_len() {
1287        let num_vars = 2;
1288        let mut trace = create_mock_trace(num_vars);
1289
1290        let col_data = vec![Block128::ZERO; 5];
1291        let err = trace
1292            .add_column(col_data.into_trace_column())
1293            .expect_err("Expected length mismatch error");
1294
1295        assert!(matches!(
1296            err,
1297            errors::Error::Trace(Error::ColumnLengthMismatch { .. })
1298        ));
1299    }
1300
1301    #[test]
1302    fn trace_get_element_mixed_types() {
1303        let num_vars = 1;
1304        let mut trace = create_mock_trace(num_vars);
1305
1306        trace
1307            .add_column(TraceColumn::Bit(vec![Bit::new(0), Bit::new(1)]))
1308            .unwrap();
1309        trace
1310            .add_column(vec![Block32::from(10u32), Block32::from(20u32)].into_trace_column())
1311            .unwrap();
1312
1313        let val0_r0: Flat<Block128> = trace.get_element(0, 0).unwrap();
1314        let val0_r1: Flat<Block128> = trace.get_element(0, 1).unwrap();
1315        let val1_r0: Flat<Block128> = trace.get_element(1, 0).unwrap();
1316        let val1_r1: Flat<Block128> = trace.get_element(1, 1).unwrap();
1317
1318        assert_eq!(val0_r0.into_raw(), Block128::ZERO);
1319        assert_eq!(val0_r1.into_raw(), Block128::ONE);
1320
1321        let expected_10 = Block128::promote_flat(Block32::from(10u32).to_hardware()).into_raw();
1322        let expected_20 = Block128::promote_flat(Block32::from(20u32).to_hardware()).into_raw();
1323
1324        assert_eq!(val1_r0.into_raw(), expected_10);
1325        assert_eq!(val1_r1.into_raw(), expected_20);
1326    }
1327
1328    #[test]
1329    fn get_element_oob_row() {
1330        let mut trace = create_mock_trace(1);
1331        trace
1332            .add_column(TraceColumn::Bit(vec![Bit::ZERO; 2]))
1333            .unwrap();
1334        trace
1335            .get_element::<Block128>(0, 2)
1336            .expect_err("Expected out-of-bounds row error");
1337    }
1338
1339    #[test]
1340    fn get_element_oob_col() {
1341        let trace = create_mock_trace(1);
1342        trace
1343            .get_element::<Block128>(0, 0)
1344            .expect_err("Expected out-of-bounds column error");
1345    }
1346
1347    // ===========================================
1348    // SAFETY & SLICE TESTS
1349    // ===========================================
1350
1351    #[test]
1352    fn get_column_slice_correct_type() {
1353        let mut trace = create_mock_trace(2);
1354        let data = vec![
1355            Block32::from(1u32),
1356            Block32::from(2u32),
1357            Block32::from(3u32),
1358            Block32::from(4u32),
1359        ];
1360        trace.add_column(data.clone().into_trace_column()).unwrap();
1361
1362        let expected_hw: Vec<Flat<Block32>> = data.into_iter().map(|x| x.to_hardware()).collect();
1363
1364        let slice: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1365        assert_eq!(slice, expected_hw.as_slice());
1366    }
1367
1368    #[test]
1369    fn get_column_slice_wrong_type() {
1370        let mut trace = create_mock_trace(1);
1371        trace
1372            .add_column(vec![Block128::ZERO; 2].into_trace_column())
1373            .unwrap();
1374
1375        trace
1376            .get_column_slice::<Flat<Block32>>(0)
1377            .expect_err("Expected column type mismatch error");
1378    }
1379
1380    #[test]
1381    fn trace_stores_hardware_basis() {
1382        let mut trace = create_mock_trace(2);
1383
1384        let tower_data = vec![
1385            Block32::from(42u32),
1386            Block32::from(13u32),
1387            Block32::from(255u32),
1388            Block32::from(1u32),
1389        ];
1390
1391        let expected_hardware: Vec<Block32> = tower_data
1392            .iter()
1393            .map(|x| x.to_hardware().into_raw())
1394            .collect();
1395
1396        trace.add_column(tower_data.into_trace_column()).unwrap();
1397
1398        let stored: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1399
1400        for (i, (&stored_val, &expected_val)) in
1401            stored.iter().zip(expected_hardware.iter()).enumerate()
1402        {
1403            assert_eq!(
1404                stored_val.into_raw(),
1405                expected_val,
1406                "Row {}: stored value {:?} != expected hardware {:?}",
1407                i,
1408                stored_val,
1409                expected_val
1410            );
1411        }
1412    }
1413
1414    #[test]
1415    fn trace_hardware_basis_homomorphism() {
1416        let mut trace = create_mock_trace(3);
1417
1418        let a_tower = vec![Block32::from(5u32); 8];
1419        let b_tower = vec![Block32::from(7u32); 8];
1420
1421        trace
1422            .add_column(a_tower.clone().into_trace_column())
1423            .unwrap();
1424        trace
1425            .add_column(b_tower.clone().into_trace_column())
1426            .unwrap();
1427
1428        let a_stored: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1429        let b_stored: &[Flat<Block32>] = trace.get_column_slice(1).unwrap();
1430
1431        let a_hw_expected = a_tower[0].to_hardware().into_raw();
1432        let b_hw_expected = b_tower[0].to_hardware().into_raw();
1433
1434        assert_eq!(a_stored[0].into_raw(), a_hw_expected);
1435        assert_eq!(b_stored[0].into_raw(), b_hw_expected);
1436
1437        let product_hw = a_stored[0] * b_stored[0];
1438        let product_expected = (a_tower[0] * b_tower[0]).to_hardware().into_raw();
1439
1440        assert_eq!(product_hw.into_raw(), product_expected);
1441    }
1442
1443    // =========================================================
1444    // TRACE BUILDER TESTS
1445    // =========================================================
1446
1447    #[test]
1448    fn trace_builder_construction_and_auto_padding() {
1449        let layout = &[ColumnType::B32, ColumnType::Bit];
1450        let tb = TraceBuilder::new(layout, 2).unwrap(); // 4 rows
1451        assert_eq!(tb.num_rows(), 4);
1452
1453        let trace = tb.build();
1454        assert_eq!(trace.num_cols(), 2);
1455        assert_eq!(trace.num_rows().unwrap(), 4);
1456
1457        // All values should be zero (auto-padding)
1458        assert!(trace.columns[0].is_all_zeros());
1459        assert!(trace.columns[1].is_all_zeros());
1460    }
1461
1462    #[test]
1463    fn trace_builder_set_b32_stores_hardware_basis() {
1464        let layout = &[ColumnType::B32];
1465        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1466
1467        tb.set_b32(0, 0, Block32::from(42u32)).unwrap();
1468        tb.set_b32(0, 1, Block32::from(13u32)).unwrap();
1469
1470        let trace = tb.build();
1471        let stored: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1472
1473        assert_eq!(stored[0], Block32::from(42u32).to_hardware());
1474        assert_eq!(stored[1], Block32::from(13u32).to_hardware());
1475    }
1476
1477    #[test]
1478    fn trace_builder_column_ordering_matches_schema() {
1479        let layout = &[ColumnType::Bit, ColumnType::B32, ColumnType::B128];
1480        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1481
1482        tb.set_bit(0, 0, Bit::ONE).unwrap();
1483        tb.set_b32(1, 0, Block32::from(99u32)).unwrap();
1484        tb.set_b128(2, 0, Block128::from(7u8)).unwrap();
1485
1486        let trace = tb.build();
1487        assert_eq!(trace.columns[0].column_type(), ColumnType::Bit);
1488        assert_eq!(trace.columns[1].column_type(), ColumnType::B32);
1489        assert_eq!(trace.columns[2].column_type(), ColumnType::B128);
1490    }
1491
1492    #[test]
1493    fn trace_builder_type_mismatch_returns_error() {
1494        let layout = &[ColumnType::Bit];
1495        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1496
1497        let err = tb.set_b32(0, 0, Block32::ZERO);
1498        assert!(err.is_err());
1499    }
1500
1501    #[test]
1502    fn trace_builder_row_out_of_bounds_returns_error() {
1503        let layout = &[ColumnType::B32];
1504        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1505
1506        let err = tb.set_b32(0, 2, Block32::ZERO);
1507        assert!(err.is_err());
1508    }
1509
1510    #[test]
1511    fn trace_builder_col_out_of_bounds_returns_error() {
1512        let layout = &[ColumnType::B32];
1513        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1514
1515        let err = tb.set_b32(1, 0, Block32::ZERO);
1516        assert!(err.is_err());
1517    }
1518
1519    #[test]
1520    fn trace_builder_fill_selector() {
1521        let layout = &[ColumnType::Bit];
1522        let mut tb = TraceBuilder::new(layout, 2).unwrap(); // 4 rows
1523
1524        tb.fill_selector(0, 3).unwrap(); // rows 0,1,2 = ONE
1525
1526        let trace = tb.build();
1527        let bits = trace.columns[0].as_bit_slice().unwrap();
1528        assert_eq!(bits[0], Bit::ONE);
1529        assert_eq!(bits[1], Bit::ONE);
1530        assert_eq!(bits[2], Bit::ONE);
1531        assert_eq!(bits[3], Bit::ZERO); // padding
1532    }
1533
1534    #[test]
1535    fn trace_builder_push_mode() {
1536        let layout = &[ColumnType::B32, ColumnType::Bit];
1537        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1538
1539        tb.push_b32(0, Block32::from(10u32)).unwrap();
1540        tb.push_b32(0, Block32::from(20u32)).unwrap();
1541        tb.push_bit(1, Bit::ONE).unwrap();
1542        // row 1 of col 1 stays zero (auto-pad)
1543
1544        let trace = tb.build();
1545        let b32s: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1546        assert_eq!(b32s[0], Block32::from(10u32).to_hardware());
1547        assert_eq!(b32s[1], Block32::from(20u32).to_hardware());
1548
1549        let bits = trace.columns[1].as_bit_slice().unwrap();
1550        assert_eq!(bits[0], Bit::ONE);
1551        assert_eq!(bits[1], Bit::ZERO);
1552    }
1553
1554    #[test]
1555    fn trace_builder_set_b32_array() {
1556        let layout = &[ColumnType::B32, ColumnType::B32, ColumnType::B32];
1557        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1558
1559        let vals = [
1560            Block32::from(1u32),
1561            Block32::from(2u32),
1562            Block32::from(3u32),
1563        ];
1564        tb.set_b32_array(0, 0, &vals).unwrap();
1565
1566        let trace = tb.build();
1567        for (i, &expected) in vals.iter().enumerate() {
1568            let stored: &[Flat<Block32>] = trace.get_column_slice(i).unwrap();
1569            assert_eq!(stored[0], expected.to_hardware());
1570        }
1571    }
1572
1573    #[test]
1574    fn trace_builder_set_b8() {
1575        let layout = &[ColumnType::B8];
1576        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1577
1578        tb.set_b8(0, 0, Block8(0xAB)).unwrap();
1579        tb.set_b8(0, 1, Block8(0xCD)).unwrap();
1580
1581        let trace = tb.build();
1582        let stored: &[Flat<Block8>] = trace.get_column_slice(0).unwrap();
1583        assert_eq!(stored[0], Block8(0xAB).to_hardware());
1584        assert_eq!(stored[1], Block8(0xCD).to_hardware());
1585    }
1586
1587    #[test]
1588    fn trace_builder_set_b16() {
1589        let layout = &[ColumnType::B16];
1590        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1591
1592        tb.set_b16(0, 0, Block16(1000)).unwrap();
1593        tb.set_b16(0, 1, Block16(2000)).unwrap();
1594
1595        let trace = tb.build();
1596        let stored: &[Flat<Block16>] = trace.get_column_slice(0).unwrap();
1597        assert_eq!(stored[0], Block16(1000).to_hardware());
1598        assert_eq!(stored[1], Block16(2000).to_hardware());
1599    }
1600
1601    #[test]
1602    fn trace_builder_set_b64() {
1603        let layout = &[ColumnType::B64];
1604        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1605
1606        tb.set_b64(0, 0, Block64(0xDEADBEEF_CAFEBABE)).unwrap();
1607
1608        let trace = tb.build();
1609        let stored: &[Flat<Block64>] = trace.get_column_slice(0).unwrap();
1610        assert_eq!(stored[0], Block64(0xDEADBEEF_CAFEBABE).to_hardware());
1611        assert_eq!(stored[1], Block64::ZERO.to_hardware()); // auto-padding
1612    }
1613
1614    #[test]
1615    fn trace_builder_set_b128() {
1616        let layout = &[ColumnType::B128];
1617        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1618
1619        let val = Block128::from(0xFFu8);
1620        tb.set_b128(0, 0, val).unwrap();
1621
1622        let trace = tb.build();
1623        let stored: &[Flat<Block128>] = trace.get_column_slice(0).unwrap();
1624        assert_eq!(stored[0], val.to_hardware());
1625    }
1626
1627    #[test]
1628    fn trace_builder_push_all_types() {
1629        let layout = &[
1630            ColumnType::Bit,
1631            ColumnType::B8,
1632            ColumnType::B16,
1633            ColumnType::B32,
1634            ColumnType::B64,
1635            ColumnType::B128,
1636        ];
1637        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1638
1639        tb.push_bit(0, Bit::ONE).unwrap();
1640        tb.push_b8(1, Block8(0x42)).unwrap();
1641        tb.push_b16(2, Block16(1234)).unwrap();
1642        tb.push_b32(3, Block32::from(5678u32)).unwrap();
1643        tb.push_b64(4, Block64(9999)).unwrap();
1644        tb.push_b128(5, Block128::from(77u8)).unwrap();
1645
1646        let trace = tb.build();
1647        assert_eq!(trace.num_cols(), 6);
1648
1649        let bits = trace.columns[0].as_bit_slice().unwrap();
1650        assert_eq!(bits[0], Bit::ONE);
1651        assert_eq!(bits[1], Bit::ZERO); // auto-pad
1652
1653        let b8s: &[Flat<Block8>] = trace.get_column_slice(1).unwrap();
1654        assert_eq!(b8s[0], Block8(0x42).to_hardware());
1655
1656        let b16s: &[Flat<Block16>] = trace.get_column_slice(2).unwrap();
1657        assert_eq!(b16s[0], Block16(1234).to_hardware());
1658
1659        let b32s: &[Flat<Block32>] = trace.get_column_slice(3).unwrap();
1660        assert_eq!(b32s[0], Block32::from(5678u32).to_hardware());
1661
1662        let b64s: &[Flat<Block64>] = trace.get_column_slice(4).unwrap();
1663        assert_eq!(b64s[0], Block64(9999).to_hardware());
1664
1665        let b128s: &[Flat<Block128>] = trace.get_column_slice(5).unwrap();
1666        assert_eq!(b128s[0], Block128::from(77u8).to_hardware());
1667    }
1668
1669    #[test]
1670    fn trace_builder_type_mismatch_all_setters() {
1671        // B32 column, every non-B32 setter should fail
1672        let layout = &[ColumnType::B32];
1673        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1674
1675        assert!(tb.set_bit(0, 0, Bit::ONE).is_err());
1676        assert!(tb.set_b8(0, 0, Block8(1)).is_err());
1677        assert!(tb.set_b16(0, 0, Block16(1)).is_err());
1678        assert!(tb.set_b64(0, 0, Block64(1)).is_err());
1679        assert!(tb.set_b128(0, 0, Block128::ONE).is_err());
1680
1681        // Correct type succeeds
1682        assert!(tb.set_b32(0, 0, Block32::ONE).is_ok());
1683    }
1684
1685    #[test]
1686    fn trace_builder_array_setters_all_types() {
1687        let layout = &[
1688            ColumnType::Bit,
1689            ColumnType::Bit,
1690            ColumnType::B8,
1691            ColumnType::B8,
1692            ColumnType::B64,
1693            ColumnType::B64,
1694            ColumnType::B128,
1695            ColumnType::B128,
1696        ];
1697        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1698
1699        tb.set_bit_array(0, 0, &[Bit::ONE, Bit::ZERO]).unwrap();
1700        tb.set_b8_array(2, 0, &[Block8(10), Block8(20)]).unwrap();
1701        tb.set_b64_array(4, 0, &[Block64(100), Block64(200)])
1702            .unwrap();
1703        tb.set_b128_array(6, 0, &[Block128::ONE, Block128::from(2u8)])
1704            .unwrap();
1705
1706        let trace = tb.build();
1707
1708        let bits = trace.columns[0].as_bit_slice().unwrap();
1709        assert_eq!(bits[0], Bit::ONE);
1710
1711        let bits1 = trace.columns[1].as_bit_slice().unwrap();
1712        assert_eq!(bits1[0], Bit::ZERO);
1713
1714        let b8s: &[Flat<Block8>] = trace.get_column_slice(2).unwrap();
1715        assert_eq!(b8s[0], Block8(10).to_hardware());
1716
1717        let b8s1: &[Flat<Block8>] = trace.get_column_slice(3).unwrap();
1718        assert_eq!(b8s1[0], Block8(20).to_hardware());
1719    }
1720
1721    #[test]
1722    fn trace_builder_invalid_num_vars() {
1723        let layout = &[ColumnType::B32];
1724        // num_vars too large to shift
1725        assert!(TraceBuilder::new(layout, 128).is_err());
1726    }
1727}