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))` on
333    /// every `Bit` used as selector, constraint operand, or
334    /// LogUp source: the commitment binds a cell only to GF(2^32).
335    Bit,
336    B8,
337    B16,
338    B32,
339    B64,
340    B128,
341}
342
343impl ColumnType {
344    #[inline]
345    pub const fn byte_size(&self) -> usize {
346        match self {
347            Self::Bit => 1,
348            Self::B8 => 1,
349            Self::B16 => 2,
350            Self::B32 => 4,
351            Self::B64 => 8,
352            Self::B128 => 16,
353        }
354    }
355
356    /// The tower field this column's cells are committed
357    /// in under the MDS Reed-Solomon row code. GF(2),
358    /// GF(2^8), and GF(2^16) admit no rate-1/2 MDS RS.
359    #[inline]
360    pub const fn rs_field(&self) -> Self {
361        match self {
362            Self::Bit | Self::B8 | Self::B16 => Self::B32,
363            other => *other,
364        }
365    }
366
367    /// Parse a field element from its on-wire
368    /// bytes (little-endian, hardware basis)
369    /// without intermediate allocation.
370    pub fn parse_from_bytes<F>(&self, bytes: &[u8]) -> Flat<F>
371    where
372        F: TraceCompatibleField,
373    {
374        match self {
375            Self::Bit => Flat::from_raw(F::from(Bit::new(bytes[0]))),
376            Self::B8 => F::promote_flat(Flat::from_raw(Block8(bytes[0]))),
377            Self::B16 => {
378                let mut buf = [0u8; 2];
379                buf.copy_from_slice(&bytes[0..2]);
380
381                F::promote_flat(Flat::from_raw(Block16(u16::from_le_bytes(buf))))
382            }
383            Self::B32 => {
384                let mut buf = [0u8; 4];
385                buf.copy_from_slice(&bytes[0..4]);
386
387                F::promote_flat(Flat::from_raw(Block32(u32::from_le_bytes(buf))))
388            }
389            Self::B64 => {
390                let mut buf = [0u8; 8];
391                buf.copy_from_slice(&bytes[0..8]);
392
393                F::promote_flat(Flat::from_raw(Block64(u64::from_le_bytes(buf))))
394            }
395            Self::B128 => {
396                let mut buf = [0u8; 16];
397                buf.copy_from_slice(&bytes[0..16]);
398
399                F::promote_flat(Flat::from_raw(Block128(u128::from_le_bytes(buf))))
400            }
401        }
402    }
403}
404
405/// One typed column of the execution trace.
406/// Stored in hardware (flat) basis for all
407/// non-`Bit` variants.
408#[derive(Clone, Debug, Zeroize)]
409#[cfg_attr(feature = "secure-memory", derive(ZeroizeOnDrop))]
410pub enum TraceColumn {
411    Bit(Vec<Bit>),
412    B8(Vec<Flat<Block8>>),
413    B16(Vec<Flat<Block16>>),
414    B32(Vec<Flat<Block32>>),
415    B64(Vec<Flat<Block64>>),
416    B128(Vec<Flat<Block128>>),
417}
418
419impl TraceColumn {
420    /// Narrow a vector of `Block128` into a column
421    /// of `target_type` by truncating to the low bytes.
422    pub fn from_data(data: Vec<Block128>, target_type: ColumnType) -> Self {
423        match target_type {
424            ColumnType::Bit => {
425                let converted: Vec<Bit> = data
426                    .iter()
427                    .map(|val| {
428                        let bytes = val.to_bytes();
429                        Bit::from(bytes[0] & 1)
430                    })
431                    .collect();
432                TraceColumn::Bit(converted)
433            }
434            ColumnType::B8 => {
435                let converted: Vec<Flat<Block8>> = data
436                    .iter()
437                    .map(|val| {
438                        let bytes = val.to_bytes();
439                        Block8::from(bytes[0]).to_hardware()
440                    })
441                    .collect();
442                TraceColumn::B8(converted)
443            }
444            ColumnType::B16 => {
445                let converted: Vec<Flat<Block16>> = data
446                    .iter()
447                    .map(|val| {
448                        let bytes = val.to_bytes();
449                        let mut chunk = [0u8; 2];
450                        chunk.copy_from_slice(&bytes[0..2]);
451
452                        Block16::from(u16::from_le_bytes(chunk)).to_hardware()
453                    })
454                    .collect();
455                TraceColumn::B16(converted)
456            }
457            ColumnType::B32 => {
458                let converted: Vec<Flat<Block32>> = data
459                    .iter()
460                    .map(|val| {
461                        let bytes = val.to_bytes();
462                        let mut chunk = [0u8; 4];
463                        chunk.copy_from_slice(&bytes[0..4]);
464
465                        Block32::from(u32::from_le_bytes(chunk)).to_hardware()
466                    })
467                    .collect();
468                TraceColumn::B32(converted)
469            }
470            ColumnType::B64 => {
471                let converted: Vec<Flat<Block64>> = data
472                    .iter()
473                    .map(|val| {
474                        let bytes = val.to_bytes();
475                        let mut chunk = [0u8; 8];
476                        chunk.copy_from_slice(&bytes[0..8]);
477
478                        Block64::from(u64::from_le_bytes(chunk)).to_hardware()
479                    })
480                    .collect();
481                TraceColumn::B64(converted)
482            }
483            ColumnType::B128 => {
484                TraceColumn::B128(data.into_iter().map(|value| value.to_hardware()).collect())
485            }
486        }
487    }
488
489    pub fn len(&self) -> usize {
490        match self {
491            Self::Bit(v) => v.len(),
492            Self::B8(v) => v.len(),
493            Self::B16(v) => v.len(),
494            Self::B32(v) => v.len(),
495            Self::B64(v) => v.len(),
496            Self::B128(v) => v.len(),
497        }
498    }
499
500    pub fn is_empty(&self) -> bool {
501        self.len() == 0
502    }
503
504    pub fn is_all_zeros(&self) -> bool {
505        match self {
506            Self::Bit(v) => v.iter().all(|x| x.get() == 0),
507            Self::B8(v) => v.iter().all(|x| x.into_raw().0 == 0),
508            Self::B16(v) => v.iter().all(|x| x.into_raw().0 == 0),
509            Self::B32(v) => v.iter().all(|x| x.into_raw().0 == 0),
510            Self::B64(v) => v.iter().all(|x| x.into_raw().0 == 0),
511            Self::B128(v) => v.iter().all(|x| x.into_raw() == Block128::ZERO),
512        }
513    }
514
515    /// The `ColumnType` tag matching this variant.
516    pub fn column_type(&self) -> ColumnType {
517        match self {
518            Self::Bit(_) => ColumnType::Bit,
519            Self::B8(_) => ColumnType::B8,
520            Self::B16(_) => ColumnType::B16,
521            Self::B32(_) => ColumnType::B32,
522            Self::B64(_) => ColumnType::B64,
523            Self::B128(_) => ColumnType::B128,
524        }
525    }
526
527    /// Append this row's little-endian
528    /// serialization to `buf`. Used for
529    /// Merkle leaf hashing and LDT opening bytes.
530    pub fn append_bytes_at(&self, row_idx: usize, buf: &mut Vec<u8>) {
531        match self {
532            Self::Bit(v) => {
533                buf.push(v[row_idx].get());
534            }
535            Self::B8(v) => {
536                buf.push(v[row_idx].into_raw().0);
537            }
538            Self::B16(v) => {
539                buf.extend_from_slice(&v[row_idx].into_raw().0.to_le_bytes());
540            }
541            Self::B32(v) => {
542                buf.extend_from_slice(&v[row_idx].into_raw().0.to_le_bytes());
543            }
544            Self::B64(v) => {
545                buf.extend_from_slice(&v[row_idx].into_raw().0.to_le_bytes());
546            }
547            Self::B128(v) => {
548                buf.extend_from_slice(&v[row_idx].into_raw().0.to_le_bytes());
549            }
550        }
551    }
552
553    // ===========================================
554    // Typed slice accessors
555    // ===========================================
556
557    pub fn as_bit_slice(&self) -> Option<&[Bit]> {
558        if let Self::Bit(v) = self {
559            Some(v)
560        } else {
561            None
562        }
563    }
564
565    pub fn as_b8_slice(&self) -> Option<&[Flat<Block8>]> {
566        if let Self::B8(v) = self {
567            Some(v)
568        } else {
569            None
570        }
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
581    pub fn as_b32_slice(&self) -> Option<&[Flat<Block32>]> {
582        if let Self::B32(v) = self {
583            Some(v)
584        } else {
585            None
586        }
587    }
588
589    pub fn as_b64_slice(&self) -> Option<&[Flat<Block64>]> {
590        if let Self::B64(v) = self {
591            Some(v)
592        } else {
593            None
594        }
595    }
596
597    pub fn as_b128_slice(&self) -> Option<&[Flat<Block128>]> {
598        if let Self::B128(v) = self {
599            Some(v)
600        } else {
601            None
602        }
603    }
604}
605
606/// A concrete implementation of
607/// Trace using column-major storage.
608#[derive(Clone, Debug, Zeroize)]
609pub struct ColumnTrace {
610    pub columns: Vec<TraceColumn>,
611    pub num_vars: usize,
612
613    #[zeroize(skip)]
614    secret: bool,
615}
616
617impl Drop for ColumnTrace {
618    fn drop(&mut self) {
619        if self.secret || cfg!(feature = "secure-memory") {
620            self.columns.zeroize();
621        }
622    }
623}
624
625#[cfg(feature = "secure-memory")]
626impl ZeroizeOnDrop for ColumnTrace {}
627
628impl Trace for ColumnTrace {
629    fn num_vars(&self) -> usize {
630        self.num_vars
631    }
632
633    fn columns(&self) -> &[TraceColumn] {
634        &self.columns
635    }
636}
637
638impl ColumnTrace {
639    pub fn new(num_vars: usize) -> errors::Result<Self> {
640        let num_rows = num_rows_from_num_vars(num_vars)?;
641        if num_rows == 0 {
642            return Err(Error::InvalidParameters {
643                message: "trace height is zero",
644            }
645            .into());
646        }
647
648        Ok(Self {
649            num_vars,
650            columns: Vec::new(),
651            secret: false,
652        })
653    }
654
655    /// An empty trace whose columns zeroize on drop.
656    pub fn new_secret(num_vars: usize) -> errors::Result<Self> {
657        let mut trace = Self::new(num_vars)?;
658        trace.secret = true;
659
660        Ok(trace)
661    }
662
663    /// Consume the trace and return its owned column storage.
664    pub fn into_columns(mut self) -> Vec<TraceColumn> {
665        core::mem::take(&mut self.columns)
666    }
667
668    pub fn add_column(&mut self, col: TraceColumn) -> errors::Result<()> {
669        let expected_len = self.num_rows()?;
670        let got_len = col.len();
671
672        if got_len != expected_len {
673            return Err(Error::ColumnLengthMismatch {
674                expected_len,
675                got_len,
676            }
677            .into());
678        }
679
680        self.columns.push(col);
681
682        Ok(())
683    }
684}
685
686pub trait IntoTraceColumn {
687    fn into_trace_column(self) -> TraceColumn;
688}
689
690impl IntoTraceColumn for Vec<Bit> {
691    fn into_trace_column(self) -> TraceColumn {
692        TraceColumn::Bit(self)
693    }
694}
695
696impl IntoTraceColumn for Vec<Block8> {
697    fn into_trace_column(self) -> TraceColumn {
698        TraceColumn::B8(self.into_iter().map(|value| value.to_hardware()).collect())
699    }
700}
701
702impl IntoTraceColumn for Vec<Block16> {
703    fn into_trace_column(self) -> TraceColumn {
704        TraceColumn::B16(self.into_iter().map(|value| value.to_hardware()).collect())
705    }
706}
707
708impl IntoTraceColumn for Vec<Block32> {
709    fn into_trace_column(self) -> TraceColumn {
710        TraceColumn::B32(self.into_iter().map(|value| value.to_hardware()).collect())
711    }
712}
713
714impl IntoTraceColumn for Vec<Block64> {
715    fn into_trace_column(self) -> TraceColumn {
716        TraceColumn::B64(self.into_iter().map(|value| value.to_hardware()).collect())
717    }
718}
719
720impl IntoTraceColumn for Vec<Block128> {
721    fn into_trace_column(self) -> TraceColumn {
722        TraceColumn::B128(self.into_iter().map(|value| value.to_hardware()).collect())
723    }
724}
725
726impl IntoTraceColumn for Vec<Flat<Block8>> {
727    fn into_trace_column(self) -> TraceColumn {
728        TraceColumn::B8(self)
729    }
730}
731
732impl IntoTraceColumn for Vec<Flat<Block16>> {
733    fn into_trace_column(self) -> TraceColumn {
734        TraceColumn::B16(self)
735    }
736}
737
738impl IntoTraceColumn for Vec<Flat<Block32>> {
739    fn into_trace_column(self) -> TraceColumn {
740        TraceColumn::B32(self)
741    }
742}
743
744impl IntoTraceColumn for Vec<Flat<Block64>> {
745    fn into_trace_column(self) -> TraceColumn {
746        TraceColumn::B64(self)
747    }
748}
749
750impl IntoTraceColumn for Vec<Flat<Block128>> {
751    fn into_trace_column(self) -> TraceColumn {
752        TraceColumn::B128(self)
753    }
754}
755
756/// Zero-copy byte views `(ptr, elem_width)` for every column in a trace.
757/// Centralizes the `#[repr(transparent)]`-dependent pointer casts.
758pub fn get_col_views(columns: &[TraceColumn]) -> Vec<(&[u8], usize)> {
759    columns
760        .iter()
761        .map(|col| match col {
762            TraceColumn::Bit(v) => (
763                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len()) },
764                1,
765            ),
766            TraceColumn::B8(v) => (
767                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len()) },
768                1,
769            ),
770            TraceColumn::B16(v) => (
771                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 2) },
772                2,
773            ),
774            TraceColumn::B32(v) => (
775                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4) },
776                4,
777            ),
778            TraceColumn::B64(v) => (
779                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 8) },
780                8,
781            ),
782            TraceColumn::B128(v) => (
783                unsafe { core::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 16) },
784                16,
785            ),
786        })
787        .collect()
788}
789
790// =========================================================
791// TRACE BUILDER
792// =========================================================
793
794macro_rules! impl_width_writes {
795    (
796        $(
797            $block:ty {
798                set: $set:ident,
799                set_array: $set_array:ident,
800                push: $push:ident,
801                set_flat: $set_flat:ident,
802                set_array_flat: $set_array_flat:ident,
803                column: $expect:ident,
804            }
805        )*
806    ) => {
807        $(
808            #[inline]
809            pub fn $set(&mut self, col: usize, row: usize, val: $block) -> errors::Result<()> {
810                self.$set_flat(col, row, val.to_hardware())
811            }
812
813            #[inline]
814            pub fn $set_flat(
815                &mut self,
816                col: usize,
817                row: usize,
818                val: Flat<$block>,
819            ) -> errors::Result<()> {
820                let num_rows = self.num_rows;
821                let data = self.$expect(col)?;
822                let slot = data.get_mut(row).ok_or(Error::RowIndexOutOfBounds {
823                    row_idx: row,
824                    num_rows,
825                })?;
826
827                *slot = val;
828
829                Ok(())
830            }
831
832            #[inline]
833            pub fn $push(&mut self, col: usize, val: $block) -> errors::Result<()> {
834                let row = self.cursor(col)?;
835                self.$set(col, row, val)?;
836
837                self.cursors[col] = row + 1;
838
839                Ok(())
840            }
841
842            pub fn $set_array(
843                &mut self,
844                base: usize,
845                row: usize,
846                values: &[$block],
847            ) -> errors::Result<()> {
848                for (i, &val) in values.iter().enumerate() {
849                    self.$set(base + i, row, val)?;
850                }
851
852                Ok(())
853            }
854
855            pub fn $set_array_flat(
856                &mut self,
857                base: usize,
858                row: usize,
859                values: &[Flat<$block>],
860            ) -> errors::Result<()> {
861                for (i, &val) in values.iter().enumerate() {
862                    self.$set_flat(base + i, row, val)?;
863                }
864
865                Ok(())
866            }
867        )*
868    };
869}
870
871/// Schema-driven builder. Every column
872/// is allocated zero-filled from the layout;
873/// unfilled rows stay zero, so padding is implicit.
874pub struct TraceBuilder {
875    columns: Vec<TraceColumn>,
876    num_vars: usize,
877    num_rows: usize,
878    cursors: Vec<usize>,
879    secret: bool,
880}
881
882impl Drop for TraceBuilder {
883    fn drop(&mut self) {
884        if self.secret || cfg!(feature = "secure-memory") {
885            self.columns.zeroize();
886        }
887    }
888}
889
890impl TraceBuilder {
891    pub fn new(layout: &[ColumnType], num_vars: usize) -> errors::Result<Self> {
892        Self::with_secrecy(layout, num_vars, false)
893    }
894
895    /// A builder whose columns and the trace it builds
896    /// zeroize on drop: for witnesses that carry secrets.
897    pub fn new_secret(layout: &[ColumnType], num_vars: usize) -> errors::Result<Self> {
898        Self::with_secrecy(layout, num_vars, true)
899    }
900
901    fn with_secrecy(layout: &[ColumnType], num_vars: usize, secret: bool) -> errors::Result<Self> {
902        let num_rows = num_rows_from_num_vars(num_vars)?;
903        let columns = layout
904            .iter()
905            .map(|ct| match ct {
906                ColumnType::Bit => TraceColumn::Bit(vec![Bit::ZERO; num_rows]),
907                ColumnType::B8 => TraceColumn::B8(vec![Block8::ZERO.to_hardware(); num_rows]),
908                ColumnType::B16 => TraceColumn::B16(vec![Block16::ZERO.to_hardware(); num_rows]),
909                ColumnType::B32 => TraceColumn::B32(vec![Block32::ZERO.to_hardware(); num_rows]),
910                ColumnType::B64 => TraceColumn::B64(vec![Block64::ZERO.to_hardware(); num_rows]),
911                ColumnType::B128 => TraceColumn::B128(vec![Block128::ZERO.to_hardware(); num_rows]),
912            })
913            .collect();
914
915        Ok(Self {
916            columns,
917            num_vars,
918            num_rows,
919            cursors: vec![0; layout.len()],
920            secret,
921        })
922    }
923
924    /// `2^num_vars`.
925    #[inline]
926    pub fn num_rows(&self) -> usize {
927        self.num_rows
928    }
929
930    // =========================================================
931    // Column writes: indexed, cursor-push, array
932    // =========================================================
933
934    #[inline]
935    pub fn set_bit(&mut self, col: usize, row: usize, val: Bit) -> errors::Result<()> {
936        let num_rows = self.num_rows;
937        let data = self.expect_bit_col(col)?;
938        let slot = data.get_mut(row).ok_or(Error::RowIndexOutOfBounds {
939            row_idx: row,
940            num_rows,
941        })?;
942
943        *slot = val;
944
945        Ok(())
946    }
947
948    pub fn set_bit_array(&mut self, base: usize, row: usize, values: &[Bit]) -> errors::Result<()> {
949        for (i, &val) in values.iter().enumerate() {
950            self.set_bit(base + i, row, val)?;
951        }
952
953        Ok(())
954    }
955
956    #[inline]
957    pub fn push_bit(&mut self, col: usize, val: Bit) -> errors::Result<()> {
958        let row = self.cursor(col)?;
959        self.set_bit(col, row, val)?;
960
961        self.cursors[col] = row + 1;
962
963        Ok(())
964    }
965
966    impl_width_writes! {
967        Block8 {
968            set: set_b8,
969            set_array: set_b8_array,
970            push: push_b8,
971            set_flat: set_b8_flat,
972            set_array_flat: set_b8_array_flat,
973            column: expect_b8_col,
974        }
975
976        Block16 {
977            set: set_b16,
978            set_array: set_b16_array,
979            push: push_b16,
980            set_flat: set_b16_flat,
981            set_array_flat: set_b16_array_flat,
982            column: expect_b16_col,
983        }
984
985        Block32 {
986            set: set_b32,
987            set_array: set_b32_array,
988            push: push_b32,
989            set_flat: set_b32_flat,
990            set_array_flat: set_b32_array_flat,
991            column: expect_b32_col,
992        }
993
994        Block64 {
995            set: set_b64,
996            set_array: set_b64_array,
997            push: push_b64,
998            set_flat: set_b64_flat,
999            set_array_flat: set_b64_array_flat,
1000            column: expect_b64_col,
1001        }
1002
1003        Block128 {
1004            set: set_b128,
1005            set_array: set_b128_array,
1006            push: push_b128,
1007            set_flat: set_b128_flat,
1008            set_array_flat: set_b128_array_flat,
1009            column: expect_b128_col,
1010        }
1011    }
1012
1013    // =========================================================
1014    // Selector helpers
1015    // =========================================================
1016
1017    /// Write `ONE` into rows `[0, active_rows)`
1018    /// of a `Bit` column. The tail stays zero from allocation.
1019    pub fn fill_selector(&mut self, col: usize, active_rows: usize) -> errors::Result<()> {
1020        let limit = active_rows.min(self.num_rows);
1021        let data = self.expect_bit_col(col)?;
1022
1023        for slot in data.iter_mut().take(limit) {
1024            *slot = Bit::ONE;
1025        }
1026
1027        Ok(())
1028    }
1029
1030    // =========================================================
1031    // Finalization
1032    // =========================================================
1033
1034    /// Consume the builder and return a `ColumnTrace`.
1035    /// Column order matches the schema passed to `new`.
1036    pub fn build(mut self) -> ColumnTrace {
1037        ColumnTrace {
1038            columns: core::mem::take(&mut self.columns),
1039            num_vars: self.num_vars,
1040            secret: self.secret,
1041        }
1042    }
1043
1044    // =========================================================
1045    // Internal helpers
1046    // =========================================================
1047
1048    #[inline]
1049    fn cursor(&self, col: usize) -> errors::Result<usize> {
1050        self.cursors.get(col).copied().ok_or_else(|| {
1051            Error::ColumnIndexOutOfBounds {
1052                col_idx: col,
1053                num_cols: self.columns.len(),
1054            }
1055            .into()
1056        })
1057    }
1058
1059    #[inline]
1060    fn expect_bit_col(&mut self, col: usize) -> errors::Result<&mut Vec<Bit>> {
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::Bit(data) => Ok(data),
1072            other => Err(Error::ColumnTypeMismatch {
1073                col_idx: col,
1074                expected: "Bit",
1075                got: other.column_type_name(),
1076            }
1077            .into()),
1078        }
1079    }
1080
1081    #[inline]
1082    fn expect_b8_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block8>>> {
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::B8(data) => Ok(data),
1094            other => Err(Error::ColumnTypeMismatch {
1095                col_idx: col,
1096                expected: "B8",
1097                got: other.column_type_name(),
1098            }
1099            .into()),
1100        }
1101    }
1102
1103    #[inline]
1104    fn expect_b16_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block16>>> {
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::B16(data) => Ok(data),
1116            other => Err(Error::ColumnTypeMismatch {
1117                col_idx: col,
1118                expected: "B16",
1119                got: other.column_type_name(),
1120            }
1121            .into()),
1122        }
1123    }
1124
1125    #[inline]
1126    fn expect_b32_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block32>>> {
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::B32(data) => Ok(data),
1138            other => Err(Error::ColumnTypeMismatch {
1139                col_idx: col,
1140                expected: "B32",
1141                got: other.column_type_name(),
1142            }
1143            .into()),
1144        }
1145    }
1146
1147    #[inline]
1148    fn expect_b64_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block64>>> {
1149        let num_cols = self.columns.len();
1150        let tc = self
1151            .columns
1152            .get_mut(col)
1153            .ok_or(Error::ColumnIndexOutOfBounds {
1154                col_idx: col,
1155                num_cols,
1156            })?;
1157
1158        match tc {
1159            TraceColumn::B64(data) => Ok(data),
1160            other => Err(Error::ColumnTypeMismatch {
1161                col_idx: col,
1162                expected: "B64",
1163                got: other.column_type_name(),
1164            }
1165            .into()),
1166        }
1167    }
1168
1169    #[inline]
1170    fn expect_b128_col(&mut self, col: usize) -> errors::Result<&mut Vec<Flat<Block128>>> {
1171        let num_cols = self.columns.len();
1172        let tc = self
1173            .columns
1174            .get_mut(col)
1175            .ok_or(Error::ColumnIndexOutOfBounds {
1176                col_idx: col,
1177                num_cols,
1178            })?;
1179
1180        match tc {
1181            TraceColumn::B128(data) => Ok(data),
1182            other => Err(Error::ColumnTypeMismatch {
1183                col_idx: col,
1184                expected: "B128",
1185                got: other.column_type_name(),
1186            }
1187            .into()),
1188        }
1189    }
1190}
1191
1192impl TraceColumn {
1193    fn column_type_name(&self) -> &'static str {
1194        match self {
1195            Self::Bit(_) => "Bit",
1196            Self::B8(_) => "B8",
1197            Self::B16(_) => "B16",
1198            Self::B32(_) => "B32",
1199            Self::B64(_) => "B64",
1200            Self::B128(_) => "B128",
1201        }
1202    }
1203}
1204
1205fn num_rows_from_num_vars(num_vars: usize) -> errors::Result<usize> {
1206    let num_vars_u32 = match u32::try_from(num_vars) {
1207        Ok(v) => v,
1208        Err(_) => {
1209            return Err(Error::InvalidParameters {
1210                message: "num_vars too large",
1211            }
1212            .into());
1213        }
1214    };
1215
1216    let Some(num_rows) = 1usize.checked_shl(num_vars_u32) else {
1217        return Err(Error::InvalidParameters {
1218            message: "num_rows overflow",
1219        }
1220        .into());
1221    };
1222
1223    if num_rows == 0 {
1224        return Err(Error::InvalidParameters {
1225            message: "num_rows is zero",
1226        }
1227        .into());
1228    }
1229
1230    Ok(num_rows)
1231}
1232
1233// =================================================================
1234// UNIT TESTS
1235// =================================================================
1236
1237#[cfg(test)]
1238mod tests {
1239    use super::*;
1240    use crate::errors;
1241    use hekate_math::HardwareField;
1242
1243    fn create_mock_trace(num_vars: usize) -> ColumnTrace {
1244        ColumnTrace::new(num_vars).unwrap()
1245    }
1246
1247    #[test]
1248    fn trace_construction_basic() {
1249        let num_vars = 3;
1250        let mut trace = create_mock_trace(num_vars);
1251
1252        let col_data = vec![Block128::from(1u8); 8];
1253        trace.add_column(col_data.into_trace_column()).unwrap();
1254
1255        assert_eq!(trace.num_rows().unwrap(), 8);
1256        assert_eq!(trace.num_cols(), 1);
1257        assert_eq!(trace.num_vars, 3);
1258    }
1259
1260    #[test]
1261    fn trace_add_column_wrong_len() {
1262        let num_vars = 2;
1263        let mut trace = create_mock_trace(num_vars);
1264
1265        let col_data = vec![Block128::ZERO; 5];
1266        let err = trace
1267            .add_column(col_data.into_trace_column())
1268            .expect_err("Expected length mismatch error");
1269
1270        assert!(matches!(
1271            err,
1272            errors::Error::Trace(Error::ColumnLengthMismatch { .. })
1273        ));
1274    }
1275
1276    #[test]
1277    fn trace_get_element_mixed_types() {
1278        let num_vars = 1;
1279        let mut trace = create_mock_trace(num_vars);
1280
1281        trace
1282            .add_column(TraceColumn::Bit(vec![Bit::new(0), Bit::new(1)]))
1283            .unwrap();
1284        trace
1285            .add_column(vec![Block32::from(10u32), Block32::from(20u32)].into_trace_column())
1286            .unwrap();
1287
1288        let val0_r0: Flat<Block128> = trace.get_element(0, 0).unwrap();
1289        let val0_r1: Flat<Block128> = trace.get_element(0, 1).unwrap();
1290        let val1_r0: Flat<Block128> = trace.get_element(1, 0).unwrap();
1291        let val1_r1: Flat<Block128> = trace.get_element(1, 1).unwrap();
1292
1293        assert_eq!(val0_r0.into_raw(), Block128::ZERO);
1294        assert_eq!(val0_r1.into_raw(), Block128::ONE);
1295
1296        let expected_10 = Block128::promote_flat(Block32::from(10u32).to_hardware()).into_raw();
1297        let expected_20 = Block128::promote_flat(Block32::from(20u32).to_hardware()).into_raw();
1298
1299        assert_eq!(val1_r0.into_raw(), expected_10);
1300        assert_eq!(val1_r1.into_raw(), expected_20);
1301    }
1302
1303    #[test]
1304    fn get_element_oob_row() {
1305        let mut trace = create_mock_trace(1);
1306        trace
1307            .add_column(TraceColumn::Bit(vec![Bit::ZERO; 2]))
1308            .unwrap();
1309        trace
1310            .get_element::<Block128>(0, 2)
1311            .expect_err("Expected out-of-bounds row error");
1312    }
1313
1314    #[test]
1315    fn get_element_oob_col() {
1316        let trace = create_mock_trace(1);
1317        trace
1318            .get_element::<Block128>(0, 0)
1319            .expect_err("Expected out-of-bounds column error");
1320    }
1321
1322    // ===========================================
1323    // SAFETY & SLICE TESTS
1324    // ===========================================
1325
1326    #[test]
1327    fn get_column_slice_correct_type() {
1328        let mut trace = create_mock_trace(2);
1329        let data = vec![
1330            Block32::from(1u32),
1331            Block32::from(2u32),
1332            Block32::from(3u32),
1333            Block32::from(4u32),
1334        ];
1335        trace.add_column(data.clone().into_trace_column()).unwrap();
1336
1337        let expected_hw: Vec<Flat<Block32>> = data.into_iter().map(|x| x.to_hardware()).collect();
1338
1339        let slice: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1340        assert_eq!(slice, expected_hw.as_slice());
1341    }
1342
1343    #[test]
1344    fn get_column_slice_wrong_type() {
1345        let mut trace = create_mock_trace(1);
1346        trace
1347            .add_column(vec![Block128::ZERO; 2].into_trace_column())
1348            .unwrap();
1349
1350        trace
1351            .get_column_slice::<Flat<Block32>>(0)
1352            .expect_err("Expected column type mismatch error");
1353    }
1354
1355    #[test]
1356    fn trace_stores_hardware_basis() {
1357        let mut trace = create_mock_trace(2);
1358
1359        let tower_data = vec![
1360            Block32::from(42u32),
1361            Block32::from(13u32),
1362            Block32::from(255u32),
1363            Block32::from(1u32),
1364        ];
1365
1366        let expected_hardware: Vec<Block32> = tower_data
1367            .iter()
1368            .map(|x| x.to_hardware().into_raw())
1369            .collect();
1370
1371        trace.add_column(tower_data.into_trace_column()).unwrap();
1372
1373        let stored: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1374
1375        for (i, (&stored_val, &expected_val)) in
1376            stored.iter().zip(expected_hardware.iter()).enumerate()
1377        {
1378            assert_eq!(
1379                stored_val.into_raw(),
1380                expected_val,
1381                "Row {}: stored value {:?} != expected hardware {:?}",
1382                i,
1383                stored_val,
1384                expected_val
1385            );
1386        }
1387    }
1388
1389    #[test]
1390    fn trace_hardware_basis_homomorphism() {
1391        let mut trace = create_mock_trace(3);
1392
1393        let a_tower = vec![Block32::from(5u32); 8];
1394        let b_tower = vec![Block32::from(7u32); 8];
1395
1396        trace
1397            .add_column(a_tower.clone().into_trace_column())
1398            .unwrap();
1399        trace
1400            .add_column(b_tower.clone().into_trace_column())
1401            .unwrap();
1402
1403        let a_stored: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1404        let b_stored: &[Flat<Block32>] = trace.get_column_slice(1).unwrap();
1405
1406        let a_hw_expected = a_tower[0].to_hardware().into_raw();
1407        let b_hw_expected = b_tower[0].to_hardware().into_raw();
1408
1409        assert_eq!(a_stored[0].into_raw(), a_hw_expected);
1410        assert_eq!(b_stored[0].into_raw(), b_hw_expected);
1411
1412        let product_hw = a_stored[0] * b_stored[0];
1413        let product_expected = (a_tower[0] * b_tower[0]).to_hardware().into_raw();
1414
1415        assert_eq!(product_hw.into_raw(), product_expected);
1416    }
1417
1418    // =========================================================
1419    // TRACE BUILDER TESTS
1420    // =========================================================
1421
1422    #[test]
1423    fn trace_builder_construction_and_auto_padding() {
1424        let layout = &[ColumnType::B32, ColumnType::Bit];
1425        let tb = TraceBuilder::new(layout, 2).unwrap(); // 4 rows
1426        assert_eq!(tb.num_rows(), 4);
1427
1428        let trace = tb.build();
1429        assert_eq!(trace.num_cols(), 2);
1430        assert_eq!(trace.num_rows().unwrap(), 4);
1431
1432        // All values should be zero (auto-padding)
1433        assert!(trace.columns[0].is_all_zeros());
1434        assert!(trace.columns[1].is_all_zeros());
1435    }
1436
1437    #[test]
1438    fn trace_builder_set_b32_stores_hardware_basis() {
1439        let layout = &[ColumnType::B32];
1440        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1441
1442        tb.set_b32(0, 0, Block32::from(42u32)).unwrap();
1443        tb.set_b32(0, 1, Block32::from(13u32)).unwrap();
1444
1445        let trace = tb.build();
1446        let stored: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1447
1448        assert_eq!(stored[0], Block32::from(42u32).to_hardware());
1449        assert_eq!(stored[1], Block32::from(13u32).to_hardware());
1450    }
1451
1452    #[test]
1453    fn trace_builder_column_ordering_matches_schema() {
1454        let layout = &[ColumnType::Bit, ColumnType::B32, ColumnType::B128];
1455        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1456
1457        tb.set_bit(0, 0, Bit::ONE).unwrap();
1458        tb.set_b32(1, 0, Block32::from(99u32)).unwrap();
1459        tb.set_b128(2, 0, Block128::from(7u8)).unwrap();
1460
1461        let trace = tb.build();
1462        assert_eq!(trace.columns[0].column_type(), ColumnType::Bit);
1463        assert_eq!(trace.columns[1].column_type(), ColumnType::B32);
1464        assert_eq!(trace.columns[2].column_type(), ColumnType::B128);
1465    }
1466
1467    #[test]
1468    fn trace_builder_type_mismatch_returns_error() {
1469        let layout = &[ColumnType::Bit];
1470        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1471
1472        let err = tb.set_b32(0, 0, Block32::ZERO);
1473        assert!(err.is_err());
1474    }
1475
1476    #[test]
1477    fn trace_builder_row_out_of_bounds_returns_error() {
1478        let layout = &[ColumnType::B32];
1479        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1480
1481        let err = tb.set_b32(0, 2, Block32::ZERO);
1482        assert!(err.is_err());
1483    }
1484
1485    #[test]
1486    fn trace_builder_col_out_of_bounds_returns_error() {
1487        let layout = &[ColumnType::B32];
1488        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1489
1490        let err = tb.set_b32(1, 0, Block32::ZERO);
1491        assert!(err.is_err());
1492    }
1493
1494    #[test]
1495    fn trace_builder_fill_selector() {
1496        let layout = &[ColumnType::Bit];
1497        let mut tb = TraceBuilder::new(layout, 2).unwrap(); // 4 rows
1498
1499        tb.fill_selector(0, 3).unwrap(); // rows 0,1,2 = ONE
1500
1501        let trace = tb.build();
1502        let bits = trace.columns[0].as_bit_slice().unwrap();
1503        assert_eq!(bits[0], Bit::ONE);
1504        assert_eq!(bits[1], Bit::ONE);
1505        assert_eq!(bits[2], Bit::ONE);
1506        assert_eq!(bits[3], Bit::ZERO); // padding
1507    }
1508
1509    #[test]
1510    fn trace_builder_push_mode() {
1511        let layout = &[ColumnType::B32, ColumnType::Bit];
1512        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1513
1514        tb.push_b32(0, Block32::from(10u32)).unwrap();
1515        tb.push_b32(0, Block32::from(20u32)).unwrap();
1516        tb.push_bit(1, Bit::ONE).unwrap();
1517        // row 1 of col 1 stays zero (auto-pad)
1518
1519        let trace = tb.build();
1520        let b32s: &[Flat<Block32>] = trace.get_column_slice(0).unwrap();
1521        assert_eq!(b32s[0], Block32::from(10u32).to_hardware());
1522        assert_eq!(b32s[1], Block32::from(20u32).to_hardware());
1523
1524        let bits = trace.columns[1].as_bit_slice().unwrap();
1525        assert_eq!(bits[0], Bit::ONE);
1526        assert_eq!(bits[1], Bit::ZERO);
1527    }
1528
1529    #[test]
1530    fn trace_builder_set_b32_array() {
1531        let layout = &[ColumnType::B32, ColumnType::B32, ColumnType::B32];
1532        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1533
1534        let vals = [
1535            Block32::from(1u32),
1536            Block32::from(2u32),
1537            Block32::from(3u32),
1538        ];
1539        tb.set_b32_array(0, 0, &vals).unwrap();
1540
1541        let trace = tb.build();
1542        for (i, &expected) in vals.iter().enumerate() {
1543            let stored: &[Flat<Block32>] = trace.get_column_slice(i).unwrap();
1544            assert_eq!(stored[0], expected.to_hardware());
1545        }
1546    }
1547
1548    #[test]
1549    fn trace_builder_set_b8() {
1550        let layout = &[ColumnType::B8];
1551        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1552
1553        tb.set_b8(0, 0, Block8(0xAB)).unwrap();
1554        tb.set_b8(0, 1, Block8(0xCD)).unwrap();
1555
1556        let trace = tb.build();
1557        let stored: &[Flat<Block8>] = trace.get_column_slice(0).unwrap();
1558        assert_eq!(stored[0], Block8(0xAB).to_hardware());
1559        assert_eq!(stored[1], Block8(0xCD).to_hardware());
1560    }
1561
1562    #[test]
1563    fn trace_builder_set_b16() {
1564        let layout = &[ColumnType::B16];
1565        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1566
1567        tb.set_b16(0, 0, Block16(1000)).unwrap();
1568        tb.set_b16(0, 1, Block16(2000)).unwrap();
1569
1570        let trace = tb.build();
1571        let stored: &[Flat<Block16>] = trace.get_column_slice(0).unwrap();
1572        assert_eq!(stored[0], Block16(1000).to_hardware());
1573        assert_eq!(stored[1], Block16(2000).to_hardware());
1574    }
1575
1576    #[test]
1577    fn trace_builder_set_b64() {
1578        let layout = &[ColumnType::B64];
1579        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1580
1581        tb.set_b64(0, 0, Block64(0xDEADBEEF_CAFEBABE)).unwrap();
1582
1583        let trace = tb.build();
1584        let stored: &[Flat<Block64>] = trace.get_column_slice(0).unwrap();
1585        assert_eq!(stored[0], Block64(0xDEADBEEF_CAFEBABE).to_hardware());
1586        assert_eq!(stored[1], Block64::ZERO.to_hardware()); // auto-padding
1587    }
1588
1589    #[test]
1590    fn trace_builder_set_b128() {
1591        let layout = &[ColumnType::B128];
1592        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1593
1594        let val = Block128::from(0xFFu8);
1595        tb.set_b128(0, 0, val).unwrap();
1596
1597        let trace = tb.build();
1598        let stored: &[Flat<Block128>] = trace.get_column_slice(0).unwrap();
1599        assert_eq!(stored[0], val.to_hardware());
1600    }
1601
1602    #[test]
1603    fn trace_builder_push_all_types() {
1604        let layout = &[
1605            ColumnType::Bit,
1606            ColumnType::B8,
1607            ColumnType::B16,
1608            ColumnType::B32,
1609            ColumnType::B64,
1610            ColumnType::B128,
1611        ];
1612        let mut tb = TraceBuilder::new(layout, 1).unwrap(); // 2 rows
1613
1614        tb.push_bit(0, Bit::ONE).unwrap();
1615        tb.push_b8(1, Block8(0x42)).unwrap();
1616        tb.push_b16(2, Block16(1234)).unwrap();
1617        tb.push_b32(3, Block32::from(5678u32)).unwrap();
1618        tb.push_b64(4, Block64(9999)).unwrap();
1619        tb.push_b128(5, Block128::from(77u8)).unwrap();
1620
1621        let trace = tb.build();
1622        assert_eq!(trace.num_cols(), 6);
1623
1624        let bits = trace.columns[0].as_bit_slice().unwrap();
1625        assert_eq!(bits[0], Bit::ONE);
1626        assert_eq!(bits[1], Bit::ZERO); // auto-pad
1627
1628        let b8s: &[Flat<Block8>] = trace.get_column_slice(1).unwrap();
1629        assert_eq!(b8s[0], Block8(0x42).to_hardware());
1630
1631        let b16s: &[Flat<Block16>] = trace.get_column_slice(2).unwrap();
1632        assert_eq!(b16s[0], Block16(1234).to_hardware());
1633
1634        let b32s: &[Flat<Block32>] = trace.get_column_slice(3).unwrap();
1635        assert_eq!(b32s[0], Block32::from(5678u32).to_hardware());
1636
1637        let b64s: &[Flat<Block64>] = trace.get_column_slice(4).unwrap();
1638        assert_eq!(b64s[0], Block64(9999).to_hardware());
1639
1640        let b128s: &[Flat<Block128>] = trace.get_column_slice(5).unwrap();
1641        assert_eq!(b128s[0], Block128::from(77u8).to_hardware());
1642    }
1643
1644    #[test]
1645    fn trace_builder_type_mismatch_all_setters() {
1646        // B32 column, every non-B32 setter should fail
1647        let layout = &[ColumnType::B32];
1648        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1649
1650        assert!(tb.set_bit(0, 0, Bit::ONE).is_err());
1651        assert!(tb.set_b8(0, 0, Block8(1)).is_err());
1652        assert!(tb.set_b16(0, 0, Block16(1)).is_err());
1653        assert!(tb.set_b64(0, 0, Block64(1)).is_err());
1654        assert!(tb.set_b128(0, 0, Block128::ONE).is_err());
1655
1656        // Correct type succeeds
1657        assert!(tb.set_b32(0, 0, Block32::ONE).is_ok());
1658    }
1659
1660    #[test]
1661    fn trace_builder_array_setters_all_types() {
1662        let layout = &[
1663            ColumnType::Bit,
1664            ColumnType::Bit,
1665            ColumnType::B8,
1666            ColumnType::B8,
1667            ColumnType::B64,
1668            ColumnType::B64,
1669            ColumnType::B128,
1670            ColumnType::B128,
1671        ];
1672        let mut tb = TraceBuilder::new(layout, 1).unwrap();
1673
1674        tb.set_bit_array(0, 0, &[Bit::ONE, Bit::ZERO]).unwrap();
1675        tb.set_b8_array(2, 0, &[Block8(10), Block8(20)]).unwrap();
1676        tb.set_b64_array(4, 0, &[Block64(100), Block64(200)])
1677            .unwrap();
1678        tb.set_b128_array(6, 0, &[Block128::ONE, Block128::from(2u8)])
1679            .unwrap();
1680
1681        let trace = tb.build();
1682
1683        let bits = trace.columns[0].as_bit_slice().unwrap();
1684        assert_eq!(bits[0], Bit::ONE);
1685
1686        let bits1 = trace.columns[1].as_bit_slice().unwrap();
1687        assert_eq!(bits1[0], Bit::ZERO);
1688
1689        let b8s: &[Flat<Block8>] = trace.get_column_slice(2).unwrap();
1690        assert_eq!(b8s[0], Block8(10).to_hardware());
1691
1692        let b8s1: &[Flat<Block8>] = trace.get_column_slice(3).unwrap();
1693        assert_eq!(b8s1[0], Block8(20).to_hardware());
1694    }
1695
1696    #[test]
1697    fn trace_builder_invalid_num_vars() {
1698        let layout = &[ColumnType::B32];
1699        // num_vars too large to shift
1700        assert!(TraceBuilder::new(layout, 128).is_err());
1701    }
1702}