Skip to main content

miden_lifted_stark/lmcs/
utils.rs

1//! Utility functions for LMCS operations.
2
3use alloc::vec::Vec;
4use core::array;
5
6use p3_field::PackedValue;
7use p3_util::log2_strict_usize;
8
9/// Strict logâ‚‚ returning `u8`.
10///
11/// Panics if `n` is not a power of two.
12#[inline]
13pub fn log2_strict_u8(n: usize) -> u8 {
14    log2_strict_usize(n) as u8
15}
16
17/// Extension trait for `PackedValue` providing columnar pack/unpack operations.
18///
19/// These methods perform transpose operations on packed data, useful for
20/// SIMD-parallelized Merkle tree construction.
21pub trait PackedValueExt: PackedValue {
22    /// Pack columns from `WIDTH` rows of scalar values.
23    ///
24    /// Given `WIDTH` rows of `N` scalar values, extract each column and pack it
25    /// into a single packed value. This performs a transpose operation.
26    #[inline]
27    #[must_use]
28    fn pack_columns<const N: usize>(rows: &[[Self::Value; N]]) -> [Self; N] {
29        assert_eq!(rows.len(), Self::WIDTH);
30        array::from_fn(|col| Self::from_fn(|lane| rows[lane][col]))
31    }
32}
33
34// Blanket implementation for all PackedValue types
35impl<T: PackedValue> PackedValueExt for T {}
36
37/// Compute the aligned length for `len` given an alignment.
38#[inline]
39pub const fn aligned_len(len: usize, alignment: usize) -> usize {
40    if alignment <= 1 {
41        len
42    } else {
43        len.next_multiple_of(alignment)
44    }
45}
46
47/// Align each width in place, returning the same `Vec`.
48pub fn aligned_widths(mut widths: Vec<usize>, alignment: usize) -> Vec<usize> {
49    for w in &mut widths {
50        *w = aligned_len(*w, alignment);
51    }
52    widths
53}