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

binius_field/underlier/
scaled.rs
