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