1use std::io::{Error, ErrorKind, Result};
2
3pub trait VarIntRead {
4 fn read_var_i32(&mut self) -> Result<i32>;
5 fn read_var_i64(&mut self) -> Result<i64>;
6}
7
8pub trait VarIntWrite {
9 fn write_var_i32(&mut self, value: i32) -> Result<usize>;
10 fn write_var_i64(&mut self, value: i64) -> Result<usize>;
11}
12
13impl<R> VarIntRead for R
14where
15 R: std::io::Read,
16{
17 fn read_var_i32(&mut self) -> Result<i32> {
18 let mut val = 0;
19 let buf = &mut [0u8];
20 for i in 0..5 {
21 self.read_exact(buf)?;
22 let byte = buf[0];
23 val |= (byte as i32 & 0b01111111) << (i * 7);
24 if byte & 0b10000000 == 0 {
25 return Ok(val);
26 }
27 }
28 Err(Error::new(ErrorKind::InvalidInput, "VarInt is too large"))
29 }
30
31 fn read_var_i64(&mut self) -> Result<i64> {
32 let mut buf = [0];
33 let mut ans = 0;
34 for i in 0..8 {
35 self.read_exact(&mut buf)?;
36 ans |= ((buf[0] & 0b0111_1111) as i64) << 7 * i;
37 if buf[0] & 0b1000_0000 == 0 {
38 break;
39 }
40 }
41 Ok(ans)
42 }
43}
44
45impl<W> VarIntWrite for W
46where
47 W: std::io::Write,
48{
49 fn write_var_i32(&mut self, value: i32) -> Result<usize> {
50 let x = value as u64;
51 let stage1 = (x & 0x000000000000007f)
52 | ((x & 0x0000000000003f80) << 1)
53 | ((x & 0x00000000001fc000) << 2)
54 | ((x & 0x000000000fe00000) << 3)
55 | ((x & 0x00000000f0000000) << 4);
56
57 let leading = stage1.leading_zeros();
58
59 let unused_bytes = (leading - 1) >> 3;
60 let bytes_needed = 8 - unused_bytes;
61
62 let msbs = 0x8080808080808080;
64 let msbmask = 0xffffffffffffffff >> (((8 - bytes_needed + 1) << 3) - 1);
65
66 let merged = stage1 | (msbs & msbmask);
67 let bytes = merged.to_le_bytes();
68
69 Ok(self.write(unsafe { bytes.get_unchecked(..bytes_needed as usize) })?)
70 }
71
72 fn write_var_i64(&mut self, mut value: i64) -> Result<usize> {
73 let mut buf = [0];
74 let mut cnt = 0;
75 while value != 0 {
76 buf[0] = (value & 0b0111_1111) as u8;
77 value = (value >> 7) & (i64::max_value() >> 6);
78 if value != 0 {
79 buf[0] |= 0b1000_0000;
80 }
81 cnt += self.write(&mut buf)?;
82 }
83 Ok(cnt)
84 }
85}
86
87#[test]
88fn test_var_i32() {
89 fn test(first: Vec<u8>, second: i32) {
90 let mut vec = Vec::new();
91 vec.write_var_i32(second).unwrap();
92 assert_eq!(first, vec);
93 let read = std::io::Cursor::new(vec).read_var_i32().unwrap();
94 assert_eq!(read, second);
95 }
96 test(vec![0x00], 0);
97 test(vec![0x01], 1);
98 test(vec![0x02], 2);
99 test(vec![0x7f], 127);
100 test(vec![0x80, 0x01], 128);
101 test(vec![0xff, 0x01], 255);
102 test(vec![0xdd, 0xc7, 0x01], 25565);
103 test(vec![0xff, 0xff, 0x7f], 2097151);
104 test(vec![0xff, 0xff, 0xff, 0xff, 0x07], 2147483647);
105 test(vec![0xff, 0xff, 0xff, 0xff, 0x0f], -1);
106 test(vec![0x80, 0x80, 0x80, 0x80, 0x08], -2147483648);
107}
108