chunkedge_binary/
var_int.rs

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    ///
41    /// This is the size of "canonical" (non-padded) encoding for this
42    /// value, not necessarily the number of bytes that were consumed
43    /// while decoding it. Minecraft `VarInts` may be padded by using
44    /// more continuation bytes than the value requires, so callers that
45    /// need the decoded wire length should compare the input slice length
46    /// before and after decoding instead.
47    pub const fn written_size(self) -> usize {
48        match self.0 {
49            0 => 1,
50            n => (31 - n.leading_zeros() as usize) / 7 + 1,
51        }
52    }
53
54    pub fn decode_partial<R: Read>(mut r: R) -> Result<i32, VarIntDecodeError> {
55        let mut val = 0;
56        for i in 0..Self::MAX_SIZE {
57            let byte = r.read_u8().map_err(|_| VarIntDecodeError::Incomplete)?;
58            val |= (i32::from(byte) & 0b01111111) << (i * 7);
59            if byte & 0b10000000 == 0 {
60                return Ok(val);
61            }
62        }
63
64        Err(VarIntDecodeError::TooLarge)
65    }
66}
67
68#[derive(Copy, Clone, PartialEq, Eq, Debug, Error)]
69pub enum VarIntDecodeError {
70    #[error("incomplete VarInt decode")]
71    Incomplete,
72    #[error("VarInt is too large")]
73    TooLarge,
74}
75
76impl Encode for VarInt {
77    // Adapted from VarInt-Simd encode
78    // https://github.com/as-com/varint-simd/blob/0f468783da8e181929b01b9c6e9f741c1fe09825/src/encode/mod.rs#L71
79    fn encode(&self, mut w: impl Write) -> anyhow::Result<()> {
80        let x = self.0 as u64;
81        let stage1 = (x & 0x000000000000007f)
82            | ((x & 0x0000000000003f80) << 1)
83            | ((x & 0x00000000001fc000) << 2)
84            | ((x & 0x000000000fe00000) << 3)
85            | ((x & 0x00000000f0000000) << 4);
86
87        let leading = stage1.leading_zeros();
88
89        let unused_bytes = (leading - 1) >> 3;
90        let bytes_needed = 8 - unused_bytes;
91
92        // set all but the last MSBs
93        let msbs = 0x8080808080808080;
94        let msbmask = 0xffffffffffffffff >> (((8 - bytes_needed + 1) << 3) - 1);
95
96        let merged = stage1 | (msbs & msbmask);
97        let bytes = merged.to_le_bytes();
98
99        w.write_all(unsafe { bytes.get_unchecked(..bytes_needed as usize) })?;
100
101        Ok(())
102    }
103}
104
105impl Decode<'_> for VarInt {
106    fn decode(r: &mut &[u8]) -> anyhow::Result<Self> {
107        let mut val = 0;
108        for i in 0..Self::MAX_SIZE {
109            let byte = r.read_u8()?;
110            val |= (i32::from(byte) & 0b01111111) << (i * 7);
111            if byte & 0b10000000 == 0 {
112                return Ok(VarInt(val));
113            }
114        }
115        bail!("VarInt is too large")
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use rand::{RngExt, rng};
122
123    use super::*;
124
125    #[test]
126    fn varint_written_size() {
127        let mut rng = rng();
128        let mut buf = vec![];
129
130        for n in (0..100_000)
131            .map(|_| rng.random::<i32>())
132            .chain([0, i32::MIN, i32::MAX])
133            .map(VarInt)
134        {
135            buf.clear();
136            n.encode(&mut buf).unwrap();
137            assert_eq!(buf.len(), n.written_size());
138        }
139    }
140
141    #[test]
142    fn varint_round_trip() {
143        let mut rng = rng();
144        let mut buf = vec![];
145
146        for n in (0..1_000_000)
147            .map(|_| rng.random::<i32>())
148            .chain([0, i32::MIN, i32::MAX])
149        {
150            VarInt(n).encode(&mut buf).unwrap();
151
152            let mut slice = buf.as_slice();
153            assert!(slice.len() <= VarInt::MAX_SIZE);
154
155            assert_eq!(n, VarInt::decode(&mut slice).unwrap().0);
156
157            assert!(slice.is_empty());
158            buf.clear();
159        }
160    }
161}
chunkedge_binary/var_int.rs

chunkedge_binary/
var_int.rs
