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    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/// Adapter for [`Var[U]Integer n`](https://docs.ton.org/develop/data-formats/msg-tlb#varuinteger-n)
93/// where `n` is *constant*.
94///
95/// ```tlb
96/// var_uint$_ {n:#} len:(#< n) value:(uint (len * 8)) = VarUInteger n;
97/// var_int$_ {n:#} len:(#< n) value:(int (len * 8)) = VarInteger n;
98/// ```
99/// See [`VarNBits`] for *dynamic* version.
100#[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            // BigUint::to_bytes_be() returns [0] instead of []
113            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/// Adapter for [`Var[U]Integer (n * 8)`](https://docs.ton.org/develop/data-formats/msg-tlb#varuinteger-n) where `n` is *dynamic*.
168/// ```tlb
169/// var_uint$_ {n:#} len:(#< n) value:(uint (len * 8)) = VarUInteger n;
170/// var_int$_ {n:#} len:(#< n) value:(int (len * 8)) = VarInteger n;
171/// ```
172/// See [`VarInt`] for *constant* version.
173#[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    /// number of bits
181    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    /// number of bits
208    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/// Adapter for [`Var[U]Integer n`](https://docs.ton.org/develop/data-formats/msg-tlb#varuinteger-n) where `n` is *dynamic*.
229/// ```tlb
230/// var_uint$_ {n:#} len:(#< n) value:(uint (len * 8)) = VarUInteger n;
231/// var_int$_ {n:#} len:(#< n) value:(int (len * 8)) = VarInteger n;
232/// ```
233/// See [`VarInt`] for *constant* version.
234#[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    /// number of bytes
242    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    /// number of bytes
269    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}