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 de::{BitReader, BitReaderExt, args::r#as::BitUnpackAsWithArgs, r#as::BitUnpackAs},
13 ser::{BitWriter, BitWriterExt, args::r#as::BitPackAsWithArgs, r#as::BitPackAs},
14};
15
16use super::{NBits, VarBytes};
17
18impl<const BITS: usize> BitPackAs<BigUint> for NBits<BITS> {
19 #[inline]
20 fn pack_as<W>(source: &BigUint, mut writer: W) -> Result<(), W::Error>
21 where
22 W: BitWriter,
23 {
24 let used_bits = source.bits() as usize;
25 if BITS < used_bits {
26 return Err(Error::custom(format!(
27 "{source:#b} cannot be packed into {BITS} bits"
28 )));
29 }
30
31 writer.repeat_bit(BITS - used_bits, false)?;
32
33 let bytes = source.to_bytes_be();
34 let mut bits = bytes.as_bits::<Msb0>();
35 bits = &bits[bits.len() - used_bits..];
36 writer.pack(bits)?;
37 Ok(())
38 }
39}
40
41impl<'de, const BITS: usize> BitUnpackAs<'de, BigUint> for NBits<BITS> {
42 #[inline]
43 fn unpack_as<R>(mut reader: R) -> Result<BigUint, R::Error>
44 where
45 R: BitReader<'de>,
46 {
47 let mut bits: BitVec<u8, Msb0> = reader.unpack_with(BITS)?;
48 let total_bits = (BITS + 7) & !7;
49 bits.resize(total_bits, false);
50 bits.shift_right(total_bits - BITS);
51 Ok(BigUint::from_bytes_be(bits.as_raw_slice()))
52 }
53}
54
55impl<const BITS: usize> BitPackAs<BigInt> for NBits<BITS> {
56 #[inline]
57 fn pack_as<W>(source: &BigInt, mut writer: W) -> Result<(), W::Error>
58 where
59 W: BitWriter,
60 {
61 let used_bits = source.bits() as usize;
62 if BITS < used_bits {
63 return Err(Error::custom(format!(
64 "{source:#b} cannot be packed into {BITS} bits"
65 )));
66 }
67
68 writer.repeat_bit(BITS - used_bits, false)?;
69
70 let bytes = source.to_signed_bytes_be();
71 let mut bits = bytes.as_bits::<Msb0>();
72 bits = &bits[bits.len() - used_bits..];
73 writer.pack(bits)?;
74 Ok(())
75 }
76}
77
78impl<'de, const BITS: usize> BitUnpackAs<'de, BigInt> for NBits<BITS> {
79 #[inline]
80 fn unpack_as<R>(mut reader: R) -> Result<BigInt, R::Error>
81 where
82 R: BitReader<'de>,
83 {
84 let mut bits: BitVec<u8, Msb0> = reader.unpack_with(BITS)?;
85 let total_bits = (BITS + 7) & !7;
86 bits.resize(total_bits, false);
87 bits.shift_right(total_bits - BITS);
88 Ok(BigInt::from_signed_bytes_be(bits.as_raw_slice()))
89 }
90}
91
92#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
101pub struct VarInt<const BITS_FOR_BYTES_LEN: usize>;
102
103impl<const BITS_FOR_BYTES_LEN: usize> BitPackAs<BigUint> for VarInt<BITS_FOR_BYTES_LEN> {
104 #[inline]
105 fn pack_as<W>(source: &BigUint, mut writer: W) -> Result<(), W::Error>
106 where
107 W: BitWriter,
108 {
109 let bytes = if source != &BigUint::ZERO {
110 source.to_bytes_be()
111 } else {
112 Vec::new()
114 };
115 writer.pack_as::<_, VarBytes<BITS_FOR_BYTES_LEN>>(bytes)?;
116 Ok(())
117 }
118}
119
120impl<'de, const BITS_FOR_BYTES_LEN: usize> BitUnpackAs<'de, BigUint>
121 for VarInt<BITS_FOR_BYTES_LEN>
122{
123 #[inline]
124 fn unpack_as<R>(mut reader: R) -> Result<BigUint, R::Error>
125 where
126 R: BitReader<'de>,
127 {
128 let mut bits = BitVec::<u8, Msb0>::from_vec(
129 reader.unpack_as::<Vec<u8>, VarBytes<BITS_FOR_BYTES_LEN>>()?,
130 );
131 let total_bits = (bits.len() + 7) & !7;
132 let shift = total_bits - bits.len();
133 bits.resize(total_bits, false);
134 bits.shift_right(shift);
135 Ok(BigUint::from_bytes_be(bits.as_raw_slice()))
136 }
137}
138
139impl<const BITS_FOR_BYTES_LEN: usize> BitPackAs<BigInt> for VarInt<BITS_FOR_BYTES_LEN> {
140 #[inline]
141 fn pack_as<W>(source: &BigInt, mut writer: W) -> Result<(), W::Error>
142 where
143 W: BitWriter,
144 {
145 writer.pack_as::<_, VarBytes<BITS_FOR_BYTES_LEN>>(source.to_signed_bytes_be())?;
146 Ok(())
147 }
148}
149
150impl<'de, const BITS_FOR_BYTES_LEN: usize> BitUnpackAs<'de, BigInt> for VarInt<BITS_FOR_BYTES_LEN> {
151 #[inline]
152 fn unpack_as<R>(mut reader: R) -> Result<BigInt, R::Error>
153 where
154 R: BitReader<'de>,
155 {
156 let mut bits = BitVec::<u8, Msb0>::from_vec(
157 reader.unpack_as::<Vec<u8>, VarBytes<BITS_FOR_BYTES_LEN>>()?,
158 );
159 let total_bits = (bits.len() + 7) & !7;
160 let shift = total_bits - bits.len();
161 bits.resize(total_bits, false);
162 bits.shift_right(shift);
163 Ok(BigInt::from_signed_bytes_be(bits.as_raw_slice()))
164 }
165}
166
167#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
174pub struct VarNBits;
175
176impl<T> BitPackAsWithArgs<T> for VarNBits
177where
178 T: PrimInt + Binary + ToBytes,
179{
180 type Args = u32;
182
183 #[inline]
184 fn pack_as_with<W>(source: &T, mut writer: W, num_bits: Self::Args) -> Result<(), W::Error>
185 where
186 W: BitWriter,
187 {
188 let size_bits: u32 = bits_of::<T>() as u32;
189 let leading_zeroes = source.leading_zeros();
190 let used_bits = size_bits - leading_zeroes;
191 if num_bits < used_bits {
192 return Err(Error::custom(format!(
193 "{source:0b} cannot be packed into {num_bits} bits",
194 )));
195 }
196 let arr = source.to_be_bytes();
197 let bits = arr.as_bits();
198 writer.write_bitslice(&bits[bits.len() - num_bits as usize..])?;
199 Ok(())
200 }
201}
202
203impl<'de, T> BitUnpackAsWithArgs<'de, T> for VarNBits
204where
205 T: PrimInt,
206{
207 type Args = u32;
209
210 #[inline]
211 fn unpack_as_with<R>(mut reader: R, num_bits: Self::Args) -> Result<T, R::Error>
212 where
213 R: BitReader<'de>,
214 {
215 let size_bits: u32 = bits_of::<T>() as u32;
216 if num_bits > size_bits {
217 return Err(Error::custom("excessive bits for the type"));
218 }
219 let mut v: T = T::zero();
220 for bit in reader.unpack_iter::<bool>().take(num_bits as usize) {
221 v = v << 1;
222 v = v | if bit? { T::one() } else { T::zero() };
223 }
224 Ok(v)
225 }
226}
227
228#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
235pub struct VarNBytes;
236
237impl<T> BitPackAsWithArgs<T> for VarNBytes
238where
239 T: PrimInt + LowerHex + ToBytes,
240{
241 type Args = u32;
243
244 #[inline]
245 fn pack_as_with<W>(source: &T, mut writer: W, num_bytes: Self::Args) -> Result<(), W::Error>
246 where
247 W: BitWriter,
248 {
249 let size_bytes: u32 = size_of::<T>() as u32;
250 let leading_zeroes = source.leading_zeros();
251 let used_bytes = size_bytes - leading_zeroes / 8;
252 if num_bytes < used_bytes {
253 return Err(Error::custom(format!(
254 "{source:0x} cannot be packed into {num_bytes} bytes",
255 )));
256 }
257 let arr = source.to_be_bytes();
258 let bytes = arr.as_ref();
259 writer.write_bitslice((&bytes[bytes.len() - num_bytes as usize..]).as_bits())?;
260 Ok(())
261 }
262}
263
264impl<'de, T> BitUnpackAsWithArgs<'de, T> for VarNBytes
265where
266 T: PrimInt,
267{
268 type Args = u32;
270
271 #[inline]
272 fn unpack_as_with<R>(mut reader: R, num_bytes: Self::Args) -> Result<T, R::Error>
273 where
274 R: BitReader<'de>,
275 {
276 let size_bytes: u32 = size_of::<T>() as u32;
277 if num_bytes > size_bytes {
278 return Err(Error::custom("excessive bits for type"));
279 }
280 let mut v: T = T::zero();
281 for byte in reader.unpack_iter::<u8>().take(num_bytes as usize) {
282 v = v << 8;
283 v = v | T::from(byte?).unwrap();
284 }
285 Ok(v)
286 }
287}