mc_network_io/
var_int.rs

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        // set all but the last MSBs
63        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