1use core::{
2 fmt::{Binary, LowerHex},
3 mem::size_of,
4};
5
6use bitvec::{mem::bits_of, order::Msb0, vec::BitVec, view::AsBits};
7use num_bigint::{BigInt, BigUint};
8use num_traits::{PrimInt, ToBytes};
9
10use crate::{
11 Error,
12 r#as::{Same, VarLen},
13 de::{BitReader, BitReaderExt, BitUnpackAs},
14 ser::{BitPackAs, BitWriter, BitWriterExt},
15};
16
17use super::NBits;
18
19impl<const BITS: usize> BitPackAs<BigUint> for NBits<BITS> {
20 type Args = ();
21
22 #[inline]
23 fn pack_as<W>(source: &BigUint, writer: &mut W, _: Self::Args) -> Result<(), W::Error>
24 where
25 W: BitWriter + ?Sized,
26 {
27 let used_bits = source.bits() as usize;
28 if BITS < used_bits {
29 return Err(Error::custom(format!(
30 "{source:#b} cannot be packed into {BITS} bits"
31 )));
32 }
33
34 writer.repeat_bit(BITS - used_bits, false)?;
35
36 let bytes = source.to_bytes_be();
37 let mut bits = bytes.as_bits::<Msb0>();
38 bits = &bits[bits.len() - used_bits..];
39 writer.write_bitslice(bits)?;
40 Ok(())
41 }
42}
43
44impl<'de, const BITS: usize> BitUnpackAs<'de, BigUint> for NBits<BITS> {
45 type Args = ();
46
47 #[inline]
48 fn unpack_as<R>(reader: &mut R, _: Self::Args) -> Result<BigUint, R::Error>
49 where
50 R: BitReader<'de> + ?Sized,
51 {
52 let total_bits = (BITS + 7) & !7;
53 let mut bits = BitVec::<u8, Msb0>::repeat(false, total_bits);
54 reader.read_bits_into(&mut bits[total_bits - BITS..])?;
55 Ok(BigUint::from_bytes_be(bits.as_raw_slice()))
56 }
57}
58
59impl<const BITS: usize> BitPackAs<BigInt> for NBits<BITS> {
60 type Args = ();
61
62 #[inline]
63 fn pack_as<W>(source: &BigInt, writer: &mut W, _: Self::Args) -> Result<(), W::Error>
64 where
65 W: BitWriter + ?Sized,
66 {
67 let used_bits = source.bits() as usize;
68 if BITS < used_bits {
69 return Err(Error::custom(format!(
70 "{source:#b} cannot be packed into {BITS} bits"
71 )));
72 }
73
74 writer.repeat_bit(BITS - used_bits, false)?;
75
76 let bytes = source.to_signed_bytes_be();
77 let mut bits = bytes.as_bits::<Msb0>();
78 bits = &bits[bits.len() - used_bits..];
79 writer.write_bitslice(bits)?;
80 Ok(())
81 }
82}
83
84impl<'de, const BITS: usize> BitUnpackAs<'de, BigInt> for NBits<BITS> {
85 type Args = ();
86
87 #[inline]
88 fn unpack_as<R>(reader: &mut R, _: Self::Args) -> Result<BigInt, R::Error>
89 where
90 R: BitReader<'de> + ?Sized,
91 {
92 let total_bits = (BITS + 7) & !7;
93 let mut bits = BitVec::<u8, Msb0>::repeat(false, total_bits);
94 reader.read_bits_into(&mut bits[total_bits - BITS..])?;
95 Ok(BigInt::from_signed_bytes_be(bits.as_raw_slice()))
96 }
97}
98
99#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
108pub struct VarInt<const BITS_FOR_BYTES_LEN: usize>;
109
110impl<const BITS_FOR_BYTES_LEN: usize> BitPackAs<BigUint> for VarInt<BITS_FOR_BYTES_LEN> {
111 type Args = ();
112
113 #[inline]
114 fn pack_as<W>(source: &BigUint, writer: &mut W, _: Self::Args) -> Result<(), W::Error>
115 where
116 W: BitWriter + ?Sized,
117 {
118 let bytes = if source != &BigUint::ZERO {
119 source.to_bytes_be()
120 } else {
121 Vec::new()
123 };
124 writer.pack_as::<_, VarLen<Vec<Same>, BITS_FOR_BYTES_LEN>>(bytes, ())?;
125 Ok(())
126 }
127}
128
129impl<'de, const BITS_FOR_BYTES_LEN: usize> BitUnpackAs<'de, BigUint>
130 for VarInt<BITS_FOR_BYTES_LEN>
131{
132 type Args = ();
133
134 #[inline]
135 fn unpack_as<R>(reader: &mut R, _: Self::Args) -> Result<BigUint, R::Error>
136 where
137 R: BitReader<'de> + ?Sized,
138 {
139 let mut bits = BitVec::<u8, Msb0>::from_vec(
140 reader.unpack_as::<_, VarLen<Vec<Same>, BITS_FOR_BYTES_LEN>>(())?,
141 );
142 let total_bits = (bits.len() + 7) & !7;
143 let shift = total_bits - bits.len();
144 bits.resize(total_bits, false);
145 bits.shift_right(shift);
146 Ok(BigUint::from_bytes_be(bits.as_raw_slice()))
147 }
148}
149
150impl<const BITS_FOR_BYTES_LEN: usize> BitPackAs<BigInt> for VarInt<BITS_FOR_BYTES_LEN> {
151 type Args = ();
152
153 #[inline]
154 fn pack_as<W>(source: &BigInt, writer: &mut W, _: Self::Args) -> Result<(), W::Error>
155 where
156 W: BitWriter + ?Sized,
157 {
158 writer.pack_as::<_, VarLen<Same, BITS_FOR_BYTES_LEN>>(source.to_signed_bytes_be(), ())?;
159 Ok(())
160 }
161}
162
163impl<'de, const BITS_FOR_BYTES_LEN: usize> BitUnpackAs<'de, BigInt> for VarInt<BITS_FOR_BYTES_LEN> {
164 type Args = ();
165
166 #[inline]
167 fn unpack_as<R>(reader: &mut R, _: Self::Args) -> Result<BigInt, R::Error>
168 where
169 R: BitReader<'de> + ?Sized,
170 {
171 let mut bits = BitVec::<u8, Msb0>::from_vec(
172 reader.unpack_as::<_, VarLen<Same, BITS_FOR_BYTES_LEN>>(())?,
173 );
174 let total_bits = (bits.len() + 7) & !7;
175 let shift = total_bits - bits.len();
176 bits.resize(total_bits, false);
177 bits.shift_right(shift);
178 Ok(BigInt::from_signed_bytes_be(bits.as_raw_slice()))
179 }
180}
181
182#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
189pub struct VarNBits;
190
191impl<T> BitPackAs<T> for VarNBits
192where
193 T: PrimInt + Binary + ToBytes,
194{
195 type Args = u32;
197
198 #[inline]
199 fn pack_as<W>(source: &T, writer: &mut W, num_bits: Self::Args) -> Result<(), W::Error>
200 where
201 W: BitWriter + ?Sized,
202 {
203 let size_bits: u32 = bits_of::<T>() as u32;
204 let leading_zeroes = source.leading_zeros();
205 let used_bits = size_bits - leading_zeroes;
206 if num_bits < used_bits {
207 return Err(Error::custom(format!(
208 "{source:0b} cannot be packed into {num_bits} bits",
209 )));
210 }
211 let arr = source.to_be_bytes();
212 let bits = arr.as_bits();
213 writer.write_bitslice(&bits[bits.len() - num_bits as usize..])?;
214 Ok(())
215 }
216}
217
218impl<'de, T> BitUnpackAs<'de, T> for VarNBits
219where
220 T: PrimInt,
221{
222 type Args = u32;
224
225 #[inline]
226 fn unpack_as<R>(reader: &mut R, num_bits: Self::Args) -> Result<T, R::Error>
227 where
228 R: BitReader<'de> + ?Sized,
229 {
230 let size_bits: u32 = bits_of::<T>() as u32;
231 if num_bits > size_bits {
232 return Err(Error::custom("excessive bits for the type"));
233 }
234 let mut v: T = T::zero();
235 for bit in reader.unpack_iter::<bool>(()).take(num_bits as usize) {
236 v = v << 1;
237 v = v | if bit? { T::one() } else { T::zero() };
238 }
239 Ok(v)
240 }
241}
242
243#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
250pub struct VarNBytes;
251
252impl<T> BitPackAs<T> for VarNBytes
253where
254 T: PrimInt + LowerHex + ToBytes,
255{
256 type Args = u32;
258
259 #[inline]
260 fn pack_as<W>(source: &T, writer: &mut W, num_bytes: Self::Args) -> Result<(), W::Error>
261 where
262 W: BitWriter + ?Sized,
263 {
264 let size_bytes: u32 = size_of::<T>() as u32;
265 let leading_zeroes = source.leading_zeros();
266 let used_bytes = size_bytes - leading_zeroes / 8;
267 if num_bytes < used_bytes {
268 return Err(Error::custom(format!(
269 "{source:0x} cannot be packed into {num_bytes} bytes",
270 )));
271 }
272 let arr = source.to_be_bytes();
273 let bytes = arr.as_ref();
274 writer.write_bitslice((&bytes[bytes.len() - num_bytes as usize..]).as_bits())?;
275 Ok(())
276 }
277}
278
279impl<'de, T> BitUnpackAs<'de, T> for VarNBytes
280where
281 T: PrimInt,
282{
283 type Args = u32;
285
286 #[inline]
287 fn unpack_as<R>(reader: &mut R, num_bytes: Self::Args) -> Result<T, R::Error>
288 where
289 R: BitReader<'de> + ?Sized,
290 {
291 let size_bytes: u32 = size_of::<T>() as u32;
292 if num_bytes > size_bytes {
293 return Err(Error::custom("excessive bits for type"));
294 }
295 let mut v: T = T::zero();
296 for byte in reader.unpack_iter::<u8>(()).take(num_bytes as usize) {
297 v = v << 8;
298 v = v | T::from(byte?).unwrap();
299 }
300 Ok(v)
301 }
302}