1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
use bytes::{Buf, BufMut};

use crate::error::*;

const SIZE_OF_SIZE_MARKER: u32 = 2;
const FITS_IN_ONE_BYTE: u64 = 2_u64.pow(8 - SIZE_OF_SIZE_MARKER) - 1;
const FITS_IN_TWO_BYTES: u64 = 2_u64.pow(16 - SIZE_OF_SIZE_MARKER) - 1;
const FITS_IN_FOUR_BYTES: u64 = 2_u64.pow(32 - SIZE_OF_SIZE_MARKER) - 1;

/// Read an epee variable sized number from `r`.
///
/// ```rust
/// use cuprate_epee_encoding::read_varint;
///
/// assert_eq!(read_varint(&mut [252].as_slice()).unwrap(), 63);
/// assert_eq!(read_varint(&mut [1, 1].as_slice()).unwrap(), 64);
/// assert_eq!(read_varint(&mut [253, 255].as_slice()).unwrap(), 16_383);
/// assert_eq!(read_varint(&mut [2, 0, 1, 0].as_slice()).unwrap(), 16_384);
/// assert_eq!(read_varint(&mut [254, 255, 255, 255].as_slice()).unwrap(), 1_073_741_823);
/// assert_eq!(read_varint(&mut [3, 0, 0, 0, 1, 0, 0, 0].as_slice()).unwrap(), 1_073_741_824);
/// ```
pub fn read_varint<B: Buf>(r: &mut B) -> Result<u64> {
    if !r.has_remaining() {
        Err(Error::IO("Not enough bytes to build VarInt"))?
    }

    let vi_start = r.get_u8();
    let len = 1 << (vi_start & 0b11);

    if r.remaining() < len - 1 {
        Err(Error::IO("Not enough bytes to build VarInt"))?
    }

    let mut vi = u64::from(vi_start >> 2);
    for i in 1..len {
        vi |= u64::from(r.get_u8()) << (((i - 1) * 8) + 6);
    }
    Ok(vi)
}

/// Write an epee variable sized number into `w`.
///
/// ```rust
/// use cuprate_epee_encoding::write_varint;
///
/// let mut buf = vec![];
///
/// for (number, expected_bytes) in [
///     (63, [252].as_slice()),
///     (64, [1, 1].as_slice()),
///     (16_383, [253, 255].as_slice()),
///     (16_384, [2, 0, 1, 0].as_slice()),
///     (1_073_741_823, [254, 255, 255, 255].as_slice()),
///     (1_073_741_824, [3, 0, 0, 0, 1, 0, 0, 0].as_slice()),
/// ] {
///     buf.clear();
///     write_varint(number, &mut buf);
///     assert_eq!(buf.as_slice(), expected_bytes);
/// }
/// ```
pub fn write_varint<B: BufMut>(number: u64, w: &mut B) -> Result<()> {
    let size_marker = match number {
        0..=FITS_IN_ONE_BYTE => 0,
        64..=FITS_IN_TWO_BYTES => 1,
        16384..=FITS_IN_FOUR_BYTES => 2,
        _ => 3,
    };

    if w.remaining_mut() < 1 << size_marker {
        Err(Error::IO("Not enough capacity to write VarInt"))?;
    }

    let number = (number << 2) | size_marker;

    // Although `as` is unsafe we just checked the length.
    match size_marker {
        0 => w.put_u8(number as u8),
        1 => w.put_u16_le(number as u16),
        2 => w.put_u32_le(number as u32),
        3 => w.put_u64_le(number),
        _ => unreachable!(),
    }

    Ok(())
}

#[cfg(test)]
mod tests {

    use alloc::vec::Vec;

    use crate::varint::*;

    fn assert_varint_length(number: u64, len: usize) {
        let mut w = Vec::new();
        write_varint(number, &mut w).unwrap();
        assert_eq!(w.len(), len);
    }

    fn assert_varint_val(mut varint: &[u8], val: u64) {
        assert_eq!(read_varint(&mut varint).unwrap(), val);
    }

    #[test]
    fn varint_write_length() {
        assert_varint_length(FITS_IN_ONE_BYTE, 1);
        assert_varint_length(FITS_IN_ONE_BYTE + 1, 2);
        assert_varint_length(FITS_IN_TWO_BYTES, 2);
        assert_varint_length(FITS_IN_TWO_BYTES + 1, 4);
        assert_varint_length(FITS_IN_FOUR_BYTES, 4);
        assert_varint_length(FITS_IN_FOUR_BYTES + 1, 8);
    }

    #[test]
    fn varint_read() {
        assert_varint_val(&[252], FITS_IN_ONE_BYTE);
        assert_varint_val(&[1, 1], FITS_IN_ONE_BYTE + 1);
        assert_varint_val(&[253, 255], FITS_IN_TWO_BYTES);
        assert_varint_val(&[2, 0, 1, 0], FITS_IN_TWO_BYTES + 1);
        assert_varint_val(&[254, 255, 255, 255], FITS_IN_FOUR_BYTES);
        assert_varint_val(&[3, 0, 0, 0, 1, 0, 0, 0], FITS_IN_FOUR_BYTES + 1);
    }
}