binius_field/
packed_extension.rs

1// Copyright 2023-2025 Irreducible Inc.
2
3use crate::{
4	as_packed_field::PackScalar,
5	underlier::{Divisible, WithUnderlier},
6	ExtensionField, Field, PackedField,
7};
8
9/// A [`PackedField`] that can be safely cast to indexable slices of scalars.
10///
11/// Not all packed fields can index individual scalar elements. Notably, packed fields of
12/// $\mathbb{F}_2$ elements can pack multiple scalars into a single byte.
13///
14///
15/// # Safety
16///
17/// In order for the above relation to be guaranteed, the memory representation of a slice of
18/// `PackedExtensionIndexable` elements must be the same as a slice of the underlying scalar
19/// elements, differing only in the slice lengths.
20pub unsafe trait PackedFieldIndexable: PackedField {
21	fn unpack_scalars(packed: &[Self]) -> &[Self::Scalar];
22	fn unpack_scalars_mut(packed: &mut [Self]) -> &mut [Self::Scalar];
23}
24
25unsafe impl<S, P> PackedFieldIndexable for P
26where
27	S: Field,
28	P: PackedDivisible<S, Scalar = S>,
29{
30	fn unpack_scalars(packed: &[Self]) -> &[Self::Scalar] {
31		P::divide(packed)
32	}
33
34	fn unpack_scalars_mut(packed: &mut [Self]) -> &mut [Self::Scalar] {
35		P::divide_mut(packed)
36	}
37}
38
39/// Trait represents a relationship between a packed struct of field elements and a packed struct
40/// of elements from an extension field.
41///
42/// This trait guarantees that one packed type has the same
43/// memory representation as the other, differing only in the scalar type and preserving the order
44/// of smaller elements.
45///
46/// This trait relation guarantees that the following iterators yield the same sequence of scalar
47/// elements:
48///
49/// ```
50/// use binius_field::{ExtensionField, PackedExtension, PackedField, Field};
51///
52/// fn ext_then_bases<'a, F, PE>(packed: &'a PE) -> impl Iterator<Item=F> + 'a
53///     where
54///         PE: PackedField<Scalar: ExtensionField<F>>,
55///         F: Field,
56/// {
57///     packed.iter().flat_map(|ext| ext.into_iter_bases())
58/// }
59///
60/// fn cast_then_iter<'a, F, PE>(packed: &'a PE) -> impl Iterator<Item=F> + 'a
61///     where
62///         PE: PackedExtension<F>,
63///         F: Field,
64/// {
65///     PE::cast_base_ref(packed).into_iter()
66/// }
67/// ```
68///
69/// # Safety
70///
71/// In order for the above relation to be guaranteed, the memory representation of
72/// `PackedExtensionField` element must be the same as a slice of the underlying `PackedField`
73/// element.
74pub trait PackedExtension<FS: Field>: PackedField<Scalar: ExtensionField<FS>> {
75	type PackedSubfield: PackedField<Scalar = FS>;
76
77	fn cast_bases(packed: &[Self]) -> &[Self::PackedSubfield];
78	fn cast_bases_mut(packed: &mut [Self]) -> &mut [Self::PackedSubfield];
79
80	fn cast_exts(packed: &[Self::PackedSubfield]) -> &[Self];
81	fn cast_exts_mut(packed: &mut [Self::PackedSubfield]) -> &mut [Self];
82
83	fn cast_base(self) -> Self::PackedSubfield;
84	fn cast_base_ref(&self) -> &Self::PackedSubfield;
85	fn cast_base_mut(&mut self) -> &mut Self::PackedSubfield;
86
87	fn cast_ext(base: Self::PackedSubfield) -> Self;
88	fn cast_ext_ref(base: &Self::PackedSubfield) -> &Self;
89	fn cast_ext_mut(base: &mut Self::PackedSubfield) -> &mut Self;
90
91	#[inline(always)]
92	fn cast_base_arr<const N: usize>(packed: [Self; N]) -> [Self::PackedSubfield; N] {
93		packed.map(Self::cast_base)
94	}
95
96	#[inline(always)]
97	fn cast_base_arr_ref<const N: usize>(packed: &[Self; N]) -> &[Self::PackedSubfield; N] {
98		Self::cast_bases(packed)
99			.try_into()
100			.expect("array has size N")
101	}
102
103	#[inline(always)]
104	fn cast_base_arr_mut<const N: usize>(packed: &mut [Self; N]) -> &mut [Self::PackedSubfield; N] {
105		Self::cast_bases_mut(packed)
106			.try_into()
107			.expect("array has size N")
108	}
109
110	#[inline(always)]
111	fn cast_ext_arr<const N: usize>(packed: [Self::PackedSubfield; N]) -> [Self; N] {
112		packed.map(Self::cast_ext)
113	}
114
115	#[inline(always)]
116	fn cast_ext_arr_ref<const N: usize>(packed: &[Self::PackedSubfield; N]) -> &[Self; N] {
117		Self::cast_exts(packed)
118			.try_into()
119			.expect("array has size N")
120	}
121
122	#[inline(always)]
123	fn cast_ext_arr_mut<const N: usize>(packed: &mut [Self::PackedSubfield; N]) -> &mut [Self; N] {
124		Self::cast_exts_mut(packed)
125			.try_into()
126			.expect("array has size N")
127	}
128}
129
130impl<PT, FS> PackedExtension<FS> for PT
131where
132	FS: Field,
133	PT: PackedField<Scalar: ExtensionField<FS>> + WithUnderlier<Underlier: PackScalar<FS>>,
134{
135	type PackedSubfield = <PT::Underlier as PackScalar<FS>>::Packed;
136
137	fn cast_bases(packed: &[Self]) -> &[Self::PackedSubfield] {
138		Self::PackedSubfield::from_underliers_ref(Self::to_underliers_ref(packed))
139	}
140
141	fn cast_bases_mut(packed: &mut [Self]) -> &mut [Self::PackedSubfield] {
142		Self::PackedSubfield::from_underliers_ref_mut(Self::to_underliers_ref_mut(packed))
143	}
144
145	fn cast_exts(base: &[Self::PackedSubfield]) -> &[Self] {
146		Self::from_underliers_ref(Self::PackedSubfield::to_underliers_ref(base))
147	}
148
149	fn cast_exts_mut(base: &mut [Self::PackedSubfield]) -> &mut [Self] {
150		Self::from_underliers_ref_mut(Self::PackedSubfield::to_underliers_ref_mut(base))
151	}
152
153	fn cast_base(self) -> Self::PackedSubfield {
154		Self::PackedSubfield::from_underlier(self.to_underlier())
155	}
156
157	fn cast_base_ref(&self) -> &Self::PackedSubfield {
158		Self::PackedSubfield::from_underlier_ref(self.to_underlier_ref())
159	}
160
161	fn cast_base_mut(&mut self) -> &mut Self::PackedSubfield {
162		Self::PackedSubfield::from_underlier_ref_mut(self.to_underlier_ref_mut())
163	}
164
165	fn cast_ext(base: Self::PackedSubfield) -> Self {
166		Self::from_underlier(base.to_underlier())
167	}
168
169	fn cast_ext_ref(base: &Self::PackedSubfield) -> &Self {
170		Self::from_underlier_ref(base.to_underlier_ref())
171	}
172
173	fn cast_ext_mut(base: &mut Self::PackedSubfield) -> &mut Self {
174		Self::from_underlier_ref_mut(base.to_underlier_ref_mut())
175	}
176}
177
178/// Convenient type alias that returns the packed field type for the scalar field `F` and packed
179/// extension `P`.
180pub type PackedSubfield<P, F> = <P as PackedExtension<F>>::PackedSubfield;
181
182/// Recast a packed field from one subfield of a packed extension to another.
183pub fn recast_packed<P, FSub1, FSub2>(elem: PackedSubfield<P, FSub1>) -> PackedSubfield<P, FSub2>
184where
185	P: PackedField + PackedExtension<FSub1> + PackedExtension<FSub2>,
186	P::Scalar: ExtensionField<FSub1> + ExtensionField<FSub2>,
187	FSub1: Field,
188	FSub2: Field,
189{
190	<P as PackedExtension<FSub2>>::cast_base(<P as PackedExtension<FSub1>>::cast_ext(elem))
191}
192
193/// Recast a slice of packed field elements from one subfield of a packed extension to another.
194pub fn recast_packed_slice<P, FSub1, FSub2>(
195	elems: &[PackedSubfield<P, FSub1>],
196) -> &[PackedSubfield<P, FSub2>]
197where
198	P: PackedField + PackedExtension<FSub1> + PackedExtension<FSub2>,
199	P::Scalar: ExtensionField<FSub1> + ExtensionField<FSub2>,
200	FSub1: Field,
201	FSub2: Field,
202{
203	<P as PackedExtension<FSub2>>::cast_bases(<P as PackedExtension<FSub1>>::cast_exts(elems))
204}
205
206/// Recast a mutable slice of packed field elements from one subfield of a packed extension to
207/// another.
208pub fn recast_packed_mut<P, FSub1, FSub2>(
209	elems: &mut [PackedSubfield<P, FSub1>],
210) -> &mut [PackedSubfield<P, FSub2>]
211where
212	P: PackedField + PackedExtension<FSub1> + PackedExtension<FSub2>,
213	P::Scalar: ExtensionField<FSub1> + ExtensionField<FSub2>,
214	FSub1: Field,
215	FSub2: Field,
216{
217	<P as PackedExtension<FSub2>>::cast_bases_mut(<P as PackedExtension<FSub1>>::cast_exts_mut(
218		elems,
219	))
220}
221
222/// This trait is a shorthand for the case `PackedExtension<P::Scalar, PackedSubfield = P>` which is a
223/// quite common case in our codebase.
224pub trait RepackedExtension<P: PackedField>:
225	PackedField<Scalar: ExtensionField<P::Scalar>> + PackedExtension<P::Scalar, PackedSubfield = P>
226{
227}
228
229impl<PT1, PT2> RepackedExtension<PT1> for PT2
230where
231	PT1: PackedField,
232	PT2: PackedExtension<PT1::Scalar, PackedSubfield = PT1, Scalar: ExtensionField<PT1::Scalar>>,
233{
234}
235
236/// This trait adds shortcut methods for the case `PackedExtension<F, PackedSubfield: PackedFieldIndexable>` which is a
237/// quite common case in our codebase.
238pub trait PackedExtensionIndexable<F: Field>:
239	PackedExtension<F, PackedSubfield: PackedFieldIndexable> + PackedField<Scalar: ExtensionField<F>>
240{
241	fn unpack_base_scalars(packed: &[Self]) -> &[F] {
242		Self::PackedSubfield::unpack_scalars(Self::cast_bases(packed))
243	}
244
245	fn unpack_base_scalars_mut(packed: &mut [Self]) -> &mut [F] {
246		Self::PackedSubfield::unpack_scalars_mut(Self::cast_bases_mut(packed))
247	}
248}
249
250impl<F, PT> PackedExtensionIndexable<F> for PT
251where
252	F: Field,
253	PT: PackedExtension<F, PackedSubfield: PackedFieldIndexable>,
254{
255}
256
257/// Trait represents a relationship between a packed struct of field elements and a smaller packed
258/// struct the same field elements.
259///
260/// This trait can be used to safely cast memory slices from larger packed fields to smaller ones.
261///
262/// # Safety
263///
264/// In order for the above relation to be guaranteed, the memory representation of a slice of
265/// `PackedDivisible` elements must be the same as a slice of the underlying `PackedField`
266/// elements, differing only in the slice lengths.
267pub unsafe trait PackedDivisible<P>: PackedField
268where
269	P: PackedField<Scalar = Self::Scalar>,
270{
271	fn divide(packed: &[Self]) -> &[P];
272	fn divide_mut(packed: &mut [Self]) -> &mut [P];
273}
274
275unsafe impl<PT1, PT2> PackedDivisible<PT2> for PT1
276where
277	PT2: PackedField + WithUnderlier,
278	PT1: PackedField<Scalar = PT2::Scalar> + WithUnderlier<Underlier: Divisible<PT2::Underlier>>,
279{
280	fn divide(packed: &[Self]) -> &[PT2] {
281		let underliers = PT1::to_underliers_ref(packed);
282		let underliers: &[PT2::Underlier] = PT1::Underlier::split_slice(underliers);
283		PT2::from_underliers_ref(underliers)
284	}
285
286	fn divide_mut(packed: &mut [Self]) -> &mut [PT2] {
287		let underliers = PT1::to_underliers_ref_mut(packed);
288		let underliers: &mut [PT2::Underlier] = PT1::Underlier::split_slice_mut(underliers);
289		PT2::from_underliers_ref_mut(underliers)
290	}
291}
binius_field/packed_extension.rs

binius_field/
packed_extension.rs
