valence_protocol/
var_long.rs

Help
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    pub fn written_size(self) -> usize {
38        match self.0 {
39            0 => 1,
40            n => (63 - n.leading_zeros() as usize) / 7 + 1,
41        }
42    }
43}
44
45impl Encode for VarLong {
46    // Adapted from VarInt-Simd encode
47    // https://github.com/as-com/varint-simd/blob/0f468783da8e181929b01b9c6e9f741c1fe09825/src/encode/mod.rs#L71
48    #[cfg(all(
49        any(target_arch = "x86", target_arch = "x86_64"),
50        not(target_os = "macos")
51    ))]
52    fn encode(&self, mut w: impl Write) -> anyhow::Result<()> {
53        #[cfg(target_arch = "x86")]
54        use std::arch::x86::*;
55        #[cfg(target_arch = "x86_64")]
56        use std::arch::x86_64::*;
57
58        // Break the number into 7-bit parts and spread them out into a vector
59        let mut res = [0_u64; 2];
60        {
61            let x = self.0 as u64;
62
63            res[0] = unsafe { _pdep_u64(x, 0x7f7f7f7f7f7f7f7f) };
64            res[1] = unsafe { _pdep_u64(x >> 56, 0x000000000000017f) }
65        };
66        let stage1: __m128i = unsafe { std::mem::transmute(res) };
67
68        // Create a mask for where there exist values
69        // This signed comparison works because all MSBs should be cleared at this point
70        // Also handle the special case when num == 0
71        let minimum =
72            unsafe { _mm_set_epi8(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff_u8 as i8) };
73        let exists = unsafe { _mm_or_si128(_mm_cmpgt_epi8(stage1, _mm_setzero_si128()), minimum) };
74        let bits = unsafe { _mm_movemask_epi8(exists) };
75
76        // Count the number of bytes used
77        let bytes_needed = 32 - bits.leading_zeros() as u8; // lzcnt on supported CPUs
78
79        // Fill that many bytes into a vector
80        let ascend = unsafe { _mm_setr_epi8(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15) };
81        let mask = unsafe { _mm_cmplt_epi8(ascend, _mm_set1_epi8(bytes_needed as i8)) };
82
83        // Shift it down 1 byte so the last MSB is the only one set, and make sure only
84        // the MSB is set
85        let shift = unsafe { _mm_bsrli_si128(mask, 1) };
86        let msbmask = unsafe { _mm_and_si128(shift, _mm_set1_epi8(128_u8 as i8)) };
87
88        // Merge the MSB bits into the vector
89        let merged = unsafe { _mm_or_si128(stage1, msbmask) };
90        let bytes = unsafe { std::mem::transmute::<__m128i, [u8; 16]>(merged) };
91
92        w.write_all(unsafe { bytes.get_unchecked(..bytes_needed as usize) })?;
93
94        Ok(())
95    }
96
97    #[cfg(any(
98        not(any(target_arch = "x86", target_arch = "x86_64")),
99        target_os = "macos"
100    ))]
101    fn encode(&self, mut w: impl Write) -> anyhow::Result<()> {
102        use byteorder::WriteBytesExt;
103
104        let mut val = self.0 as u64;
105        loop {
106            if val & 0b1111111111111111111111111111111111111111111111111111111110000000 == 0 {
107                w.write_u8(val as u8)?;
108                return Ok(());
109            }
110            w.write_u8(val as u8 & 0b01111111 | 0b10000000)?;
111            val >>= 7;
112        }
113    }
114}
115
116impl Decode<'_> for VarLong {
117    fn decode(r: &mut &[u8]) -> anyhow::Result<Self> {
118        let mut val = 0;
119        for i in 0..Self::MAX_SIZE {
120            let byte = r.read_u8()?;
121            val |= (i64::from(byte) & 0b01111111) << (i * 7);
122            if byte & 0b10000000 == 0 {
123                return Ok(VarLong(val));
124            }
125        }
126        bail!("VarInt is too large")
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use rand::{thread_rng, Rng};
133
134    use super::*;
135
136    #[test]
137    fn encode_decode() {
138        let mut rng = thread_rng();
139        let mut buf = vec![];
140
141        for n in (0..1_000_000)
142            .map(|_| rng.gen())
143            .chain([0, i64::MIN, i64::MAX])
144        {
145            VarLong(n).encode(&mut buf).unwrap();
146
147            let mut slice = buf.as_slice();
148            assert!(slice.len() <= VarLong::MAX_SIZE);
149
150            assert_eq!(n, VarLong::decode(&mut slice).unwrap().0);
151            assert!(slice.is_empty());
152            buf.clear();
153        }
154    }
155}
valence_protocol/var_long.rs

valence_protocol/
var_long.rs
