Skip to main content

miden_processor/trace/
utils.rs

1use alloc::{string::ToString, vec::Vec};
2use core::{mem::MaybeUninit, slice};
3
4use miden_air::trace::{Challenges, MIN_TRACE_LEN, MainTrace};
5
6use super::chiplets::Chiplets;
7use crate::{
8    Felt, RowIndex,
9    debug::BusDebugger,
10    field::ExtensionField,
11    utils::{assume_init_vec, uninit_vector},
12};
13#[cfg(test)]
14use crate::{operation::Operation, utils::ToElements};
15
16// ROW-MAJOR TRACE WRITER
17// ================================================================================================
18
19/// Row-major flat buffer writer (`write_row` is a single `copy_from_slice`).
20#[derive(Debug)]
21pub struct RowMajorTraceWriter<'a, E> {
22    data: &'a mut [E],
23    width: usize,
24}
25
26impl<'a, E: Copy> RowMajorTraceWriter<'a, E> {
27    pub fn new(data: &'a mut [E], width: usize) -> Self {
28        debug_assert_eq!(data.len() % width, 0, "buffer length must be a multiple of width");
29        Self { data, width }
30    }
31
32    /// Writes one row; `values.len()` must equal `width`.
33    #[inline(always)]
34    pub fn write_row(&mut self, row: usize, values: &[E]) {
35        debug_assert_eq!(values.len(), self.width);
36        let start = row * self.width;
37        self.data[start..start + self.width].copy_from_slice(values);
38    }
39}
40
41// TRACE FRAGMENT
42// ================================================================================================
43
44/// TODO: add docs
45pub struct TraceFragment<'a> {
46    data: Vec<&'a mut [Felt]>,
47    num_rows: usize,
48}
49
50impl<'a> TraceFragment<'a> {
51    /// Creates a new [TraceFragment] with the expected number of columns and rows.
52    ///
53    /// The memory needed to hold the trace fragment data is not allocated during construction.
54    pub fn new(num_columns: usize, num_rows: usize) -> Self {
55        TraceFragment {
56            data: Vec::with_capacity(num_columns),
57            num_rows,
58        }
59    }
60
61    // PUBLIC ACCESSORS
62    // --------------------------------------------------------------------------------------------
63
64    /// Returns the number of columns in this execution trace fragment.
65    pub fn width(&self) -> usize {
66        self.data.len()
67    }
68
69    /// Returns the number of rows in this execution trace fragment.
70    pub fn len(&self) -> usize {
71        self.num_rows
72    }
73
74    // DATA MUTATORS
75    // --------------------------------------------------------------------------------------------
76
77    /// Updates a single cell in this fragment with provided value.
78    #[inline(always)]
79    pub fn set(&mut self, row_idx: RowIndex, col_idx: usize, value: Felt) {
80        self.data[col_idx][row_idx] = value;
81    }
82
83    /// Returns a mutable iterator to the columns of this fragment.
84    pub fn columns(&mut self) -> slice::IterMut<'_, &'a mut [Felt]> {
85        self.data.iter_mut()
86    }
87
88    /// Adds a new column to this fragment by pushing a mutable slice with the first `self.len()`
89    /// elements of the provided column.
90    ///
91    /// Returns the rest of the provided column as a separate mutable slice.
92    pub fn push_column_slice(&mut self, column: &'a mut [Felt]) -> &'a mut [Felt] {
93        let (column_fragment, rest) = column.split_at_mut(self.num_rows);
94        self.data.push(column_fragment);
95        rest
96    }
97
98    // TEST METHODS
99    // --------------------------------------------------------------------------------------------
100
101    #[cfg(test)]
102    pub fn trace_to_fragment(trace: &'a mut [Vec<Felt>]) -> Self {
103        assert!(!trace.is_empty(), "expected trace to have at least one column");
104        let mut data = Vec::new();
105        for column in trace.iter_mut() {
106            data.push(column.as_mut_slice());
107        }
108
109        let num_rows = data[0].len();
110        Self { data, num_rows }
111    }
112}
113
114// TRACE LENGTH SUMMARY
115// ================================================================================================
116
117/// Contains the data about lengths of the trace parts.
118///
119/// - `main_trace_len` contains the length of the main trace.
120/// - `range_trace_len` contains the length of the range checker trace.
121/// - `chiplets_trace_len` contains the trace lengths of the all chiplets (hash, bitwise, memory,
122///   kernel ROM)
123#[derive(Debug, Default, Eq, PartialEq, Clone, Copy)]
124pub struct TraceLenSummary {
125    main_trace_len: usize,
126    range_trace_len: usize,
127    chiplets_trace_len: ChipletsLengths,
128}
129
130impl TraceLenSummary {
131    pub fn new(
132        main_trace_len: usize,
133        range_trace_len: usize,
134        chiplets_trace_len: ChipletsLengths,
135    ) -> Self {
136        TraceLenSummary {
137            main_trace_len,
138            range_trace_len,
139            chiplets_trace_len,
140        }
141    }
142
143    /// Returns length of the main trace.
144    pub fn main_trace_len(&self) -> usize {
145        self.main_trace_len
146    }
147
148    /// Returns length of the range checker trace.
149    pub fn range_trace_len(&self) -> usize {
150        self.range_trace_len
151    }
152
153    /// Returns [ChipletsLengths] which contains trace lengths of all chilplets.
154    pub fn chiplets_trace_len(&self) -> ChipletsLengths {
155        self.chiplets_trace_len
156    }
157
158    /// Returns the maximum of all component lengths.
159    pub fn trace_len(&self) -> usize {
160        self.range_trace_len
161            .max(self.main_trace_len)
162            .max(self.chiplets_trace_len.trace_len())
163    }
164
165    /// Returns `trace_len` rounded up to the next power of two, clamped to `MIN_TRACE_LEN`.
166    pub fn padded_trace_len(&self) -> usize {
167        self.trace_len().next_power_of_two().max(MIN_TRACE_LEN)
168    }
169
170    /// Returns the percent (0 - 100) of the steps that were added to the trace to pad it to the
171    /// next power of tow.
172    pub fn padding_percentage(&self) -> usize {
173        (self.padded_trace_len() - self.trace_len()) * 100 / self.padded_trace_len()
174    }
175}
176
177// CHIPLET LENGTHS
178// ================================================================================================
179
180/// Contains trace lengths of all chiplets: hash, bitwise, memory, ACE, and kernel ROM.
181#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)]
182pub struct ChipletsLengths {
183    hash_chiplet_len: usize,
184    bitwise_chiplet_len: usize,
185    memory_chiplet_len: usize,
186    ace_chiplet_len: usize,
187    kernel_rom_len: usize,
188}
189
190impl ChipletsLengths {
191    pub fn new(chiplets: &Chiplets) -> Self {
192        ChipletsLengths {
193            hash_chiplet_len: chiplets.bitwise_start().into(),
194            bitwise_chiplet_len: chiplets.memory_start() - chiplets.bitwise_start(),
195            memory_chiplet_len: chiplets.ace_start() - chiplets.memory_start(),
196            ace_chiplet_len: chiplets.kernel_rom_start() - chiplets.ace_start(),
197            kernel_rom_len: chiplets.padding_start() - chiplets.kernel_rom_start(),
198        }
199    }
200
201    pub fn from_parts(
202        hash_len: usize,
203        bitwise_len: usize,
204        memory_len: usize,
205        ace_len: usize,
206        kernel_len: usize,
207    ) -> Self {
208        ChipletsLengths {
209            hash_chiplet_len: hash_len,
210            bitwise_chiplet_len: bitwise_len,
211            memory_chiplet_len: memory_len,
212            ace_chiplet_len: ace_len,
213            kernel_rom_len: kernel_len,
214        }
215    }
216
217    /// Returns the length of the hash chiplet trace
218    pub fn hash_chiplet_len(&self) -> usize {
219        self.hash_chiplet_len
220    }
221
222    /// Returns the length of the bitwise trace
223    pub fn bitwise_chiplet_len(&self) -> usize {
224        self.bitwise_chiplet_len
225    }
226
227    /// Returns the length of the memory trace
228    pub fn memory_chiplet_len(&self) -> usize {
229        self.memory_chiplet_len
230    }
231
232    /// Returns the length of the ACE chiplet trace
233    pub fn ace_chiplet_len(&self) -> usize {
234        self.ace_chiplet_len
235    }
236
237    /// Returns the length of the kernel ROM trace
238    pub fn kernel_rom_len(&self) -> usize {
239        self.kernel_rom_len
240    }
241
242    /// Returns the length of the trace required to accommodate chiplet components and 1
243    /// mandatory padding row required for ensuring sufficient trace length for auxiliary connector
244    /// columns that rely on the memory chiplet.
245    pub fn trace_len(&self) -> usize {
246        self.hash_chiplet_len()
247            + self.bitwise_chiplet_len()
248            + self.memory_chiplet_len()
249            + self.ace_chiplet_len()
250            + self.kernel_rom_len()
251            + 1
252    }
253}
254
255// AUXILIARY COLUMN BUILDER
256// ================================================================================================
257
258/// Defines a builder responsible for building a single auxiliary bus column in the execution
259/// trace.
260///
261/// Columns are initialized to the multiset identity. Public-input-dependent boundary
262/// terms are used in the check by the verifier in `reduced_aux_values` for the final values of
263/// the auxiliary columns.
264pub(crate) trait AuxColumnBuilder<E: ExtensionField<Felt>> {
265    // REQUIRED METHODS
266    // --------------------------------------------------------------------------------------------
267
268    fn get_requests_at(
269        &self,
270        main_trace: &MainTrace,
271        challenges: &Challenges<E>,
272        row: RowIndex,
273        debugger: &mut BusDebugger<E>,
274    ) -> E;
275
276    fn get_responses_at(
277        &self,
278        main_trace: &MainTrace,
279        challenges: &Challenges<E>,
280        row: RowIndex,
281        debugger: &mut BusDebugger<E>,
282    ) -> E;
283
284    /// Whether to assert that all requests/responses balance in debug mode.
285    ///
286    /// Buses whose final value encodes a public-input-dependent boundary term (checked
287    /// via `reduced_aux_values`) will NOT balance to identity and should return `false`.
288    #[cfg(any(test, feature = "bus-debugger"))]
289    fn enforce_bus_balance(&self) -> bool;
290
291    // PROVIDED METHODS
292    // --------------------------------------------------------------------------------------------
293
294    /// Builds an auxiliary bus trace column as a running product of responses over requests.
295    ///
296    /// The column is initialized to 1; boundary terms are checked via `reduced_aux_values`
297    /// in the verifier.
298    fn build_aux_column(&self, main_trace: &MainTrace, challenges: &Challenges<E>) -> Vec<E> {
299        let mut bus_debugger = BusDebugger::new("aux bus".to_string());
300
301        let mut requests: Vec<MaybeUninit<E>> = uninit_vector(main_trace.num_rows());
302        requests[0].write(E::ONE);
303
304        let mut responses_prod: Vec<MaybeUninit<E>> = uninit_vector(main_trace.num_rows());
305        responses_prod[0].write(E::ONE);
306
307        let mut requests_running_prod = E::ONE;
308        let mut prev_prod = E::ONE;
309
310        // Product of all requests to be inverted, used to compute inverses of requests.
311        for row_idx in 0..main_trace.num_rows() - 1 {
312            let row = row_idx.into();
313
314            let response = self.get_responses_at(main_trace, challenges, row, &mut bus_debugger);
315            prev_prod *= response;
316            responses_prod[row_idx + 1].write(prev_prod);
317
318            let request = self.get_requests_at(main_trace, challenges, row, &mut bus_debugger);
319            requests[row_idx + 1].write(request);
320            requests_running_prod *= request;
321        }
322
323        // all elements are now initialized
324        let requests = unsafe { assume_init_vec(requests) };
325        let mut result_aux_column = unsafe { assume_init_vec(responses_prod) };
326
327        // Use batch-inversion method to compute running product of `response[i]/request[i]`.
328        let mut requests_running_divisor = requests_running_prod.inverse();
329        for i in (0..main_trace.num_rows()).rev() {
330            result_aux_column[i] *= requests_running_divisor;
331            requests_running_divisor *= requests[i];
332        }
333
334        #[cfg(any(test, feature = "bus-debugger"))]
335        if self.enforce_bus_balance() {
336            assert!(bus_debugger.is_empty(), "{bus_debugger}");
337        }
338
339        result_aux_column
340    }
341}
342
343// U32 HELPERS
344// ================================================================================================
345
346/// Splits an element into two 16 bit integer limbs. It assumes that the field element contains a
347/// valid 32-bit integer value.
348pub(crate) fn split_element_u32_into_u16(value: Felt) -> (Felt, Felt) {
349    let (hi, lo) = split_u32_into_u16(value.as_canonical_u64());
350    (Felt::new(hi as u64), Felt::new(lo as u64))
351}
352
353/// Splits a u64 integer assumed to contain a 32-bit value into two u16 integers.
354///
355/// # Errors
356/// Fails in debug mode if the provided value is not a 32-bit value.
357pub(crate) fn split_u32_into_u16(value: u64) -> (u16, u16) {
358    const U32MAX: u64 = u32::MAX as u64;
359    debug_assert!(value <= U32MAX, "not a 32-bit value");
360
361    let lo = value as u16;
362    let hi = (value >> 16) as u16;
363
364    (hi, lo)
365}
366
367// TEST HELPERS
368// ================================================================================================
369
370#[cfg(test)]
371pub fn build_span_with_respan_ops() -> (Vec<Operation>, Vec<Felt>) {
372    let iv = [1, 3, 5, 7, 9, 11, 13, 15, 17].to_elements();
373    let ops = vec![
374        Operation::Push(iv[0]),
375        Operation::Push(iv[1]),
376        Operation::Push(iv[2]),
377        Operation::Push(iv[3]),
378        Operation::Push(iv[4]),
379        Operation::Push(iv[5]),
380        Operation::Push(iv[6]),
381        // next batch
382        Operation::Push(iv[7]),
383        Operation::Push(iv[8]),
384        Operation::Add,
385        // drops to make sure stack overflow is empty on exit
386        Operation::Drop,
387        Operation::Drop,
388        Operation::Drop,
389        Operation::Drop,
390        Operation::Drop,
391        Operation::Drop,
392        Operation::Drop,
393        Operation::Drop,
394    ];
395    (ops, iv)
396}