binserde/
varint.rs

1use std::io;
2use std::io::{Read, Write};
3
4use byteorder::ReadBytesExt;
5
6pub fn encode_min(num: i64) -> u64 {
7    let u_num = num as u64;
8    (u_num << 1 ^ (num >> 63) as u64) | u_num >> 63
9}
10
11pub fn decode_min(num: u64) -> i64 {
12    (num >> 1) as i64 ^ ((num << 63) as i64) >> 63
13}
14
15pub fn varint_write<W: Write>(num: u64, mut pipe: W) -> io::Result<usize> {
16    let mut num_pos = 0;
17    let mut idx = 0;
18    let mut buf = [0; 9];
19    let data_bits = 64 - num.leading_zeros();
20
21    loop {
22        let next = if data_bits - num_pos > 7 {
23            0b10000000
24        } else {
25            0
26        };
27        let piece = (num >> num_pos) as u8 & 0b01111111 | next;
28        buf[idx] = piece;
29        idx += 1;
30        num_pos += 7;
31
32        if !(num_pos < data_bits) {
33            break;
34        }
35    }
36
37    pipe.write(&buf[..idx])?;
38
39    Ok(idx)
40}
41
42pub fn varint_read<R: Read>(mut pipe: R) -> io::Result<u64> {
43    let mut offset = 0;
44    let mut num = 0;
45
46    loop {
47        let byte = pipe.read_u8()?;
48        let has_next = byte & 0b10000000 != 0;
49        num |= (byte as u64 & 0b01111111) << offset;
50        offset += 7;
51
52        if !has_next {
53            break;
54        }
55    }
56
57    Ok(num)
58}
59
60#[cfg(test)]
61mod test {
62    use crate::varint::{decode_min, encode_min, varint_read, varint_write};
63
64    #[test]
65    fn test_encode_min() {
66        for i in -5..5 {
67            let r = encode_min(i);
68            println!("{} = {:016X}", i, r);
69            assert_eq!(i, decode_min(r));
70        }
71
72        for i in i64::MAX - 5..=i64::MAX {
73            let r = encode_min(i);
74            println!("{} = {:016X}", i, r);
75            assert_eq!(i, decode_min(r));
76        }
77
78        for i in i64::MIN..i64::MIN + 5 {
79            let r = encode_min(i);
80            println!("{} = {:016X}", i, r);
81            assert_eq!(i, decode_min(r));
82        }
83    }
84
85    #[test]
86    fn test_varint() {
87        use std::io::Cursor;
88
89        let mut buf = Cursor::new(Vec::new());
90
91        varint_write(0, &mut buf).unwrap();
92        varint_write(16, &mut buf).unwrap();
93        varint_write(234567892322414124, &mut buf).unwrap();
94        varint_write(encode_min(-18), &mut buf).unwrap();
95        varint_write(encode_min(20000000), &mut buf).unwrap();
96
97        buf.set_position(0);
98        for item in buf.get_ref().iter() {
99            print!("{:02X}", item);
100            if item & 0b10000000 == 0 {
101                print!("] ");
102            } else {
103                print!(" ");
104            }
105        }
106        println!();
107
108        assert_eq!(0, varint_read(&mut buf).unwrap());
109        assert_eq!(16, varint_read(&mut buf).unwrap());
110        assert_eq!(234567892322414124, varint_read(&mut buf).unwrap());
111        assert_eq!(encode_min(-18), varint_read(&mut buf).unwrap());
112        assert_eq!(encode_min(20000000), varint_read(&mut buf).unwrap());
113    }
114}