p3_field/packed/packed_traits.rs
1use core::iter::{Product, Sum};
2use core::mem::MaybeUninit;
3use core::ops::{Div, DivAssign};
4use core::{array, slice};
5
6use crate::field::Field;
7use crate::{Algebra, BasedVectorSpace, ExtensionField, Powers, PrimeCharacteristicRing};
8
9/// A trait to constrain types that can be packed into a packed value.
10///
11/// The `Packable` trait allows us to specify implementations for potentially conflicting types.
12pub trait Packable: 'static + Default + Copy + Send + Sync + PartialEq + Eq {}
13
14/// A trait for array-like structs made up of multiple scalar elements.
15///
16/// # Safety
17/// - If `P` implements `PackedField` then `P` must be castable to/from `[P::Value; P::WIDTH]`
18/// without UB.
19pub unsafe trait PackedValue: 'static + Copy + Send + Sync {
20 /// The scalar type that is packed into this value.
21 type Value: Packable;
22
23 /// Number of scalar values packed together.
24 const WIDTH: usize;
25
26 /// Constructs a packed value using a function to generate each element.
27 ///
28 /// Similar to [`core::array::from_fn`].
29 #[must_use]
30 fn from_fn<F>(f: F) -> Self
31 where
32 F: FnMut(usize) -> Self::Value;
33
34 /// Create a packed value with all lanes set to the same scalar value.
35 #[inline]
36 #[must_use]
37 fn broadcast(value: Self::Value) -> Self {
38 Self::from_fn(|_| value)
39 }
40
41 /// Interprets a slice of scalar values as a packed value reference.
42 ///
43 /// # Panics:
44 /// This function will panic if `slice.len() != Self::WIDTH`
45 #[must_use]
46 fn from_slice(slice: &[Self::Value]) -> &Self;
47
48 /// Interprets a mutable slice of scalar values as a mutable packed value.
49 ///
50 /// # Panics:
51 /// This function will panic if `slice.len() != Self::WIDTH`
52 #[must_use]
53 fn from_slice_mut(slice: &mut [Self::Value]) -> &mut Self;
54
55 /// Returns the underlying scalar values as an immutable slice.
56 #[must_use]
57 fn as_slice(&self) -> &[Self::Value];
58
59 /// Returns the underlying scalar values as a mutable slice.
60 #[must_use]
61 fn as_slice_mut(&mut self) -> &mut [Self::Value];
62
63 /// Extract the scalar value at the given SIMD lane.
64 ///
65 /// This is equivalent to `self.as_slice()[lane]` but more explicit about the
66 /// SIMD extraction semantics.
67 #[inline]
68 #[must_use]
69 fn extract(&self, lane: usize) -> Self::Value {
70 self.as_slice()[lane]
71 }
72
73 /// Packs a slice of scalar values into a slice of packed values.
74 ///
75 /// # Panics
76 /// Panics if the slice length is not divisible by `WIDTH`.
77 #[inline]
78 #[must_use]
79 fn pack_slice(buf: &[Self::Value]) -> &[Self] {
80 // Sources vary, but this should be true on all platforms we care about.
81 const {
82 assert!(align_of::<Self>() <= align_of::<Self::Value>());
83 }
84 assert!(
85 buf.len().is_multiple_of(Self::WIDTH),
86 "Slice length (got {}) must be a multiple of packed field width ({}).",
87 buf.len(),
88 Self::WIDTH
89 );
90 let buf_ptr = buf.as_ptr().cast::<Self>();
91 let n = buf.len() / Self::WIDTH;
92 // SAFETY: `buf_ptr` is valid for `n * WIDTH` values of `Self::Value`, which is
93 // the same region as `n` values of `Self` given the alignment and length checks above.
94 unsafe { slice::from_raw_parts(buf_ptr, n) }
95 }
96
97 /// Converts a mutable slice of scalar values into a mutable slice of packed values.
98 ///
99 /// # Panics
100 /// Panics if the slice length is not divisible by `WIDTH`.
101 #[inline]
102 #[must_use]
103 fn pack_slice_mut(buf: &mut [Self::Value]) -> &mut [Self] {
104 const {
105 assert!(align_of::<Self>() <= align_of::<Self::Value>());
106 }
107 assert!(
108 buf.len().is_multiple_of(Self::WIDTH),
109 "Slice length (got {}) must be a multiple of packed field width ({}).",
110 buf.len(),
111 Self::WIDTH
112 );
113 let buf_ptr = buf.as_mut_ptr().cast::<Self>();
114 let n = buf.len() / Self::WIDTH;
115 // SAFETY: `buf_ptr` is valid for `n * WIDTH` values of `Self::Value`, which is
116 // the same region as `n` values of `Self` given the alignment and length checks above.
117 unsafe { slice::from_raw_parts_mut(buf_ptr, n) }
118 }
119
120 /// Converts a mutable slice of possibly uninitialized scalar values into
121 /// a mutable slice of possibly uninitialized packed values.
122 ///
123 /// # Panics
124 /// Panics if the slice length is not divisible by `WIDTH`.
125 #[inline]
126 #[must_use]
127 fn pack_maybe_uninit_slice_mut(
128 buf: &mut [MaybeUninit<Self::Value>],
129 ) -> &mut [MaybeUninit<Self>] {
130 const {
131 assert!(align_of::<Self>() <= align_of::<Self::Value>());
132 }
133 assert!(
134 buf.len().is_multiple_of(Self::WIDTH),
135 "Slice length (got {}) must be a multiple of packed field width ({}).",
136 buf.len(),
137 Self::WIDTH
138 );
139 let buf_ptr = buf.as_mut_ptr().cast::<MaybeUninit<Self>>();
140 let n = buf.len() / Self::WIDTH;
141 // SAFETY: `buf_ptr` is valid for `n * WIDTH` values of `MaybeUninit<Self::Value>`,
142 // which is the same region as `n` values of `MaybeUninit<Self>` given the
143 // alignment and length checks above.
144 unsafe { slice::from_raw_parts_mut(buf_ptr, n) }
145 }
146
147 /// Packs a slice into packed values and returns the packed portion and any remaining suffix.
148 #[inline]
149 #[must_use]
150 fn pack_slice_with_suffix(buf: &[Self::Value]) -> (&[Self], &[Self::Value]) {
151 let (packed, suffix) = buf.split_at(buf.len() - buf.len() % Self::WIDTH);
152 (Self::pack_slice(packed), suffix)
153 }
154
155 /// Converts a mutable slice of scalar values into a pair:
156 /// - a slice of packed values covering the largest aligned portion,
157 /// - and a remainder slice of scalar values that couldn't be packed.
158 #[inline]
159 #[must_use]
160 fn pack_slice_with_suffix_mut(buf: &mut [Self::Value]) -> (&mut [Self], &mut [Self::Value]) {
161 let (packed, suffix) = buf.split_at_mut(buf.len() - buf.len() % Self::WIDTH);
162 (Self::pack_slice_mut(packed), suffix)
163 }
164
165 /// Converts a mutable slice of possibly uninitialized scalar values into a pair:
166 /// - a slice of possibly uninitialized packed values covering the largest aligned portion,
167 /// - and a remainder slice of possibly uninitialized scalar values that couldn't be packed.
168 #[inline]
169 #[must_use]
170 fn pack_maybe_uninit_slice_with_suffix_mut(
171 buf: &mut [MaybeUninit<Self::Value>],
172 ) -> (&mut [MaybeUninit<Self>], &mut [MaybeUninit<Self::Value>]) {
173 let (packed, suffix) = buf.split_at_mut(buf.len() - buf.len() % Self::WIDTH);
174 (Self::pack_maybe_uninit_slice_mut(packed), suffix)
175 }
176
177 /// Reinterprets a slice of packed values as a flat slice of scalar values.
178 ///
179 /// Each packed value contains `Self::WIDTH` scalar values, which are laid out
180 /// contiguously in memory. This function allows direct access to those scalars.
181 #[inline]
182 #[must_use]
183 fn unpack_slice(buf: &[Self]) -> &[Self::Value] {
184 const {
185 assert!(align_of::<Self>() >= align_of::<Self::Value>());
186 }
187 let buf_ptr = buf.as_ptr().cast::<Self::Value>();
188 let n = buf.len() * Self::WIDTH;
189 unsafe { slice::from_raw_parts(buf_ptr, n) }
190 }
191
192 /// Pack columns from `WIDTH` rows of scalar values into `N` packed values.
193 ///
194 /// Given `WIDTH` rows of `N` scalar values, extract each column and pack it
195 /// into a single packed value. This is the inverse of `unpack_into`.
196 ///
197 /// ## Panics
198 /// Panics if `rows.len() != WIDTH`.
199 #[inline]
200 #[must_use]
201 fn pack_columns<const N: usize>(rows: &[[Self::Value; N]]) -> [Self; N] {
202 assert_eq!(rows.len(), Self::WIDTH);
203 array::from_fn(|col| Self::from_fn(|lane| rows[lane][col]))
204 }
205
206 /// Pack columns using a closure that provides each row's data.
207 ///
208 /// Calls `row_fn(lane)` for each lane `0..WIDTH` to get `[Self::Value; N]`,
209 /// then transposes columns into packed values. Useful when rows aren't
210 /// contiguous in memory (e.g., strided access).
211 #[inline]
212 #[must_use]
213 fn pack_columns_fn<const N: usize>(row_fn: impl Fn(usize) -> [Self::Value; N]) -> [Self; N] {
214 array::from_fn(|col| Self::from_fn(|lane| row_fn(lane)[col]))
215 }
216
217 /// Unpack `N` packed values into `WIDTH` rows of `N` scalars.
218 ///
219 /// ## Inputs
220 /// - `packed`: An array of `N` packed values.
221 /// - `rows`: A mutable slice of exactly `WIDTH` arrays to write the unpacked values.
222 ///
223 /// ## Panics
224 /// Panics if `rows.len() != WIDTH`.
225 #[inline]
226 fn unpack_into<const N: usize>(packed: &[Self; N], rows: &mut [[Self::Value; N]]) {
227 assert_eq!(rows.len(), Self::WIDTH);
228 for (lane, row) in rows.iter_mut().enumerate() {
229 *row = array::from_fn(|col| packed[col].extract(lane));
230 }
231 }
232
233 /// Unpack `N` packed values into an iterator of `WIDTH` rows.
234 ///
235 /// This is the iterator equivalent of `unpack_into`, yielding each row
236 /// without requiring a pre-allocated buffer.
237 #[inline]
238 fn unpack_iter<const N: usize>(packed: [Self; N]) -> impl Iterator<Item = [Self::Value; N]> {
239 (0..Self::WIDTH).map(move |lane| array::from_fn(|col| packed[col].extract(lane)))
240 }
241}
242
243unsafe impl<T: Packable, const WIDTH: usize> PackedValue for [T; WIDTH] {
244 type Value = T;
245 const WIDTH: usize = WIDTH;
246
247 #[inline]
248 fn from_slice(slice: &[Self::Value]) -> &Self {
249 assert_eq!(slice.len(), Self::WIDTH);
250 unsafe { &*slice.as_ptr().cast() }
251 }
252
253 #[inline]
254 fn from_slice_mut(slice: &mut [Self::Value]) -> &mut Self {
255 assert_eq!(slice.len(), Self::WIDTH);
256 unsafe { &mut *slice.as_mut_ptr().cast() }
257 }
258
259 #[inline]
260 fn from_fn<Fn>(f: Fn) -> Self
261 where
262 Fn: FnMut(usize) -> Self::Value,
263 {
264 core::array::from_fn(f)
265 }
266
267 #[inline]
268 fn as_slice(&self) -> &[Self::Value] {
269 self
270 }
271
272 #[inline]
273 fn as_slice_mut(&mut self) -> &mut [Self::Value] {
274 self
275 }
276}
277
278/// An array of field elements which can be packed into a vector for SIMD operations.
279///
280/// # Safety
281/// - See `PackedValue` above.
282pub unsafe trait PackedField:
283 Algebra<Self::Scalar>
284 + PackedValue<Value = Self::Scalar>
285 + Div<Self, Output = Self>
286 + Div<Self::Scalar, Output = Self>
287 + DivAssign<Self>
288 + DivAssign<Self::Scalar>
289 + Sum<Self::Scalar>
290 + Product<Self::Scalar>
291{
292 type Scalar: Field;
293
294 /// Construct an iterator which returns powers of `base` packed into packed field elements.
295 ///
296 /// E.g. if `Self::WIDTH = 4`, returns: `[base^0, base^1, base^2, base^3], [base^4, base^5, base^6, base^7], ...`.
297 #[must_use]
298 fn packed_powers(base: Self::Scalar) -> Powers<Self> {
299 Self::packed_shifted_powers(base, Self::Scalar::ONE)
300 }
301
302 /// Construct an iterator which returns powers of `base` multiplied by `start` and packed into packed field elements.
303 ///
304 /// E.g. if `Self::WIDTH = 4`, returns: `[start, start*base, start*base^2, start*base^3], [start*base^4, start*base^5, start*base^6, start*base^7], ...`.
305 #[must_use]
306 fn packed_shifted_powers(base: Self::Scalar, start: Self::Scalar) -> Powers<Self> {
307 let mut current: Self = start.into();
308 let slice = current.as_slice_mut();
309 for i in 1..Self::WIDTH {
310 slice[i] = slice[i - 1] * base;
311 }
312
313 Powers {
314 base: base.exp_u64(Self::WIDTH as u64).into(),
315 current,
316 }
317 }
318
319 /// Accumulate the products of `d` coefficient streams against a shared stream of
320 /// packed base values.
321 ///
322 /// For each `(coeffs, base)` pair produced by the iterator, performs
323 /// `acc[k] += coeffs[k] * base` for all `k < d`, and returns the accumulators
324 /// (entries at `d..` are zero). Coefficient slices shorter than `d` only
325 /// contribute their available entries.
326 ///
327 /// This is the inner kernel of mixed base-times-extension dot products, where each
328 /// extension element contributes `d` base-field coefficient words. Implementations
329 /// may override it to defer modular reductions across iterations.
330 ///
331 /// # Panics
332 /// Debug builds panic if `d > 8`.
333 #[must_use]
334 fn coeffwise_dot_product<'a, I>(d: usize, pairs: I) -> [Self; 8]
335 where
336 Self: 'a,
337 I: Iterator<Item = (&'a [Self], Self)>,
338 {
339 debug_assert!(d <= 8, "Extension degree > 8 not supported");
340 let mut acc = [Self::ZERO; 8];
341 for (coeffs, base) in pairs {
342 for (acc_k, &coeff) in acc[..d].iter_mut().zip(coeffs) {
343 *acc_k += coeff * base;
344 }
345 }
346 acc
347 }
348}
349
350/// # Safety
351/// - `WIDTH` is assumed to be a power of 2.
352pub unsafe trait PackedFieldPow2: PackedField {
353 /// Take interpret two vectors as chunks of `block_len` elements. Unpack and interleave those
354 /// chunks. This is best seen with an example. If we have:
355 /// ```text
356 /// A = [x0, y0, x1, y1]
357 /// B = [x2, y2, x3, y3]
358 /// ```
359 ///
360 /// then
361 ///
362 /// ```text
363 /// interleave(A, B, 1) = ([x0, x2, x1, x3], [y0, y2, y1, y3])
364 /// ```
365 ///
366 /// Pairs that were adjacent in the input are at corresponding positions in the output.
367 ///
368 /// `r` lets us set the size of chunks we're interleaving. If we set `block_len = 2`, then for
369 ///
370 /// ```text
371 /// A = [x0, x1, y0, y1]
372 /// B = [x2, x3, y2, y3]
373 /// ```
374 ///
375 /// we obtain
376 ///
377 /// ```text
378 /// interleave(A, B, block_len) = ([x0, x1, x2, x3], [y0, y1, y2, y3])
379 /// ```
380 ///
381 /// We can also think about this as stacking the vectors, dividing them into 2x2 matrices, and
382 /// transposing those matrices.
383 ///
384 /// When `block_len = WIDTH`, this operation is a no-op.
385 ///
386 /// # Panics
387 /// This may panic if `block_len` does not divide `WIDTH`. Since `WIDTH` is specified to be a power of 2,
388 /// `block_len` must also be a power of 2. It cannot be 0 and it cannot exceed `WIDTH`.
389 #[must_use]
390 fn interleave(&self, other: Self, block_len: usize) -> (Self, Self);
391}
392
393/// Fix a field `F` a packing width `W` and an extension field `EF` of `F`.
394///
395/// By choosing a basis `B`, `EF` can be transformed into an array `[F; D]`.
396///
397/// A type should implement PackedFieldExtension if it can be transformed into `[F::Packing; D] ~ [[F; W]; D]`
398///
399/// This is interpreted by taking a transpose to get `[[F; D]; W]` which can then be reinterpreted
400/// as `[EF; W]` by making use of the chosen basis `B` again.
401pub trait PackedFieldExtension<
402 BaseField: Field,
403 ExtField: ExtensionField<BaseField, ExtensionPacking = Self>,
404>: Algebra<ExtField> + Algebra<BaseField::Packing> + BasedVectorSpace<BaseField::Packing>
405{
406 /// Construct a packed extension by applying `f` to each lane.
407 ///
408 /// This is the extension-field analog of [`PackedValue::from_fn`] and the canonical
409 /// primitive constructor for packed extensions: every other constructor in this
410 /// trait (`from_ext_slice`, `pack_ext_columns`, etc.) routes through it.
411 ///
412 /// `f` is called once per `(basis_coefficient, lane)` pair (`D * W` calls total),
413 /// hence the [`Fn`] bound — closures with side effects are unsuitable.
414 ///
415 /// The default impl uses only the [`BasedVectorSpace`] machinery the trait already
416 /// requires. Concrete impls should override when the extension struct exposes its
417 /// base packings directly, e.g. `Self::new(F::Packing::pack_columns_fn(|l| f(l).value))`.
418 #[inline]
419 #[must_use]
420 fn from_ext_fn(f: impl Fn(usize) -> ExtField) -> Self {
421 Self::from_basis_coefficients_fn(|d| {
422 BaseField::Packing::from_fn(|lane| f(lane).as_basis_coefficients_slice()[d])
423 })
424 }
425
426 /// Pack a length-`WIDTH` slice of extension field elements into one packed extension.
427 ///
428 /// ## Panics
429 /// Panics if `slice.len() != BaseField::Packing::WIDTH`.
430 #[inline]
431 #[must_use]
432 fn from_ext_slice(slice: &[ExtField]) -> Self {
433 assert_eq!(slice.len(), BaseField::Packing::WIDTH);
434 Self::from_ext_fn(|lane| slice[lane])
435 }
436
437 /// Pack `N` columns from `W` rows of extension field elements into `N` packed extensions.
438 ///
439 /// This is the extension-field analog of [`PackedValue::pack_columns`]: given `W` rows
440 /// of `N` extension elements, lane `lane` of output column `col` is `rows[lane][col]`.
441 ///
442 /// ## Panics
443 /// Panics if `rows.len() != BaseField::Packing::WIDTH`.
444 #[inline]
445 #[must_use]
446 fn pack_ext_columns<const N: usize>(rows: &[[ExtField; N]]) -> [Self; N] {
447 assert_eq!(rows.len(), BaseField::Packing::WIDTH);
448 array::from_fn(|col| Self::from_ext_fn(|lane| rows[lane][col]))
449 }
450
451 /// Pack `N` columns using a closure that produces each row.
452 ///
453 /// Analog of [`PackedValue::pack_columns_fn`].
454 #[inline]
455 #[must_use]
456 fn pack_ext_columns_fn<const N: usize>(row_fn: impl Fn(usize) -> [ExtField; N]) -> [Self; N] {
457 array::from_fn(|col| Self::from_ext_fn(|lane| row_fn(lane)[col]))
458 }
459
460 /// Extract the extension field element at the given SIMD lane.
461 #[inline]
462 #[must_use]
463 fn extract(&self, lane: usize) -> ExtField {
464 ExtField::from_basis_coefficients_fn(|d| {
465 self.as_basis_coefficients_slice()[d].as_slice()[lane]
466 })
467 }
468
469 /// Accumulate `value` into a single SIMD lane, leaving the other `W - 1` lanes unchanged.
470 ///
471 /// This is the accumulating dual of the per-lane read [`PackedFieldExtension::extract`].
472 /// It scatters one scalar extension element into a packed buffer, one lane at a time.
473 ///
474 /// The default rebuilds a full packed element and adds it.
475 /// Concrete types override it to touch only the `D` base lanes at the target lane.
476 #[inline]
477 fn add_assign_lane(&mut self, lane: usize, value: ExtField) {
478 *self += Self::from_ext_fn(|l| if l == lane { value } else { ExtField::ZERO });
479 }
480
481 /// Write all `W` lanes into the given slice.
482 ///
483 /// This is the extension-field analog of [`PackedValue::as_slice`], but the lanes of
484 /// a packed extension are not contiguous in memory (the layout is `[[F; W]; D]`,
485 /// indexed first by basis coefficient), so the lanes must be copied rather than
486 /// borrowed.
487 ///
488 /// ## Panics
489 /// Panics if `out.len() != BaseField::Packing::WIDTH`.
490 #[inline]
491 fn to_ext_slice(&self, out: &mut [ExtField]) {
492 assert_eq!(out.len(), BaseField::Packing::WIDTH);
493 for (lane, slot) in out.iter_mut().enumerate() {
494 *slot = self.extract(lane);
495 }
496 }
497
498 /// Unpack `N` packed extensions into `W` rows of `N` extension elements.
499 ///
500 /// Inverse of [`PackedFieldExtension::pack_ext_columns`]. Lane `lane` of input
501 /// column `col` is written to `rows[lane][col]`.
502 ///
503 /// ## Panics
504 /// Panics if `rows.len() != BaseField::Packing::WIDTH`.
505 #[inline]
506 fn unpack_ext_into<const N: usize>(packed: &[Self; N], rows: &mut [[ExtField; N]]) {
507 assert_eq!(rows.len(), BaseField::Packing::WIDTH);
508 for (lane, row) in rows.iter_mut().enumerate() {
509 *row = array::from_fn(|col| {
510 ExtField::from_basis_coefficients_fn(|d| {
511 packed[col].as_basis_coefficients_slice()[d].as_slice()[lane]
512 })
513 });
514 }
515 }
516
517 /// Iterator equivalent of [`PackedFieldExtension::unpack_ext_into`].
518 ///
519 /// Yields `WIDTH` rows of `N` extension elements without requiring a pre-allocated
520 /// buffer. Analog of [`PackedValue::unpack_iter`].
521 #[inline]
522 fn unpack_ext_iter<const N: usize>(packed: [Self; N]) -> impl Iterator<Item = [ExtField; N]> {
523 (0..BaseField::Packing::WIDTH).map(move |lane| {
524 array::from_fn(|col| {
525 ExtField::from_basis_coefficients_fn(|d| {
526 packed[col].as_basis_coefficients_slice()[d].as_slice()[lane]
527 })
528 })
529 })
530 }
531
532 /// Convert an iterator of packed extension field elements to an iterator of
533 /// extension field elements (flat — one `ExtField` per lane per packed value).
534 #[inline]
535 #[must_use]
536 fn to_ext_iter(iter: impl IntoIterator<Item = Self>) -> impl Iterator<Item = ExtField> {
537 iter.into_iter()
538 .flat_map(|x| (0..BaseField::Packing::WIDTH).map(move |lane| x.extract(lane)))
539 }
540
541 /// Unpacks a packed row-major matrix into its scalar transpose, in one pass.
542 ///
543 /// `src` is a row-major matrix of logical scalars, `src_width` columns wide.
544 /// `WIDTH` consecutive logical scalars share one packed element:
545 ///
546 /// ```text
547 /// src_logical[i] = src[i / WIDTH].extract(i % WIDTH)
548 /// ```
549 ///
550 /// The scalar transpose lands in `dst`, now `src_height` columns wide:
551 ///
552 /// ```text
553 /// dst[c * src_height + r] = src_logical[r * src_width + c]
554 /// ```
555 ///
556 /// Every scalar is read once and written once: half the traffic of unpack-then-transpose.
557 ///
558 /// ## Panics
559 /// - `dst.len()` must equal `src.len() * WIDTH`.
560 /// - `src_width` must divide `dst.len()`.
561 fn unpack_transpose_into(src: &[Self], dst: &mut [ExtField], src_width: usize) {
562 let w = BaseField::Packing::WIDTH;
563 assert!(src_width != 0);
564 assert_eq!(dst.len(), src.len() * w);
565 assert!(
566 dst.len().is_multiple_of(src_width),
567 "src_width must divide the scalar count"
568 );
569 let src_height = dst.len() / src_width;
570
571 // Fast path: both dimensions split into whole `W`-blocks.
572 // src_width % W == 0 -> every row starts on a packed boundary
573 // src_height % W == 0 -> every dst run is exactly `W` long, no overrun
574 if w > 1 && src_width.is_multiple_of(w) && src_height.is_multiple_of(w) {
575 let packed_width = src_width / w;
576 // Walk `W x W` scalar blocks; reads stream `W` rows, writes are contiguous runs.
577 for r0 in (0..src_height).step_by(w) {
578 for cp in 0..packed_width {
579 for l in 0..w {
580 // Lane `l` of packed (row r0+j, packed-col cp) = src_logical[r0+j][cp*W + l].
581 // It lands at column `cp*W + l`, row `r0+j`.
582 let run = &mut dst[(cp * w + l) * src_height + r0..][..w];
583 for (j, slot) in run.iter_mut().enumerate() {
584 *slot = src[(r0 + j) * packed_width + cp].extract(l);
585 }
586 }
587 }
588 }
589 } else {
590 // Fallback: gather each scalar straight to its transposed slot.
591 for r in 0..src_height {
592 for c in 0..src_width {
593 let idx = r * src_width + c;
594 dst[c * src_height + r] = src[idx / w].extract(idx % w);
595 }
596 }
597 }
598 }
599
600 /// Similar to `packed_powers`, construct an iterator which returns
601 /// powers of `base` packed into `PackedFieldExtension` elements.
602 #[must_use]
603 fn packed_ext_powers(base: ExtField) -> Powers<Self>;
604
605 /// Similar to `packed_ext_powers` but only returns `unpacked_len` powers of `base`.
606 ///
607 /// Note that the length of the returned iterator will be `unpacked_len / WIDTH` and
608 /// not `len` as the iterator is over packed extension field elements. If `unpacked_len`
609 /// is not divisible by `WIDTH`, `unpacked_len` will be rounded up to the next multiple of `WIDTH`.
610 #[must_use]
611 fn packed_ext_powers_capped(base: ExtField, unpacked_len: usize) -> impl Iterator<Item = Self> {
612 Self::packed_ext_powers(base).take(unpacked_len.div_ceil(BaseField::Packing::WIDTH))
613 }
614}
615
616unsafe impl<T: Packable> PackedValue for T {
617 type Value = Self;
618
619 const WIDTH: usize = 1;
620
621 #[inline]
622 fn from_slice(slice: &[Self::Value]) -> &Self {
623 assert_eq!(slice.len(), Self::WIDTH);
624 &slice[0]
625 }
626
627 #[inline]
628 fn from_slice_mut(slice: &mut [Self::Value]) -> &mut Self {
629 assert_eq!(slice.len(), Self::WIDTH);
630 &mut slice[0]
631 }
632
633 #[inline]
634 fn from_fn<Fn>(mut f: Fn) -> Self
635 where
636 Fn: FnMut(usize) -> Self::Value,
637 {
638 f(0)
639 }
640
641 #[inline]
642 fn as_slice(&self) -> &[Self::Value] {
643 slice::from_ref(self)
644 }
645
646 #[inline]
647 fn as_slice_mut(&mut self) -> &mut [Self::Value] {
648 slice::from_mut(self)
649 }
650}
651
652unsafe impl<F: Field> PackedField for F {
653 type Scalar = Self;
654}
655
656unsafe impl<F: Field> PackedFieldPow2 for F {
657 #[inline]
658 fn interleave(&self, other: Self, block_len: usize) -> (Self, Self) {
659 match block_len {
660 1 => (*self, other),
661 _ => panic!("unsupported block length"),
662 }
663 }
664}
665
666impl<F: Field> PackedFieldExtension<F, F> for F::Packing {
667 #[inline]
668 fn from_ext_fn(f: impl Fn(usize) -> F) -> Self {
669 F::Packing::from_fn(f)
670 }
671
672 #[inline]
673 fn from_ext_slice(slice: &[F]) -> Self {
674 *F::Packing::from_slice(slice)
675 }
676
677 #[inline]
678 fn packed_ext_powers(base: F) -> Powers<Self> {
679 F::Packing::packed_powers(base)
680 }
681
682 #[inline]
683 fn add_assign_lane(&mut self, lane: usize, value: F) {
684 // Degree-1 case: the lane is a single base element.
685 self.as_slice_mut()[lane] += value;
686 }
687}
688
689impl Packable for u8 {}
690
691impl Packable for u16 {}
692
693impl Packable for u32 {}
694
695impl Packable for u64 {}
696
697impl Packable for u128 {}