binius_field/
packed_extension.rs

1// Copyright 2023-2025 Irreducible Inc.
2
3use crate::{
4	ExtensionField, Field, PackedField, as_packed_field::PackScalar, underlier::WithUnderlier,
5};
6
7/// Trait represents a relationship between a packed struct of field elements and a packed struct
8/// of elements from an extension field.
9///
10/// This trait guarantees that one packed type has the same
11/// memory representation as the other, differing only in the scalar type and preserving the order
12/// of smaller elements.
13///
14/// This trait relation guarantees that the following iterators yield the same sequence of scalar
15/// elements:
16///
17/// ```
18/// use binius_field::{ExtensionField, PackedExtension, PackedField, Field};
19///
20/// fn ext_then_bases<'a, F, PE>(packed: &'a PE) -> impl Iterator<Item=F> + 'a
21///     where
22///         PE: PackedField<Scalar: ExtensionField<F>>,
23///         F: Field,
24/// {
25///     packed.iter().flat_map(|ext| ext.into_iter_bases())
26/// }
27///
28/// fn cast_then_iter<'a, F, PE>(packed: &'a PE) -> impl Iterator<Item=F> + 'a
29///     where
30///         PE: PackedExtension<F>,
31///         F: Field,
32/// {
33///     PE::cast_base_ref(packed).into_iter()
34/// }
35/// ```
36///
37/// # Safety
38///
39/// In order for the above relation to be guaranteed, the memory representation of
40/// `PackedExtensionField` element must be the same as a slice of the underlying `PackedField`
41/// element.
42pub trait PackedExtension<FS: Field>: PackedField<Scalar: ExtensionField<FS>> {
43	type PackedSubfield: PackedField<Scalar = FS>;
44
45	fn cast_bases(packed: &[Self]) -> &[Self::PackedSubfield];
46	fn cast_bases_mut(packed: &mut [Self]) -> &mut [Self::PackedSubfield];
47
48	fn cast_exts(packed: &[Self::PackedSubfield]) -> &[Self];
49	fn cast_exts_mut(packed: &mut [Self::PackedSubfield]) -> &mut [Self];
50
51	fn cast_base(self) -> Self::PackedSubfield;
52	fn cast_base_ref(&self) -> &Self::PackedSubfield;
53	fn cast_base_mut(&mut self) -> &mut Self::PackedSubfield;
54
55	fn cast_ext(base: Self::PackedSubfield) -> Self;
56	fn cast_ext_ref(base: &Self::PackedSubfield) -> &Self;
57	fn cast_ext_mut(base: &mut Self::PackedSubfield) -> &mut Self;
58
59	#[inline(always)]
60	fn cast_base_arr<const N: usize>(packed: [Self; N]) -> [Self::PackedSubfield; N] {
61		packed.map(Self::cast_base)
62	}
63
64	#[inline(always)]
65	fn cast_base_arr_ref<const N: usize>(packed: &[Self; N]) -> &[Self::PackedSubfield; N] {
66		Self::cast_bases(packed)
67			.try_into()
68			.expect("array has size N")
69	}
70
71	#[inline(always)]
72	fn cast_base_arr_mut<const N: usize>(packed: &mut [Self; N]) -> &mut [Self::PackedSubfield; N] {
73		Self::cast_bases_mut(packed)
74			.try_into()
75			.expect("array has size N")
76	}
77
78	#[inline(always)]
79	fn cast_ext_arr<const N: usize>(packed: [Self::PackedSubfield; N]) -> [Self; N] {
80		packed.map(Self::cast_ext)
81	}
82
83	#[inline(always)]
84	fn cast_ext_arr_ref<const N: usize>(packed: &[Self::PackedSubfield; N]) -> &[Self; N] {
85		Self::cast_exts(packed)
86			.try_into()
87			.expect("array has size N")
88	}
89
90	#[inline(always)]
91	fn cast_ext_arr_mut<const N: usize>(packed: &mut [Self::PackedSubfield; N]) -> &mut [Self; N] {
92		Self::cast_exts_mut(packed)
93			.try_into()
94			.expect("array has size N")
95	}
96}
97
98impl<PT, FS> PackedExtension<FS> for PT
99where
100	FS: Field,
101	PT: PackedField<Scalar: ExtensionField<FS>> + WithUnderlier<Underlier: PackScalar<FS>>,
102{
103	type PackedSubfield = <PT::Underlier as PackScalar<FS>>::Packed;
104
105	fn cast_bases(packed: &[Self]) -> &[Self::PackedSubfield] {
106		Self::PackedSubfield::from_underliers_ref(Self::to_underliers_ref(packed))
107	}
108
109	fn cast_bases_mut(packed: &mut [Self]) -> &mut [Self::PackedSubfield] {
110		Self::PackedSubfield::from_underliers_ref_mut(Self::to_underliers_ref_mut(packed))
111	}
112
113	fn cast_exts(base: &[Self::PackedSubfield]) -> &[Self] {
114		Self::from_underliers_ref(Self::PackedSubfield::to_underliers_ref(base))
115	}
116
117	fn cast_exts_mut(base: &mut [Self::PackedSubfield]) -> &mut [Self] {
118		Self::from_underliers_ref_mut(Self::PackedSubfield::to_underliers_ref_mut(base))
119	}
120
121	fn cast_base(self) -> Self::PackedSubfield {
122		Self::PackedSubfield::from_underlier(self.to_underlier())
123	}
124
125	fn cast_base_ref(&self) -> &Self::PackedSubfield {
126		Self::PackedSubfield::from_underlier_ref(self.to_underlier_ref())
127	}
128
129	fn cast_base_mut(&mut self) -> &mut Self::PackedSubfield {
130		Self::PackedSubfield::from_underlier_ref_mut(self.to_underlier_ref_mut())
131	}
132
133	fn cast_ext(base: Self::PackedSubfield) -> Self {
134		Self::from_underlier(base.to_underlier())
135	}
136
137	fn cast_ext_ref(base: &Self::PackedSubfield) -> &Self {
138		Self::from_underlier_ref(base.to_underlier_ref())
139	}
140
141	fn cast_ext_mut(base: &mut Self::PackedSubfield) -> &mut Self {
142		Self::from_underlier_ref_mut(base.to_underlier_ref_mut())
143	}
144}
145
146/// Convenient type alias that returns the packed field type for the scalar field `F` and packed
147/// extension `P`.
148pub type PackedSubfield<P, F> = <P as PackedExtension<F>>::PackedSubfield;
149
150/// Recast a packed field from one subfield of a packed extension to another.
151pub fn recast_packed<P, FSub1, FSub2>(elem: PackedSubfield<P, FSub1>) -> PackedSubfield<P, FSub2>
152where
153	P: PackedField + PackedExtension<FSub1> + PackedExtension<FSub2>,
154	P::Scalar: ExtensionField<FSub1> + ExtensionField<FSub2>,
155	FSub1: Field,
156	FSub2: Field,
157{
158	<P as PackedExtension<FSub2>>::cast_base(<P as PackedExtension<FSub1>>::cast_ext(elem))
159}
160
161/// Recast a slice of packed field elements from one subfield of a packed extension to another.
162pub fn recast_packed_slice<P, FSub1, FSub2>(
163	elems: &[PackedSubfield<P, FSub1>],
164) -> &[PackedSubfield<P, FSub2>]
165where
166	P: PackedField + PackedExtension<FSub1> + PackedExtension<FSub2>,
167	P::Scalar: ExtensionField<FSub1> + ExtensionField<FSub2>,
168	FSub1: Field,
169	FSub2: Field,
170{
171	<P as PackedExtension<FSub2>>::cast_bases(<P as PackedExtension<FSub1>>::cast_exts(elems))
172}
173
174/// Recast a mutable slice of packed field elements from one subfield of a packed extension to
175/// another.
176pub fn recast_packed_mut<P, FSub1, FSub2>(
177	elems: &mut [PackedSubfield<P, FSub1>],
178) -> &mut [PackedSubfield<P, FSub2>]
179where
180	P: PackedField + PackedExtension<FSub1> + PackedExtension<FSub2>,
181	P::Scalar: ExtensionField<FSub1> + ExtensionField<FSub2>,
182	FSub1: Field,
183	FSub2: Field,
184{
185	<P as PackedExtension<FSub2>>::cast_bases_mut(<P as PackedExtension<FSub1>>::cast_exts_mut(
186		elems,
187	))
188}
189
190/// This trait is a shorthand for the case `PackedExtension<P::Scalar, PackedSubfield = P>` which is
191/// a quite common case in our codebase.
192pub trait RepackedExtension<P: PackedField>:
193	PackedField<Scalar: ExtensionField<P::Scalar>> + PackedExtension<P::Scalar, PackedSubfield = P>
194{
195}
196
197impl<PT1, PT2> RepackedExtension<PT1> for PT2
198where
199	PT1: PackedField,
200	PT2: PackedExtension<PT1::Scalar, PackedSubfield = PT1, Scalar: ExtensionField<PT1::Scalar>>,
201{
202}
binius_field/packed_extension.rs

binius_field/
packed_extension.rs
