Skip to main content

tlbits/as/
integer.rs

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/// Adapter for [`Var[U]Integer n`](https://docs.ton.org/develop/data-formats/msg-tlb#varuinteger-n)
100/// where `n` is *constant*.
101///
102/// ```tlb
103/// var_uint$_ {n:#} len:(#< n) value:(uint (len * 8)) = VarUInteger n;
104/// var_int$_ {n:#} len:(#< n) value:(int (len * 8)) = VarInteger n;
105/// ```
106/// See [`VarNBits`] for *dynamic* version.
107#[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            // BigUint::to_bytes_be() returns [0] instead of []
122            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/// Adapter for [`Var[U]Integer (n * 8)`](https://docs.ton.org/develop/data-formats/msg-tlb#varuinteger-n) where `n` is *dynamic*.
183/// ```tlb
184/// var_uint$_ {n:#} len:(#< n) value:(uint (len * 8)) = VarUInteger n;
185/// var_int$_ {n:#} len:(#< n) value:(int (len * 8)) = VarInteger n;
186/// ```
187/// See [`VarInt`] for *constant* version.
188#[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    /// number of bits
196    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    /// number of bits
223    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/// Adapter for [`Var[U]Integer n`](https://docs.ton.org/develop/data-formats/msg-tlb#varuinteger-n) where `n` is *dynamic*.
244/// ```tlb
245/// var_uint$_ {n:#} len:(#< n) value:(uint (len * 8)) = VarUInteger n;
246/// var_int$_ {n:#} len:(#< n) value:(int (len * 8)) = VarInteger n;
247/// ```
248/// See [`VarInt`] for *constant* version.
249#[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    /// number of bytes
257    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    /// number of bytes
284    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}