binius_field/
packed_extension.rs

1// Copyright 2023-2025 Irreducible Inc.
2
3use crate::{
4	BinaryField, ExtensionField, Field, PackedField, arch::PackedPrimitiveType,
5	underlier::WithUnderlier,
6};
7
8/// Trait represents a relationship between a packed struct of field elements and a packed struct
9/// of elements from an extension field.
10///
11/// This trait guarantees that one packed type has the same
12/// memory representation as the other, differing only in the scalar type and preserving the order
13/// of smaller elements.
14///
15/// This trait relation guarantees that the following iterators yield the same sequence of scalar
16/// elements:
17///
18/// ```
19/// use binius_field::{ExtensionField, PackedExtension, PackedField, Field};
20///
21/// fn ext_then_bases<'a, F, PE>(packed: &'a PE) -> impl Iterator<Item=F> + 'a
22///     where
23///         PE: PackedField<Scalar: ExtensionField<F>>,
24///         F: Field,
25/// {
26///     packed.iter().flat_map(|ext| ext.into_iter_bases())
27/// }
28///
29/// fn cast_then_iter<'a, F, PE>(packed: &'a PE) -> impl Iterator<Item=F> + 'a
30///     where
31///         PE: PackedExtension<F>,
32///         F: Field,
33/// {
34///     PE::cast_base_ref(packed).into_iter()
35/// }
36/// ```
37///
38/// # Safety
39///
40/// In order for the above relation to be guaranteed, the memory representation of
41/// `PackedExtensionField` element must be the same as a slice of the underlying `PackedField`
42/// element.
43pub trait PackedExtension<FS: Field>: PackedField<Scalar: ExtensionField<FS>> {
44	type PackedSubfield: PackedField<Scalar = FS>;
45
46	fn cast_bases(packed: &[Self]) -> &[Self::PackedSubfield];
47	fn cast_bases_mut(packed: &mut [Self]) -> &mut [Self::PackedSubfield];
48
49	fn cast_exts(packed: &[Self::PackedSubfield]) -> &[Self];
50	fn cast_exts_mut(packed: &mut [Self::PackedSubfield]) -> &mut [Self];
51
52	fn cast_base(self) -> Self::PackedSubfield;
53	fn cast_base_ref(&self) -> &Self::PackedSubfield;
54	fn cast_base_mut(&mut self) -> &mut Self::PackedSubfield;
55
56	fn cast_ext(base: Self::PackedSubfield) -> Self;
57	fn cast_ext_ref(base: &Self::PackedSubfield) -> &Self;
58	fn cast_ext_mut(base: &mut Self::PackedSubfield) -> &mut Self;
59
60	#[inline(always)]
61	fn cast_base_arr<const N: usize>(packed: [Self; N]) -> [Self::PackedSubfield; N] {
62		packed.map(Self::cast_base)
63	}
64
65	#[inline(always)]
66	fn cast_base_arr_ref<const N: usize>(packed: &[Self; N]) -> &[Self::PackedSubfield; N] {
67		Self::cast_bases(packed)
68			.try_into()
69			.expect("array has size N")
70	}
71
72	#[inline(always)]
73	fn cast_base_arr_mut<const N: usize>(packed: &mut [Self; N]) -> &mut [Self::PackedSubfield; N] {
74		Self::cast_bases_mut(packed)
75			.try_into()
76			.expect("array has size N")
77	}
78
79	#[inline(always)]
80	fn cast_ext_arr<const N: usize>(packed: [Self::PackedSubfield; N]) -> [Self; N] {
81		packed.map(Self::cast_ext)
82	}
83
84	#[inline(always)]
85	fn cast_ext_arr_ref<const N: usize>(packed: &[Self::PackedSubfield; N]) -> &[Self; N] {
86		Self::cast_exts(packed)
87			.try_into()
88			.expect("array has size N")
89	}
90
91	#[inline(always)]
92	fn cast_ext_arr_mut<const N: usize>(packed: &mut [Self::PackedSubfield; N]) -> &mut [Self; N] {
93		Self::cast_exts_mut(packed)
94			.try_into()
95			.expect("array has size N")
96	}
97}
98
99impl<PT, FS> PackedExtension<FS> for PT
100where
101	FS: BinaryField,
102	PT: PackedField<Scalar: ExtensionField<FS>> + WithUnderlier,
103	PackedPrimitiveType<PT::Underlier, FS>: PackedField<Scalar = FS>,
104{
105	type PackedSubfield = PackedPrimitiveType<PT::Underlier, FS>;
106
107	fn cast_bases(packed: &[Self]) -> &[Self::PackedSubfield] {
108		Self::PackedSubfield::from_underliers_ref(Self::to_underliers_ref(packed))
109	}
110
111	fn cast_bases_mut(packed: &mut [Self]) -> &mut [Self::PackedSubfield] {
112		Self::PackedSubfield::from_underliers_ref_mut(Self::to_underliers_ref_mut(packed))
113	}
114
115	fn cast_exts(base: &[Self::PackedSubfield]) -> &[Self] {
116		Self::from_underliers_ref(Self::PackedSubfield::to_underliers_ref(base))
117	}
118
119	fn cast_exts_mut(base: &mut [Self::PackedSubfield]) -> &mut [Self] {
120		Self::from_underliers_ref_mut(Self::PackedSubfield::to_underliers_ref_mut(base))
121	}
122
123	fn cast_base(self) -> Self::PackedSubfield {
124		Self::PackedSubfield::from_underlier(self.to_underlier())
125	}
126
127	fn cast_base_ref(&self) -> &Self::PackedSubfield {
128		Self::PackedSubfield::from_underlier_ref(self.to_underlier_ref())
129	}
130
131	fn cast_base_mut(&mut self) -> &mut Self::PackedSubfield {
132		Self::PackedSubfield::from_underlier_ref_mut(self.to_underlier_ref_mut())
133	}
134
135	fn cast_ext(base: Self::PackedSubfield) -> Self {
136		Self::from_underlier(base.to_underlier())
137	}
138
139	fn cast_ext_ref(base: &Self::PackedSubfield) -> &Self {
140		Self::from_underlier_ref(base.to_underlier_ref())
141	}
142
143	fn cast_ext_mut(base: &mut Self::PackedSubfield) -> &mut Self {
144		Self::from_underlier_ref_mut(base.to_underlier_ref_mut())
145	}
146}
147
148/// Convenient type alias that returns the packed field type for the scalar field `F` and packed
149/// extension `P`.
150pub type PackedSubfield<P, F> = <P as PackedExtension<F>>::PackedSubfield;