binius_field/
packed.rs

1// Copyright 2023-2025 Irreducible Inc.
2// Copyright 2026 The Binius Developers
3
4//! Traits for packed field elements which support SIMD implementations.
5//!
6//! Interfaces are derived from [`plonky2`](https://github.com/mir-protocol/plonky2).
7
8use std::{
9	fmt::Debug,
10	iter,
11	ops::{Add, AddAssign, Mul, MulAssign, Sub, SubAssign},
12};
13
14use binius_utils::iter::IterExtensions;
15use bytemuck::Zeroable;
16
17use super::{Random, arithmetic_traits::Square};
18use crate::{BinaryField, Divisible, Maskable, WideMul, field::FieldOps};
19
20/// A packed field represents a vector of underlying field elements.
21///
22/// Arithmetic operations on packed field elements can be accelerated with SIMD CPU instructions.
23/// The vector width is a constant, `WIDTH`. This trait requires that the width must be a power of
24/// two.
25pub trait PackedField:
26	Default
27	+ Debug
28	+ Clone
29	+ Copy
30	+ Eq
31	+ Sized
32	+ FieldOps
33	+ Add<Self::Scalar, Output = Self>
34	+ Sub<Self::Scalar, Output = Self>
35	+ Mul<Self::Scalar, Output = Self>
36	+ AddAssign<Self::Scalar>
37	+ SubAssign<Self::Scalar>
38	+ MulAssign<Self::Scalar>
39	+ Send
40	+ Sync
41	+ Zeroable
42	+ Random
43	+ WideMul<Output: Debug + Send + Sync + 'static>
44	+ 'static
45	// A packed field divides into its `WIDTH` scalars. Scalar element access (`get`/`set` and
46	// their `_unchecked` variants), broadcast, and the scalar iterators are all provided by this
47	// supertrait.
48	+ Divisible<Self::Scalar>
49	// A packed field supports branchless per-lane masking over its scalars.
50	+ Maskable<Self::Scalar>
51{
52	/// Base-2 logarithm of the number of field elements packed into one packed element.
53	///
54	/// This is the number of scalars the packed field divides into, i.e. its `Divisible` log-count.
55	const LOG_WIDTH: usize = <Self as Divisible<Self::Scalar>>::LOG_N;
56
57	/// The number of field elements packed into one packed element.
58	///
59	/// WIDTH is guaranteed to equal 2^LOG_WIDTH.
60	const WIDTH: usize = 1 << Self::LOG_WIDTH;
61
62	#[inline]
63	fn into_iter(self) -> impl Iterator<Item = Self::Scalar> + Send + Clone {
64		(0..Self::WIDTH).map_skippable(move |i|
65			// Safety: `i` is always less than `WIDTH`
66			unsafe { self.get_unchecked(i) })
67	}
68
69	#[inline]
70	fn iter(&self) -> impl Iterator<Item = Self::Scalar> + Send + Clone + '_ {
71		(0..Self::WIDTH).map_skippable(move |i|
72			// Safety: `i` is always less than `WIDTH`
73			unsafe { self.get_unchecked(i) })
74	}
75
76	#[inline]
77	fn iter_slice(slice: &[Self]) -> impl Iterator<Item = Self::Scalar> + Send + Clone + '_ {
78		slice.iter().flat_map(Self::iter)
79	}
80
81	/// Initialize zero position with `scalar`, set other elements to zero.
82	#[inline(always)]
83	fn set_single(scalar: Self::Scalar) -> Self {
84		let mut result = Self::default();
85		result.set(0, scalar);
86		result
87	}
88
89	/// Construct a packed field element from a function that returns scalar values by index.
90	fn from_fn(f: impl FnMut(usize) -> Self::Scalar) -> Self;
91
92	/// Construct a packed field element from a sequence of scalars.
93	///
94	/// If the number of values in the sequence is less than the packing width, the remaining
95	/// elements are set to zero. If greater than the packing width, the excess elements are
96	/// ignored.
97	#[inline]
98	fn from_scalars(values: impl IntoIterator<Item = Self::Scalar>) -> Self {
99		let mut result = Self::default();
100		for (i, val) in values.into_iter().take(Self::WIDTH).enumerate() {
101			result.set(i, val);
102		}
103		result
104	}
105
106	/// Returns the value to the power `exp`.
107	fn pow(self, exp: u64) -> Self {
108		let mut res = Self::one();
109		for i in (0..64).rev() {
110			res = Square::square(res);
111			if ((exp >> i) & 1) == 1 {
112				res.mul_assign(self)
113			}
114		}
115		res
116	}
117
118	/// Interleaves blocks of this packed vector with another packed vector.
119	///
120	/// The operation can be seen as stacking the two vectors, dividing them into 2x2 matrices of
121	/// blocks, where each block is 2^`log_block_width` elements, and transposing the matrices.
122	///
123	/// Consider this example, where `LOG_WIDTH` is 3 and `log_block_len` is 1:
124	///     A = [a0, a1, a2, a3, a4, a5, a6, a7]
125	///     B = [b0, b1, b2, b3, b4, b5, b6, b7]
126	///
127	/// The interleaved result is
128	///     A' = [a0, a1, b0, b1, a4, a5, b4, b5]
129	///     B' = [a2, a3, b2, b3, a6, a7, b6, b7]
130	///
131	/// ## Preconditions
132	/// * `log_block_len` must be strictly less than `LOG_WIDTH`.
133	fn interleave(self, other: Self, log_block_len: usize) -> (Self, Self);
134
135	/// Unzips interleaved blocks of this packed vector with another packed vector.
136	///
137	/// Consider this example, where `LOG_WIDTH` is 3 and `log_block_len` is 1:
138	///    A = [a0, a1, b0, b1, a2, a3, b2, b3]
139	///    B = [a4, a5, b4, b5, a6, a7, b6, b7]
140	///
141	/// The transposed result is
142	///    A' = [a0, a1, a2, a3, a4, a5, a6, a7]
143	///    B' = [b0, b1, b2, b3, b4, b5, b6, b7]
144	///
145	/// ## Preconditions
146	/// * `log_block_len` must be strictly less than `LOG_WIDTH`.
147	fn unzip(self, other: Self, log_block_len: usize) -> (Self, Self);
148
149	/// Spread takes a block of elements within a packed field and repeats them to the full packing
150	/// width.
151	///
152	/// Spread can be seen as an extension of the functionality of [`Divisible::broadcast`].
153	///
154	/// ## Examples
155	///
156	/// ```
157	/// use binius_field::{BinaryField1b, PackedField, PackedBinaryField8x1b};
158	///
159	/// let input =
160	///     PackedBinaryField8x1b::from_scalars([0, 1, 0, 1, 0, 1, 0, 1].map(BinaryField1b::from));
161	/// assert_eq!(
162	///     input.spread(0, 1),
163	///     PackedBinaryField8x1b::from_scalars([1, 1, 1, 1, 1, 1, 1, 1].map(BinaryField1b::from))
164	/// );
165	/// assert_eq!(
166	///     input.spread(1, 0),
167	///     PackedBinaryField8x1b::from_scalars([0, 0, 0, 0, 1, 1, 1, 1].map(BinaryField1b::from))
168	/// );
169	/// assert_eq!(
170	///     input.spread(2, 0),
171	///     PackedBinaryField8x1b::from_scalars([0, 0, 1, 1, 0, 0, 1, 1].map(BinaryField1b::from))
172	/// );
173	/// assert_eq!(input.spread(3, 0), input);
174	/// ```
175	///
176	/// ## Preconditions
177	///
178	/// * `log_block_len` must be less than or equal to `LOG_WIDTH`.
179	/// * `block_idx` must be less than `2^(Self::LOG_WIDTH - log_block_len)`.
180	#[inline]
181	fn spread(self, log_block_len: usize, block_idx: usize) -> Self {
182		assert!(log_block_len <= Self::LOG_WIDTH);
183		assert!(block_idx < 1 << (Self::LOG_WIDTH - log_block_len));
184
185		// Safety: is guaranteed by the preconditions.
186		unsafe { self.spread_unchecked(log_block_len, block_idx) }
187	}
188
189	/// Unsafe version of [`Self::spread`].
190	///
191	/// # Safety
192	/// The caller must ensure that `log_block_len` is less than or equal to `LOG_WIDTH` and
193	/// `block_idx` is less than `2^(Self::LOG_WIDTH - log_block_len)`.
194	#[inline]
195	unsafe fn spread_unchecked(self, log_block_len: usize, block_idx: usize) -> Self {
196		let block_len = 1 << log_block_len;
197		let repeat = 1 << (Self::LOG_WIDTH - log_block_len);
198
199		Self::from_scalars(
200			self.iter()
201				.skip(block_idx * block_len)
202				.take(block_len)
203				.flat_map(|elem| iter::repeat_n(elem, repeat)),
204		)
205	}
206}
207
208#[inline(always)]
209pub fn get_packed_slice<P: PackedField>(packed: &[P], i: usize) -> P::Scalar {
210	assert!(i >> P::LOG_WIDTH < packed.len(), "index out of bounds");
211
212	unsafe { get_packed_slice_unchecked(packed, i) }
213}
214
215/// Returns the scalar at the given index without bounds checking.
216/// # Safety
217/// The caller must ensure that `i` is less than `P::WIDTH * packed.len()`.
218#[inline(always)]
219pub unsafe fn get_packed_slice_unchecked<P: PackedField>(packed: &[P], i: usize) -> P::Scalar {
220	// TODO: Consider putting a get_in_slice method on Divisible
221
222	// Safety:
223	// - `i / P::WIDTH` is within the bounds of `packed` if `i` is less than `P::WIDTH *
224	//   packed.len()`
225	// - `i % P::WIDTH` is always less than `P::WIDTH
226	unsafe {
227		packed
228			.get_unchecked(i >> P::LOG_WIDTH)
229			.get_unchecked(i % P::WIDTH)
230	}
231}
232
233/// Sets the scalar at the given index without bounds checking.
234/// # Safety
235/// The caller must ensure that `i` is less than `P::WIDTH * packed.len()`.
236#[inline]
237pub unsafe fn set_packed_slice_unchecked<P: PackedField>(
238	packed: &mut [P],
239	i: usize,
240	scalar: P::Scalar,
241) {
242	// TODO: Consider putting a set_in_slice method on Divisible
243
244	// Safety: if `i` is less than `P::WIDTH * packed.len()`, then
245	// - `i / P::WIDTH` is within the bounds of `packed`
246	// - `i % P::WIDTH` is always less than `P::WIDTH
247	unsafe {
248		packed
249			.get_unchecked_mut(i >> P::LOG_WIDTH)
250			.set_unchecked(i % P::WIDTH, scalar)
251	}
252}
253
254/// A helper trait to make the generic bounds shorter
255pub trait PackedBinaryField: PackedField<Scalar: BinaryField> {}
256
257impl<PT> PackedBinaryField for PT where PT: PackedField<Scalar: BinaryField> {}
258
259#[cfg(test)]
260mod tests {
261	use rand::prelude::*;
262
263	use crate::{
264		AESTowerField8b, BinaryField1b, BinaryField128bGhash, PackedAESBinaryField1x8b,
265		PackedAESBinaryField16x8b, PackedAESBinaryField32x8b, PackedAESBinaryField64x8b,
266		PackedBinaryField1x1b, PackedBinaryField2x1b, PackedBinaryField4x1b, PackedBinaryField8x1b,
267		PackedBinaryField16x1b, PackedBinaryField32x1b, PackedBinaryField64x1b,
268		PackedBinaryField128x1b, PackedBinaryField256x1b, PackedBinaryField512x1b,
269		PackedBinaryGhash1x128b, PackedBinaryGhash2x128b, PackedBinaryGhash4x128b, PackedField,
270		SlicedGhashSq1x256b, SlicedGhashSq2x256b, SlicedGhashSq4x256b,
271	};
272
273	trait PackedFieldTest {
274		fn run<P: PackedField>(&self);
275	}
276
277	/// Run the test for all the packed fields defined in this crate.
278	fn run_for_all_packed_fields(test: &impl PackedFieldTest) {
279		// B1
280		test.run::<BinaryField1b>();
281		test.run::<PackedBinaryField1x1b>();
282		test.run::<PackedBinaryField2x1b>();
283		test.run::<PackedBinaryField4x1b>();
284		test.run::<PackedBinaryField8x1b>();
285		test.run::<PackedBinaryField16x1b>();
286		test.run::<PackedBinaryField32x1b>();
287		test.run::<PackedBinaryField64x1b>();
288		test.run::<PackedBinaryField128x1b>();
289		test.run::<PackedBinaryField256x1b>();
290		test.run::<PackedBinaryField512x1b>();
291
292		// AES
293		test.run::<AESTowerField8b>();
294		test.run::<PackedAESBinaryField1x8b>();
295		test.run::<PackedAESBinaryField16x8b>();
296		test.run::<PackedAESBinaryField32x8b>();
297		test.run::<PackedAESBinaryField64x8b>();
298
299		// GHASH
300		test.run::<BinaryField128bGhash>();
301		test.run::<PackedBinaryGhash1x128b>();
302		test.run::<PackedBinaryGhash2x128b>();
303		test.run::<PackedBinaryGhash4x128b>();
304
305		// GHASH² in a sliced layout
306		test.run::<SlicedGhashSq1x256b>();
307		test.run::<SlicedGhashSq2x256b>();
308		test.run::<SlicedGhashSq4x256b>();
309	}
310
311	fn check_value_iteration<P: PackedField>(mut rng: impl Rng) {
312		let packed = P::random(&mut rng);
313		let mut iter = packed.iter();
314		for i in 0..P::WIDTH {
315			assert_eq!(packed.get(i), iter.next().unwrap());
316		}
317		assert!(iter.next().is_none());
318	}
319
320	fn check_ref_iteration<P: PackedField>(mut rng: impl Rng) {
321		let packed = P::random(&mut rng);
322		let mut iter = packed.into_iter();
323		for i in 0..P::WIDTH {
324			assert_eq!(packed.get(i), iter.next().unwrap());
325		}
326		assert!(iter.next().is_none());
327	}
328
329	struct PackedFieldIterationTest;
330
331	impl PackedFieldTest for PackedFieldIterationTest {
332		fn run<P: PackedField>(&self) {
333			let mut rng = StdRng::seed_from_u64(0);
334
335			check_value_iteration::<P>(&mut rng);
336			check_ref_iteration::<P>(&mut rng);
337		}
338	}
339
340	#[test]
341	fn test_iteration() {
342		run_for_all_packed_fields(&PackedFieldIterationTest);
343	}
344}
binius_field/packed.rs

binius_field/
packed.rs
