binius_field/underlier/
small_uint.rs

1// Copyright 2024-2025 Irreducible Inc.
2
3use std::{
4	fmt::{Debug, Display, LowerHex},
5	hash::{Hash, Hasher},
6	ops::{Not, Shl, Shr},
7};
8
9use binius_utils::{
10	SerializationError, SerializationMode, SerializeBytes,
11	bytes::{Buf, BufMut},
12	checked_arithmetics::checked_log_2,
13	serialization::DeserializeBytes,
14};
15use bytemuck::{NoUninit, Zeroable};
16use derive_more::{BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign};
17use rand::Rng;
18use subtle::{ConditionallySelectable, ConstantTimeEq};
19
20use super::{Random, UnderlierType, underlier_with_bit_ops::UnderlierWithBitOps};
21
22/// Unsigned type with a size strictly less than 8 bits.
23#[derive(
24	Default,
25	Zeroable,
26	Clone,
27	Copy,
28	PartialEq,
29	Eq,
30	PartialOrd,
31	Ord,
32	BitAnd,
33	BitAndAssign,
34	BitOr,
35	BitOrAssign,
36	BitXor,
37	BitXorAssign,
38)]
39#[repr(transparent)]
40pub struct SmallU<const N: usize>(u8);
41
42impl<const N: usize> SmallU<N> {
43	const _CHECK_SIZE: () = {
44		assert!(N < 8);
45	};
46
47	#[inline(always)]
48	pub const fn new(val: u8) -> Self {
49		Self(val & Self::ONES.0)
50	}
51
52	#[inline(always)]
53	pub const fn new_unchecked(val: u8) -> Self {
54		Self(val)
55	}
56
57	#[inline(always)]
58	pub const fn val(&self) -> u8 {
59		self.0
60	}
61
62	pub fn checked_add(self, rhs: Self) -> Option<Self> {
63		self.val()
64			.checked_add(rhs.val())
65			.and_then(|value| (value < Self::ONES.0).then_some(Self(value)))
66	}
67
68	pub fn checked_sub(self, rhs: Self) -> Option<Self> {
69		let a = self.val();
70		let b = rhs.val();
71		(b > a).then_some(Self(b - a))
72	}
73}
74
75impl<const N: usize> Debug for SmallU<N> {
76	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77		Debug::fmt(&self.val(), f)
78	}
79}
80
81impl<const N: usize> Display for SmallU<N> {
82	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83		Display::fmt(&self.val(), f)
84	}
85}
86
87impl<const N: usize> LowerHex for SmallU<N> {
88	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89		LowerHex::fmt(&self.0, f)
90	}
91}
92impl<const N: usize> Hash for SmallU<N> {
93	#[inline]
94	fn hash<H: Hasher>(&self, state: &mut H) {
95		self.val().hash(state);
96	}
97}
98
99impl<const N: usize> ConstantTimeEq for SmallU<N> {
100	fn ct_eq(&self, other: &Self) -> subtle::Choice {
101		self.val().ct_eq(&other.val())
102	}
103}
104
105impl<const N: usize> ConditionallySelectable for SmallU<N> {
106	fn conditional_select(a: &Self, b: &Self, choice: subtle::Choice) -> Self {
107		Self(u8::conditional_select(&a.0, &b.0, choice))
108	}
109}
110
111impl<const N: usize> Random for SmallU<N> {
112	fn random(mut rng: impl Rng) -> Self {
113		Self(rng.random_range(0..1u8 << N))
114	}
115}
116
117impl<const N: usize> Shr<usize> for SmallU<N> {
118	type Output = Self;
119
120	#[inline(always)]
121	fn shr(self, rhs: usize) -> Self::Output {
122		Self(self.val() >> rhs)
123	}
124}
125
126impl<const N: usize> Shl<usize> for SmallU<N> {
127	type Output = Self;
128
129	#[inline(always)]
130	fn shl(self, rhs: usize) -> Self::Output {
131		Self(self.val() << rhs) & Self::ONES
132	}
133}
134
135impl<const N: usize> Not for SmallU<N> {
136	type Output = Self;
137
138	fn not(self) -> Self::Output {
139		self ^ Self::ONES
140	}
141}
142
143unsafe impl<const N: usize> NoUninit for SmallU<N> {}
144
145impl<const N: usize> UnderlierType for SmallU<N> {
146	const LOG_BITS: usize = checked_log_2(N);
147}
148
149impl<const N: usize> UnderlierWithBitOps for SmallU<N> {
150	const ZERO: Self = Self(0);
151	const ONE: Self = Self(1);
152	const ONES: Self = Self((1u8 << N) - 1);
153
154	fn fill_with_bit(val: u8) -> Self {
155		Self(u8::fill_with_bit(val)) & Self::ONES
156	}
157
158	fn shl_128b_lanes(self, rhs: usize) -> Self {
159		self << rhs
160	}
161
162	fn shr_128b_lanes(self, rhs: usize) -> Self {
163		self >> rhs
164	}
165}
166
167impl<const N: usize> From<SmallU<N>> for u8 {
168	#[inline(always)]
169	fn from(value: SmallU<N>) -> Self {
170		value.val()
171	}
172}
173
174impl<const N: usize> From<SmallU<N>> for u16 {
175	#[inline(always)]
176	fn from(value: SmallU<N>) -> Self {
177		u8::from(value) as _
178	}
179}
180
181impl<const N: usize> From<SmallU<N>> for u32 {
182	#[inline(always)]
183	fn from(value: SmallU<N>) -> Self {
184		u8::from(value) as _
185	}
186}
187
188impl<const N: usize> From<SmallU<N>> for u64 {
189	#[inline(always)]
190	fn from(value: SmallU<N>) -> Self {
191		u8::from(value) as _
192	}
193}
194
195impl<const N: usize> From<SmallU<N>> for usize {
196	#[inline(always)]
197	fn from(value: SmallU<N>) -> Self {
198		u8::from(value) as _
199	}
200}
201
202impl<const N: usize> From<SmallU<N>> for u128 {
203	#[inline(always)]
204	fn from(value: SmallU<N>) -> Self {
205		u8::from(value) as _
206	}
207}
208
209impl From<SmallU<1>> for SmallU<2> {
210	#[inline(always)]
211	fn from(value: SmallU<1>) -> Self {
212		Self(value.val())
213	}
214}
215
216impl From<SmallU<1>> for SmallU<4> {
217	#[inline(always)]
218	fn from(value: SmallU<1>) -> Self {
219		Self(value.val())
220	}
221}
222
223impl From<SmallU<2>> for SmallU<4> {
224	#[inline(always)]
225	fn from(value: SmallU<2>) -> Self {
226		Self(value.val())
227	}
228}
229
230pub type U1 = SmallU<1>;
231pub type U2 = SmallU<2>;
232pub type U4 = SmallU<4>;
233
234impl From<bool> for U1 {
235	fn from(value: bool) -> Self {
236		Self::new_unchecked(value as u8)
237	}
238}
239
240impl From<U1> for bool {
241	fn from(value: U1) -> Self {
242		value == U1::ONE
243	}
244}
245
246impl<const N: usize> SerializeBytes for SmallU<N> {
247	fn serialize(
248		&self,
249		write_buf: impl BufMut,
250		mode: SerializationMode,
251	) -> Result<(), SerializationError> {
252		self.val().serialize(write_buf, mode)
253	}
254}
255
256impl<const N: usize> DeserializeBytes for SmallU<N> {
257	fn deserialize(read_buf: impl Buf, mode: SerializationMode) -> Result<Self, SerializationError>
258	where
259		Self: Sized,
260	{
261		Ok(Self::new(DeserializeBytes::deserialize(read_buf, mode)?))
262	}
263}
264
265#[cfg(test)]
266impl<const N: usize> proptest::arbitrary::Arbitrary for SmallU<N> {
267	type Parameters = ();
268	type Strategy = proptest::strategy::BoxedStrategy<Self>;
269
270	fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
271		use proptest::strategy::Strategy;
272
273		(0u8..(1u8 << N)).prop_map(Self::new_unchecked).boxed()
274	}
275}