ordered_varint/
unsigned.rs

1use std::io::{Read, Write};
2use std::num::TryFromIntError;
3
4use crate::Variable;
5
6/// An unsigned integer value
7///
8/// This type encodes values in the range `0..2.pow(124)` by using the first 4
9/// bits to denote an unsigned byte `length`. This length ranges from `0..=15`.
10/// The remaining 4 bits of the first byte and any additional bytes are then
11/// used to store the integer in big-endian encoding.
12#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Ord, PartialOrd)]
13pub struct Unsigned(pub(crate) u128);
14
15impl Unsigned {
16    pub(crate) fn encode_be_bytes<W: Write, const N: usize>(
17        mut value: [u8; N],
18        mut output: W,
19    ) -> std::io::Result<usize> {
20        // Because we encode "extra bytes" in 4 bits, we must keep the extra
21        // bytes to 15 or less. This only affects 128-bit encoding
22        if N == 16 && value[0] >> 4 != 0 {
23            return Err(std::io::Error::from(std::io::ErrorKind::InvalidData));
24        }
25
26        let (total_length, extra_bytes) = value
27            .iter()
28            .enumerate()
29            .find_map(|(index, &byte)| {
30                if byte > 0 {
31                    let extra_bytes = N - 1 - index;
32                    if byte < 16 {
33                        Some((extra_bytes + 1, extra_bytes))
34                    } else {
35                        Some((extra_bytes + 2, extra_bytes + 1))
36                    }
37                } else {
38                    None
39                }
40            })
41            .unwrap_or((0, 0));
42        let total_length = total_length.max(1);
43
44        let extra_bytes_encoded = (extra_bytes as u8) << 4;
45        if total_length > N {
46            // We need an extra byte to store the length information
47            output.write_all(&[extra_bytes_encoded])?;
48            output.write_all(&value)?;
49        } else {
50            value[N - total_length] |= extra_bytes_encoded;
51            output.write_all(&value[N - total_length..])?;
52        }
53
54        Ok(total_length)
55    }
56
57    pub(crate) fn decode_variable_bytes<R: Read, const N: usize>(
58        mut input: R,
59    ) -> std::io::Result<[u8; N]> {
60        let mut buffer = [0_u8; N];
61        input.read_exact(&mut buffer[0..1])?;
62        let first_byte = buffer[0];
63        let length = (first_byte >> 4) as usize;
64        if length > N {
65            return Err(std::io::Error::from(std::io::ErrorKind::InvalidData));
66        }
67        input.read_exact(&mut buffer[N - length..])?;
68        match N - length {
69            0 => {
70                // We overwrite the first byte with the read operation, so we need
71                // to fill back in the bits from the first byte.
72                buffer[0] |= first_byte & 0b1111;
73            }
74            1 => {
75                // Clear the top 4 bits of the first byte. The lower 4 bits may
76                // still have data in them.
77                buffer[0] &= 0b1111;
78            }
79            _ => {
80                // Move the first byte's data into the last byte read, then
81                // clear our initial byte.
82                buffer[N - 1 - length] |= first_byte & 0b1111;
83                buffer[0] = 0;
84            }
85        }
86        Ok(buffer)
87    }
88}
89
90impl Variable for Unsigned {
91    fn encode_variable<W: Write>(&self, output: W) -> std::io::Result<usize> {
92        Self::encode_be_bytes(self.0.to_be_bytes(), output)
93    }
94
95    fn decode_variable<R: Read>(input: R) -> std::io::Result<Self> {
96        let buffer = Self::decode_variable_bytes(input)?;
97
98        Ok(Self(u128::from_be_bytes(buffer)))
99    }
100}
101
102macro_rules! impl_primitive_from_varint {
103    ($ty:ty) => {
104        impl TryFrom<Unsigned> for $ty {
105            type Error = TryFromIntError;
106
107            fn try_from(value: Unsigned) -> Result<Self, Self::Error> {
108                value.0.try_into()
109            }
110        }
111    };
112}
113
114macro_rules! impl_varint_from_primitive {
115    ($ty:ty, $dest:ty) => {
116        impl From<$ty> for Unsigned {
117            fn from(value: $ty) -> Self {
118                Self(<$dest>::from(value))
119            }
120        }
121    };
122}
123
124impl_varint_from_primitive!(u8, u128);
125impl_varint_from_primitive!(u16, u128);
126impl_varint_from_primitive!(u32, u128);
127impl_varint_from_primitive!(u64, u128);
128impl_varint_from_primitive!(u128, u128);
129
130impl_primitive_from_varint!(u8);
131impl_primitive_from_varint!(u16);
132impl_primitive_from_varint!(u32);
133impl_primitive_from_varint!(u64);
134impl_primitive_from_varint!(usize);
135
136impl From<Unsigned> for u128 {
137    fn from(value: Unsigned) -> Self {
138        value.0
139    }
140}
141
142impl From<usize> for Unsigned {
143    fn from(value: usize) -> Self {
144        Self(value as u128)
145    }
146}