chunkedge_binary/
var_long.rs

1use std::io::Write;
2
3use anyhow::bail;
4use byteorder::ReadBytesExt;
5use derive_more::{From, Into};
6use serde::{Deserialize, Serialize};
7
8use crate::{Decode, Encode};
9
10/// An `i64` encoded with variable length.
11#[derive(
12    Clone,
13    Copy,
14    Default,
15    PartialEq,
16    Eq,
17    PartialOrd,
18    Ord,
19    Hash,
20    Debug,
21    From,
22    Into,
23    Serialize,
24    Deserialize,
25)]
26#[serde(transparent)]
27#[repr(transparent)]
28pub struct VarLong(pub i64);
29
30impl VarLong {
31    /// The maximum number of bytes a `VarLong` can occupy when read from and
32    /// written to the Minecraft protocol.
33    pub const MAX_SIZE: usize = 10;
34
35    /// Returns the exact number of bytes this varlong will write when
36    /// [`Encode::encode`] is called, assuming no error occurs.
37    ///
38    /// This is the size of "canonical" (non-padded) encoding for this
39    /// value, not necessarily the number of bytes that were consumed
40    /// while decoding it. Minecraft `VarLongs` may be padded by using
41    /// more continuation bytes than the value requires, so callers that
42    /// need the decoded wire length should compare the input slice length
43    /// before and after decoding instead.
44    pub fn written_size(self) -> usize {
45        match self.0 {
46            0 => 1,
47            n => (63 - n.leading_zeros() as usize) / 7 + 1,
48        }
49    }
50}
51
52impl Encode for VarLong {
53    // Adapted from VarInt-Simd encode
54    // https://github.com/as-com/varint-simd/blob/0f468783da8e181929b01b9c6e9f741c1fe09825/src/encode/mod.rs#L71
55    #[cfg(all(
56        any(target_arch = "x86", target_arch = "x86_64"),
57        not(target_os = "macos")
58    ))]
59    fn encode(&self, mut w: impl Write) -> anyhow::Result<()> {
60        #[cfg(target_arch = "x86")]
61        use std::arch::x86::*;
62        #[cfg(target_arch = "x86_64")]
63        use std::arch::x86_64::*;
64
65        // Break the number into 7-bit parts and spread them out into a vector
66        let mut res = [0_u64; 2];
67        {
68            let x = self.0 as u64;
69
70            res[0] = unsafe { _pdep_u64(x, 0x7f7f7f7f7f7f7f7f) };
71            res[1] = unsafe { _pdep_u64(x >> 56, 0x000000000000017f) }
72        };
73        let stage1: __m128i = unsafe { std::mem::transmute(res) };
74
75        // Create a mask for where there exist values
76        // This signed comparison works because all MSBs should be cleared at this point
77        // Also handle the special case when num == 0
78        let minimum =
79            unsafe { _mm_set_epi8(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff_u8 as i8) };
80        let exists = unsafe { _mm_or_si128(_mm_cmpgt_epi8(stage1, _mm_setzero_si128()), minimum) };
81        let bits = unsafe { _mm_movemask_epi8(exists) };
82
83        // Count the number of bytes used
84        let bytes_needed = 32 - bits.leading_zeros() as u8; // lzcnt on supported CPUs
85
86        // Fill that many bytes into a vector
87        let ascend = unsafe { _mm_setr_epi8(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15) };
88        let mask = unsafe { _mm_cmplt_epi8(ascend, _mm_set1_epi8(bytes_needed as i8)) };
89
90        // Shift it down 1 byte so the last MSB is the only one set, and make sure only
91        // the MSB is set
92        let shift = unsafe { _mm_bsrli_si128(mask, 1) };
93        let msbmask = unsafe { _mm_and_si128(shift, _mm_set1_epi8(128_u8 as i8)) };
94
95        // Merge the MSB bits into the vector
96        let merged = unsafe { _mm_or_si128(stage1, msbmask) };
97        let bytes = unsafe { std::mem::transmute::<__m128i, [u8; 16]>(merged) };
98
99        w.write_all(unsafe { bytes.get_unchecked(..bytes_needed as usize) })?;
100
101        Ok(())
102    }
103
104    #[cfg(any(
105        not(any(target_arch = "x86", target_arch = "x86_64")),
106        target_os = "macos"
107    ))]
108    fn encode(&self, mut w: impl Write) -> anyhow::Result<()> {
109        use byteorder::WriteBytesExt;
110
111        let mut val = self.0 as u64;
112        loop {
113            if val & 0b1111111111111111111111111111111111111111111111111111111110000000 == 0 {
114                w.write_u8(val as u8)?;
115                return Ok(());
116            }
117            w.write_u8(val as u8 & 0b01111111 | 0b10000000)?;
118            val >>= 7;
119        }
120    }
121}
122
123impl Decode<'_> for VarLong {
124    fn decode(r: &mut &[u8]) -> anyhow::Result<Self> {
125        let mut val = 0;
126        for i in 0..Self::MAX_SIZE {
127            let byte = r.read_u8()?;
128            val |= (i64::from(byte) & 0b01111111) << (i * 7);
129            if byte & 0b10000000 == 0 {
130                return Ok(VarLong(val));
131            }
132        }
133        bail!("VarInt is too large")
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use rand::{RngExt, rng};
140
141    use super::*;
142
143    #[test]
144    fn encode_decode() {
145        let mut rng = rng();
146        let mut buf = vec![];
147
148        for n in (0..1_000_000)
149            .map(|_| rng.random::<i64>())
150            .chain([0, i64::MIN, i64::MAX])
151        {
152            VarLong(n).encode(&mut buf).unwrap();
153
154            let mut slice = buf.as_slice();
155            assert!(slice.len() <= VarLong::MAX_SIZE);
156
157            assert_eq!(n, VarLong::decode(&mut slice).unwrap().0);
158            assert!(slice.is_empty());
159            buf.clear();
160        }
161    }
162}
chunkedge_binary/var_long.rs

chunkedge_binary/
var_long.rs
