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
124
125
126
127
128

use bytes::Buf;
use std::cmp::min;
use std::io::{
    self,
    Write,
};

#[inline]
pub const fn serialized_length(mut val: u64) -> usize {
    let mut ans = 1;
    while val & !0x7Fu64 != 0 {
        val >>= 7;
        ans += 1;
    }

    ans
}

#[inline]
pub fn write<W: Write>(mut val: u64, w: &mut W) -> io::Result<()> {
    // the maximum length in bytes of an encoded (64-bit) varint
    const MAX_LEN: usize = serialized_length(u64::MAX);

    let mut len = 0;
    let mut buf = [0; MAX_LEN];
    loop {
        if val & !0x7Fu64 == 0 {
            buf[len] = val as u8;
            len += 1;
            break;
        } else {
            buf[len] = ((val & 0x7F) | 0x80) as u8;
            val >>= 7;
            len += 1;
        }
    }
    w.write_all(&buf[..len])
}

#[inline]
pub fn read<B: Buf>(buf: &mut B) -> io::Result<Option<u64>> {
    if buf.remaining() == 0 {
        return Ok(None);
    }

    // Find offset of first byte with MSB not set
    let idx = match buf
        .chunk()
        .iter()
        .enumerate()
        .find(|&(_, val)| *val < 0x80 as u8)
        .map(|(idx, _)| idx)
    {
        Some(idx) => idx,
        // Fallback to per-byte read if data is span across multiple slices.
        None => return read_slow(buf),
    };

    let varint = {
        let buf = &buf.chunk()[..=idx];

        let mut r: u64 = 0;
        for byte in buf.iter().rev() {
            r = (r << 7) | u64::from(byte & 0x7f);
        }

        r
    };

    buf.advance(idx + 1);
    Ok(Some(varint))
}

fn read_slow<B: Buf>(buf: &mut B) -> io::Result<Option<u64>> {
    let mut r: u64 = 0;
    for index in 0..min(10, buf.remaining()) {
        let byte = buf.get_u8();
        r |= u64::from(byte & 0x7f) << (index * 7);
        if byte < 0x80 {
            return Ok(Some(r));
        }
    }
    Err(super::unexpected_eof())
}

#[inline]
pub fn ensure_read<B: Buf>(buf: &mut B) -> io::Result<u64> {
    match read(buf)? {
        Some(n) => Ok(n),
        None => Err(super::unexpected_eof()),
    }
}

#[test]
fn test_basic() {
    use crate::Message;
    use std::io::Cursor;

    let from_vec = |vec| read(&mut Cursor::new(vec)).unwrap().unwrap();

    let from_vec_split = |mut first_vec: Vec<u8>| {
        let at = first_vec.len() / 2;
        let second_vec = first_vec.split_off(at);
        read(&mut Cursor::new(first_vec).chain(Cursor::new(second_vec)))
            .unwrap()
            .unwrap()
    };

    assert_eq!(1000.serialize_to_vec(), vec![232, 7]);
    assert_eq!(from_vec(vec![232, 7]), 1000);
    assert_eq!(from_vec_split(vec![232, 7]), 1000);

    let minus1 = vec![255, 255, 255, 255, 255, 255, 255, 255, 255, 1];
    assert_eq!((-1 as i32).serialize_to_vec(), minus1);
    assert_eq!((-1 as i64).serialize_to_vec(), minus1);
    assert_eq!(from_vec(minus1.clone()) as i32, -1);
    assert_eq!(from_vec(minus1.clone()) as i64, -1);
    assert_eq!(from_vec_split(minus1.clone()) as i32, -1);
    assert_eq!(from_vec_split(minus1) as i64, -1);

    let minus1000 = vec![152, 248, 255, 255, 255, 255, 255, 255, 255, 1];
    assert_eq!((-1000 as i32).serialize_to_vec(), minus1000);
    assert_eq!((-1000 as i64).serialize_to_vec(), minus1000);
    assert_eq!(from_vec(minus1000.clone()) as i32, -1000);
    assert_eq!(from_vec(minus1000.clone()) as i64, -1000);
    assert_eq!(from_vec_split(minus1000.clone()) as i32, -1000);
    assert_eq!(from_vec_split(minus1000) as i64, -1000);
}