binius_field/underlier/
scaled.rs

1// Copyright 2024-2025 Irreducible Inc.
2// Copyright 2026 The Binius Developers
3
4use std::{
5	array,
6	fmt::{self, LowerHex},
7	mem,
8	ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Not},
9};
10
11use binius_utils::{
12	DeserializeBytes, FixedSizeSerializeBytes, SerializationError, SerializeBytes,
13	bytes::{Buf, BufMut},
14	checked_arithmetics::checked_log_2,
15};
16use bytemuck::{Pod, Zeroable};
17use rand::{
18	Rng,
19	distr::{Distribution, StandardUniform},
20};
21
22use super::{U1, Underlier};
23use crate::{
24	Random,
25	divisible::{Divisible, mapget},
26};
27
28/// A type that represents N elements of the same underlier type.
29/// Used as an underlier for 256-bit and 512-bit packed fields in the portable implementation.
30#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
31#[repr(transparent)]
32pub struct ScaledUnderlier<U, const N: usize>(pub [U; N]);
33
34impl<U: Default, const N: usize> Default for ScaledUnderlier<U, N> {
35	fn default() -> Self {
36		Self(array::from_fn(|_| U::default()))
37	}
38}
39
40impl<U: Random, const N: usize> Distribution<ScaledUnderlier<U, N>> for StandardUniform {
41	fn sample<R: Rng + ?Sized>(&self, mut rng: &mut R) -> ScaledUnderlier<U, N> {
42		ScaledUnderlier(array::from_fn(|_| U::random(&mut rng)))
43	}
44}
45
46impl<U, const N: usize> From<ScaledUnderlier<U, N>> for [U; N] {
47	fn from(val: ScaledUnderlier<U, N>) -> Self {
48		val.0
49	}
50}
51
52impl<T, U: From<T>, const N: usize> From<[T; N]> for ScaledUnderlier<U, N> {
53	fn from(value: [T; N]) -> Self {
54		Self(value.map(U::from))
55	}
56}
57
58impl<T: Copy, U: From<[T; 2]>> From<[T; 4]> for ScaledUnderlier<U, 2> {
59	fn from(value: [T; 4]) -> Self {
60		Self([[value[0], value[1]], [value[2], value[3]]].map(Into::into))
61	}
62}
63
64unsafe impl<U: Zeroable, const N: usize> Zeroable for ScaledUnderlier<U, N> {}
65unsafe impl<U: Pod, const N: usize> Pod for ScaledUnderlier<U, N> {}
66
67impl<U: BitAnd<Output = U> + Copy, const N: usize> BitAnd for ScaledUnderlier<U, N> {
68	type Output = Self;
69
70	fn bitand(self, rhs: Self) -> Self::Output {
71		Self(array::from_fn(|i| self.0[i] & rhs.0[i]))
72	}
73}
74
75impl<U: BitAndAssign + Copy, const N: usize> BitAndAssign for ScaledUnderlier<U, N> {
76	fn bitand_assign(&mut self, rhs: Self) {
77		for i in 0..N {
78			self.0[i] &= rhs.0[i];
79		}
80	}
81}
82
83impl<U: BitOr<Output = U> + Copy, const N: usize> BitOr for ScaledUnderlier<U, N> {
84	type Output = Self;
85
86	fn bitor(self, rhs: Self) -> Self::Output {
87		Self(array::from_fn(|i| self.0[i] | rhs.0[i]))
88	}
89}
90
91impl<U: BitOrAssign + Copy, const N: usize> BitOrAssign for ScaledUnderlier<U, N> {
92	fn bitor_assign(&mut self, rhs: Self) {
93		for i in 0..N {
94			self.0[i] |= rhs.0[i];
95		}
96	}
97}
98
99impl<U: BitXor<Output = U> + Copy, const N: usize> BitXor for ScaledUnderlier<U, N> {
100	type Output = Self;
101
102	fn bitxor(self, rhs: Self) -> Self::Output {
103		Self(array::from_fn(|i| self.0[i] ^ rhs.0[i]))
104	}
105}
106
107impl<U: BitXorAssign + Copy, const N: usize> BitXorAssign for ScaledUnderlier<U, N> {
108	fn bitxor_assign(&mut self, rhs: Self) {
109		for i in 0..N {
110			self.0[i] ^= rhs.0[i];
111		}
112	}
113}
114
115impl<U: Not<Output = U>, const N: usize> Not for ScaledUnderlier<U, N> {
116	type Output = Self;
117
118	fn not(self) -> Self::Output {
119		Self(self.0.map(U::not))
120	}
121}
122
123impl<U: Underlier + Pod, const N: usize> Underlier for ScaledUnderlier<U, N> {
124	const LOG_BITS: usize = U::LOG_BITS + checked_log_2(N);
125
126	const ZERO: Self = Self([U::ZERO; N]);
127	const ONE: Self = {
128		let mut arr = [U::ZERO; N];
129		arr[0] = U::ONE;
130		Self(arr)
131	};
132	const ONES: Self = Self([U::ONES; N]);
133
134	fn interleave(self, other: Self, log_block_len: usize) -> (Self, Self) {
135		if log_block_len < U::LOG_BITS {
136			// Case 1: Delegate to element-wise interleave
137			let pairs: [(U, U); N] =
138				array::from_fn(|i| self.0[i].interleave(other.0[i], log_block_len));
139			(Self(array::from_fn(|i| pairs[i].0)), Self(array::from_fn(|i| pairs[i].1)))
140		} else {
141			// Case 2: Interleave at element level by swapping array elements
142			// Each super-block of 2*block_len elements gets transposed as a 2x2 matrix of blocks
143			let block_len = 1 << (log_block_len - U::LOG_BITS);
144
145			let mut a = self.0;
146			let mut b = other.0;
147			for super_block in 0..(N / (2 * block_len)) {
148				let base = super_block * 2 * block_len;
149				for offset in 0..block_len {
150					mem::swap(&mut a[base + block_len + offset], &mut b[base + offset]);
151				}
152			}
153
154			(Self(a), Self(b))
155		}
156	}
157}
158
159impl<U, const N: usize> From<u8> for ScaledUnderlier<U, N>
160where
161	U: From<u8>,
162{
163	fn from(val: u8) -> Self {
164		Self(array::from_fn(|_| U::from(val)))
165	}
166}
167
168/// Zero-extends an `M128` into the least-significant limb, leaving the rest zero.
169///
170/// This is the embedding of a base-field underlier into a `ScaledUnderlier<M128, _>`-backed
171/// extension field, as used by `impl_field_extension!`'s `from_bases_sparse` (`GhashSq256b` off
172/// the AVX2 path).
173impl<const N: usize> From<crate::arch::M128> for ScaledUnderlier<crate::arch::M128, N> {
174	fn from(val: crate::arch::M128) -> Self {
175		let mut limbs = [<crate::arch::M128 as Underlier>::ZERO; N];
176		limbs[0] = val;
177		Self(limbs)
178	}
179}
180
181/// Zero-extends a single bit into the least-significant `M128` limb, leaving the rest zero.
182impl<const N: usize> From<U1> for ScaledUnderlier<crate::arch::M128, N> {
183	fn from(val: U1) -> Self {
184		Self::from(crate::arch::M128::from(val))
185	}
186}
187
188impl<U: Underlier + LowerHex, const N: usize> LowerHex for ScaledUnderlier<U, N> {
189	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
190		// Most-significant limb first. Print from the highest non-zero limb so there are no
191		// spurious leading zeros, then zero-pad each remaining limb to its full bit width.
192		let width = U::BITS / 4;
193		let top = self
194			.0
195			.iter()
196			.rposition(|limb| *limb != U::ZERO)
197			.unwrap_or(0);
198		write!(f, "{:x}", self.0[top])?;
199		for limb in self.0[..top].iter().rev() {
200			write!(f, "{limb:0width$x}")?;
201		}
202		Ok(())
203	}
204}
205
206impl<U, T, const N: usize> Divisible<T> for ScaledUnderlier<U, N>
207where
208	U: Divisible<T> + Pod + Send + Sync,
209	T: Send + 'static,
210{
211	const LOG_N: usize = U::LOG_N + checked_log_2(N);
212
213	#[inline]
214	fn value_iter(value: Self) -> impl ExactSizeIterator<Item = T> + Send + Clone {
215		mapget::value_iter(value)
216	}
217
218	#[inline]
219	fn ref_iter(value: &Self) -> impl ExactSizeIterator<Item = T> + Send + Clone + '_ {
220		mapget::value_iter(*value)
221	}
222
223	#[inline]
224	fn slice_iter(slice: &[Self]) -> impl ExactSizeIterator<Item = T> + Send + Clone + '_ {
225		mapget::slice_iter(slice)
226	}
227
228	#[inline]
229	unsafe fn get_unchecked(&self, index: usize) -> T {
230		let u_index = index >> U::LOG_N;
231		let sub_index = index & (U::N - 1);
232		// Safety: `index < Self::N` by the caller's contract, so `sub_index < <U as
233		// Divisible<T>>::N` and `u_index < N`.
234		unsafe { Divisible::<T>::get_unchecked(self.0.get_unchecked(u_index), sub_index) }
235	}
236
237	#[inline]
238	unsafe fn set_unchecked(&mut self, index: usize, val: T) {
239		let u_index = index >> U::LOG_N;
240		let sub_index = index & (U::N - 1);
241		// Safety: `index < Self::N` by the caller's contract, so `sub_index < <U as
242		// Divisible<T>>::N` and `u_index < N`.
243		unsafe { Divisible::<T>::set_unchecked(self.0.get_unchecked_mut(u_index), sub_index, val) };
244	}
245
246	#[inline]
247	fn broadcast(val: T) -> Self {
248		Self([Divisible::<T>::broadcast(val); N])
249	}
250
251	#[inline]
252	fn from_iter(mut iter: impl Iterator<Item = T>) -> Self {
253		Self(array::from_fn(|_| Divisible::<T>::from_iter(&mut iter)))
254	}
255}
256
257impl<U: SerializeBytes, const N: usize> SerializeBytes for ScaledUnderlier<U, N> {
258	fn serialize(&self, write_buf: impl BufMut) -> Result<(), SerializationError> {
259		self.0.serialize(write_buf)
260	}
261}
262
263impl<U: DeserializeBytes, const N: usize> DeserializeBytes for ScaledUnderlier<U, N> {
264	fn deserialize(read_buf: impl Buf) -> Result<Self, SerializationError> {
265		<[U; N]>::deserialize(read_buf).map(Self)
266	}
267}
268
269impl<U: FixedSizeSerializeBytes, const N: usize> FixedSizeSerializeBytes for ScaledUnderlier<U, N> {
270	const BYTE_SIZE: usize = U::BYTE_SIZE * N;
271}
272
273#[cfg(test)]
274mod tests {
275	use super::*;
276
277	#[test]
278	fn test_interleave_within_element() {
279		// Test case 1: log_block_len < U::LOG_BITS
280		// ScaledUnderlier<u8, 4> has LOG_BITS = 5 (32 bits total)
281		// u8 has LOG_BITS = 3
282		let a = ScaledUnderlier::<u8, 4>([0b01010101, 0b11110000, 0b00001111, 0b10101010]);
283		let b = ScaledUnderlier::<u8, 4>([0b10101010, 0b00001111, 0b11110000, 0b01010101]);
284
285		// At log_block_len = 0 (1-bit blocks), should delegate to u8::interleave
286		let (c, d) = a.interleave(b, 0);
287
288		// Verify element-wise interleave occurred
289		for i in 0..4 {
290			let (expected_c, expected_d) = a.0[i].interleave(b.0[i], 0);
291			assert_eq!(c.0[i], expected_c);
292			assert_eq!(d.0[i], expected_d);
293		}
294	}
295
296	#[test]
297	fn test_interleave_across_elements() {
298		// Test case 2: log_block_len >= U::LOG_BITS
299		let a = ScaledUnderlier::<u8, 4>([0, 1, 2, 3]);
300		let b = ScaledUnderlier::<u8, 4>([4, 5, 6, 7]);
301
302		// At log_block_len = 3 (8-bit blocks = 1 element), swap individual elements
303		let (c, d) = a.interleave(b, 3);
304		assert_eq!(c.0, [0, 4, 2, 6]);
305		assert_eq!(d.0, [1, 5, 3, 7]);
306
307		// At log_block_len = 4 (16-bit blocks = 2 elements), swap pairs
308		let (c, d) = a.interleave(b, 4);
309		assert_eq!(c.0, [0, 1, 4, 5]);
310		assert_eq!(d.0, [2, 3, 6, 7]);
311	}
312}
binius_field/underlier/scaled.rs

binius_field/underlier/
scaled.rs
