cuprate_epee_encoding/
varint.rs

1use bytes::{Buf, BufMut};
2
3use crate::error::*;
4
5const SIZE_OF_SIZE_MARKER: u32 = 2;
6const FITS_IN_ONE_BYTE: u64 = 2_u64.pow(8 - SIZE_OF_SIZE_MARKER) - 1;
7const FITS_IN_TWO_BYTES: u64 = 2_u64.pow(16 - SIZE_OF_SIZE_MARKER) - 1;
8const FITS_IN_FOUR_BYTES: u64 = 2_u64.pow(32 - SIZE_OF_SIZE_MARKER) - 1;
9
10/// Read an epee variable sized number from `r`.
11///
12/// ```rust
13/// use cuprate_epee_encoding::read_varint;
14///
15/// assert_eq!(read_varint::<_, u64>(&mut [252].as_slice()).unwrap(), 63);
16/// assert_eq!(read_varint::<_, u64>(&mut [1, 1].as_slice()).unwrap(), 64);
17/// assert_eq!(read_varint::<_, u64>(&mut [253, 255].as_slice()).unwrap(), 16_383);
18/// assert_eq!(read_varint::<_, u64>(&mut [2, 0, 1, 0].as_slice()).unwrap(), 16_384);
19/// assert_eq!(read_varint::<_, u64>(&mut [254, 255, 255, 255].as_slice()).unwrap(), 1_073_741_823);
20/// assert_eq!(read_varint::<_, u64>(&mut [3, 0, 0, 0, 1, 0, 0, 0].as_slice()).unwrap(), 1_073_741_824);
21/// ```
22pub fn read_varint<B: Buf, T: TryFrom<u64>>(r: &mut B) -> Result<T> {
23    if !r.has_remaining() {
24        return Err(Error::IO("Not enough bytes to build VarInt"));
25    }
26
27    let vi_start = r.get_u8();
28    let len = 1 << (vi_start & 0b11);
29
30    if r.remaining() < len - 1 {
31        return Err(Error::IO("Not enough bytes to build VarInt"));
32    }
33
34    let mut vi = u64::from(vi_start >> 2);
35    for i in 1..len {
36        vi |= u64::from(r.get_u8()) << (((i - 1) * 8) + 6);
37    }
38
39    vi.try_into().map_err(|_| Error::IO("VarInt is too big"))
40}
41
42/// Write an epee variable sized number into `w`.
43///
44/// ```rust
45/// use cuprate_epee_encoding::write_varint;
46///
47/// let mut buf = vec![];
48///
49/// for (number, expected_bytes) in [
50///     (63, [252].as_slice()),
51///     (64, [1, 1].as_slice()),
52///     (16_383, [253, 255].as_slice()),
53///     (16_384, [2, 0, 1, 0].as_slice()),
54///     (1_073_741_823, [254, 255, 255, 255].as_slice()),
55///     (1_073_741_824, [3, 0, 0, 0, 1, 0, 0, 0].as_slice()),
56/// ] {
57///     buf.clear();
58///     write_varint(number, &mut buf);
59///     assert_eq!(buf.as_slice(), expected_bytes);
60/// }
61/// ```
62pub fn write_varint<B: BufMut, T: TryInto<u64>>(number: T, w: &mut B) -> Result<()> {
63    let number = number
64        .try_into()
65        .map_err(|_| "Tried to write a varint bigger than 64-bits")
66        .unwrap();
67
68    let size_marker = match number {
69        0..=FITS_IN_ONE_BYTE => 0,
70        64..=FITS_IN_TWO_BYTES => 1,
71        16384..=FITS_IN_FOUR_BYTES => 2,
72        _ => 3,
73    };
74
75    if w.remaining_mut() < 1 << size_marker {
76        return Err(Error::IO("Not enough capacity to write VarInt"));
77    }
78
79    let number = (number << 2) | size_marker;
80
81    #[expect(
82        clippy::cast_possible_truncation,
83        reason = "Although `as` is unsafe we just checked the length."
84    )]
85    match size_marker {
86        0 => w.put_u8(number as u8),
87        1 => w.put_u16_le(number as u16),
88        2 => w.put_u32_le(number as u32),
89        3 => w.put_u64_le(number),
90        _ => unreachable!(),
91    }
92
93    Ok(())
94}
95
96#[cfg(test)]
97mod tests {
98
99    use alloc::vec::Vec;
100
101    use crate::varint::*;
102
103    fn assert_varint_length(number: u64, len: usize) {
104        let mut w = Vec::new();
105        write_varint(number, &mut w).unwrap();
106        assert_eq!(w.len(), len);
107    }
108
109    fn assert_varint_val(mut varint: &[u8], val: u64) {
110        assert_eq!(read_varint::<_, u64>(&mut varint).unwrap(), val);
111    }
112
113    #[test]
114    fn varint_write_length() {
115        assert_varint_length(FITS_IN_ONE_BYTE, 1);
116        assert_varint_length(FITS_IN_ONE_BYTE + 1, 2);
117        assert_varint_length(FITS_IN_TWO_BYTES, 2);
118        assert_varint_length(FITS_IN_TWO_BYTES + 1, 4);
119        assert_varint_length(FITS_IN_FOUR_BYTES, 4);
120        assert_varint_length(FITS_IN_FOUR_BYTES + 1, 8);
121    }
122
123    #[test]
124    fn varint_read() {
125        assert_varint_val(&[252], FITS_IN_ONE_BYTE);
126        assert_varint_val(&[1, 1], FITS_IN_ONE_BYTE + 1);
127        assert_varint_val(&[253, 255], FITS_IN_TWO_BYTES);
128        assert_varint_val(&[2, 0, 1, 0], FITS_IN_TWO_BYTES + 1);
129        assert_varint_val(&[254, 255, 255, 255], FITS_IN_FOUR_BYTES);
130        assert_varint_val(&[3, 0, 0, 0, 1, 0, 0, 0], FITS_IN_FOUR_BYTES + 1);
131    }
132}
cuprate_epee_encoding/varint.rs

cuprate_epee_encoding/
varint.rs
