valence_protocol/
var_int.rs

Help
1use std::io::{Read, Write};
2
3use anyhow::bail;
4use byteorder::ReadBytesExt;
5use derive_more::{Deref, DerefMut, From, Into};
6use serde::{Deserialize, Serialize};
7use thiserror::Error;
8
9use crate::{Decode, Encode};
10
11/// An `i32` encoded with variable length.
12#[derive(
13    Clone,
14    Copy,
15    Default,
16    PartialEq,
17    Eq,
18    PartialOrd,
19    Ord,
20    Hash,
21    Debug,
22    Deref,
23    DerefMut,
24    From,
25    Into,
26    Serialize,
27    Deserialize,
28)]
29#[serde(transparent)]
30#[repr(transparent)]
31pub struct VarInt(pub i32);
32
33impl VarInt {
34    /// The maximum number of bytes a `VarInt` could occupy when read from and
35    /// written to the Minecraft protocol.
36    pub const MAX_SIZE: usize = 5;
37
38    /// Returns the exact number of bytes this varint will write when
39    /// [`Encode::encode`] is called, assuming no error occurs.
40    pub const fn written_size(self) -> usize {
41        match self.0 {
42            0 => 1,
43            n => (31 - n.leading_zeros() as usize) / 7 + 1,
44        }
45    }
46
47    pub fn decode_partial<R: Read>(mut r: R) -> Result<i32, VarIntDecodeError> {
48        let mut val = 0;
49        for i in 0..Self::MAX_SIZE {
50            let byte = r.read_u8().map_err(|_| VarIntDecodeError::Incomplete)?;
51            val |= (i32::from(byte) & 0b01111111) << (i * 7);
52            if byte & 0b10000000 == 0 {
53                return Ok(val);
54            }
55        }
56
57        Err(VarIntDecodeError::TooLarge)
58    }
59}
60
61#[derive(Copy, Clone, PartialEq, Eq, Debug, Error)]
62pub enum VarIntDecodeError {
63    #[error("incomplete VarInt decode")]
64    Incomplete,
65    #[error("VarInt is too large")]
66    TooLarge,
67}
68
69impl Encode for VarInt {
70    // Adapted from VarInt-Simd encode
71    // https://github.com/as-com/varint-simd/blob/0f468783da8e181929b01b9c6e9f741c1fe09825/src/encode/mod.rs#L71
72    fn encode(&self, mut w: impl Write) -> anyhow::Result<()> {
73        let x = self.0 as u64;
74        let stage1 = (x & 0x000000000000007f)
75            | ((x & 0x0000000000003f80) << 1)
76            | ((x & 0x00000000001fc000) << 2)
77            | ((x & 0x000000000fe00000) << 3)
78            | ((x & 0x00000000f0000000) << 4);
79
80        let leading = stage1.leading_zeros();
81
82        let unused_bytes = (leading - 1) >> 3;
83        let bytes_needed = 8 - unused_bytes;
84
85        // set all but the last MSBs
86        let msbs = 0x8080808080808080;
87        let msbmask = 0xffffffffffffffff >> (((8 - bytes_needed + 1) << 3) - 1);
88
89        let merged = stage1 | (msbs & msbmask);
90        let bytes = merged.to_le_bytes();
91
92        w.write_all(unsafe { bytes.get_unchecked(..bytes_needed as usize) })?;
93
94        Ok(())
95    }
96}
97
98impl Decode<'_> for VarInt {
99    fn decode(r: &mut &[u8]) -> anyhow::Result<Self> {
100        let mut val = 0;
101        for i in 0..Self::MAX_SIZE {
102            let byte = r.read_u8()?;
103            val |= (i32::from(byte) & 0b01111111) << (i * 7);
104            if byte & 0b10000000 == 0 {
105                return Ok(VarInt(val));
106            }
107        }
108        bail!("VarInt is too large")
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use rand::{thread_rng, Rng};
115
116    use super::*;
117
118    #[test]
119    fn varint_written_size() {
120        let mut rng = thread_rng();
121        let mut buf = vec![];
122
123        for n in (0..100_000)
124            .map(|_| rng.gen())
125            .chain([0, i32::MIN, i32::MAX])
126            .map(VarInt)
127        {
128            buf.clear();
129            n.encode(&mut buf).unwrap();
130            assert_eq!(buf.len(), n.written_size());
131        }
132    }
133
134    #[test]
135    fn varint_round_trip() {
136        let mut rng = thread_rng();
137        let mut buf = vec![];
138
139        for n in (0..1_000_000)
140            .map(|_| rng.gen())
141            .chain([0, i32::MIN, i32::MAX])
142        {
143            VarInt(n).encode(&mut buf).unwrap();
144
145            let mut slice = buf.as_slice();
146            assert!(slice.len() <= VarInt::MAX_SIZE);
147
148            assert_eq!(n, VarInt::decode(&mut slice).unwrap().0);
149
150            assert!(slice.is_empty());
151            buf.clear();
152        }
153    }
154}
valence_protocol/var_int.rs

valence_protocol/
var_int.rs
