steel_utils/codec/
variable_integer.rs

1use std::io::{Cursor, Error, Write};
2
3#[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
4use std::arch::x86_64::_pdep_u64;
5
6use tokio::io::{AsyncRead, AsyncReadExt};
7
8use crate::{
9    FrontVec,
10    serial::{ReadFrom, WriteTo},
11};
12
13// Pack the low 56 payload bits into eight byte lanes, leaving continuation bits clear
14#[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
15#[inline]
16fn pack_payload(value: u64) -> u64 {
17    // SAFETY: ts impl is compiled only for targets with BMI2
18    unsafe { _pdep_u64(value, 0x7f7f_7f7f_7f7f_7f7f) }
19}
20
21#[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))]
22const fn pack_payload(value: u64) -> u64 {
23    (value & 0x7f)
24        | ((value & 0x3f80) << 1)
25        | ((value & 0x001f_c000) << 2)
26        | ((value & 0x0fe0_0000) << 3)
27        | ((value & 0x0000_0007_f000_0000) << 4)
28        | ((value & 0x0000_03f8_0000_0000) << 5)
29        | ((value & 0x0001_fc00_0000_0000) << 6)
30        | ((value & 0x00fe_0000_0000_0000) << 7)
31}
32
33/// A variable-length integer.
34#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
35pub struct VarInt(pub i32);
36
37impl VarInt {
38    /// The maximum number of bytes a `VarInt` can be.
39    pub const MAX_SIZE: usize = 5;
40
41    /// Returns the exact number of bytes this `VarInt` will write when
42    /// [`WriteTo::write`] is called, assuming no error occurs.
43    #[must_use]
44    pub const fn written_size(val: i32) -> usize {
45        match val {
46            0 => 1,
47            n => (31 - n.leading_zeros() as usize) / 7 + 1,
48        }
49    }
50
51    /// Reads a `VarInt` from an async reader.
52    ///
53    /// # Errors
54    /// - If the `VarInt` is too long.
55    pub async fn read_async(read: &mut (impl AsyncRead + Unpin)) -> Result<i32, Error> {
56        let mut val = 0;
57        for i in 0..Self::MAX_SIZE {
58            let byte = read
59                .read_u8()
60                .await
61                .map_err(|err| Error::new(err.kind(), "VarInt"))?;
62            val |= (i32::from(byte) & 0x7F) << (i * 7);
63            if byte & 0x80 == 0 {
64                return Ok(val);
65            }
66        }
67        Err(Error::other("VarInt"))
68    }
69
70    /// Sets the `VarInt` in front of a `FrontVec`.
71    /// `varint_size` must match [`Self::written_size`] for this value
72    ///
73    /// # Panics
74    /// - If the size exceeds [`Self::MAX_SIZE`] or the reserved front space.
75    /// - in debug builds if `varint_size` does not match the encoded size
76    pub fn set_in_front(self, vec: &mut FrontVec, varint_size: usize) {
77        debug_assert_eq!(varint_size, Self::written_size(self.0));
78        let continuation = 0x0080_8080_8080_8080_u64 >> ((8 - varint_size) * 8);
79        let bytes = (pack_payload(u64::from(self.0 as u32)) | continuation).to_le_bytes();
80
81        vec.set_in_front(&bytes[..Self::MAX_SIZE][..varint_size]);
82    }
83}
84
85impl ReadFrom for VarInt {
86    fn read(read: &mut Cursor<&[u8]>) -> Result<Self, Error> {
87        let mut val = 0;
88        for i in 0..Self::MAX_SIZE {
89            let byte = u8::read(read)?;
90            val |= (i32::from(byte) & 0x7F) << (i * 7);
91
92            if byte & 0x80 == 0 {
93                return Ok(Self(val));
94            }
95        }
96
97        Err(Error::other("VarInt to long"))
98    }
99}
100
101impl WriteTo for VarInt {
102    #[expect(
103        clippy::inline_always,
104        reason = "codec benchmarks show these fixed-size writes need to inline into packet serializers"
105    )]
106    #[inline(always)]
107    fn write(&self, writer: &mut impl Write) -> Result<(), Error> {
108        let value = self.0 as u32;
109        if value < 0x80 {
110            return writer.write_all(&[value as u8]);
111        }
112
113        if value < 0x4000 {
114            return writer.write_all(&[value as u8 | 0x80, (value >> 7) as u8]);
115        }
116
117        if value < 0x20_0000 {
118            return writer.write_all(&[
119                value as u8 | 0x80,
120                (value >> 7) as u8 | 0x80,
121                (value >> 14) as u8,
122            ]);
123        }
124
125        let packed = pack_payload(u64::from(value));
126        if value < (1 << 28) {
127            writer.write_all(&(packed | 0x0080_8080).to_le_bytes()[..4])
128        } else {
129            writer.write_all(&(packed | 0x8080_8080).to_le_bytes()[..5])
130        }
131    }
132}
133
134impl From<usize> for VarInt {
135    fn from(value: usize) -> Self {
136        Self(value as _)
137    }
138}
139
140#[expect(
141    clippy::cast_sign_loss,
142    reason = "VarInt values used as lengths are always non-negative"
143)]
144impl From<VarInt> for usize {
145    fn from(value: VarInt) -> usize {
146        value.0 as _
147    }
148}
149
150impl From<i32> for VarInt {
151    fn from(value: i32) -> Self {
152        Self(value as _)
153    }
154}
155
156impl From<VarInt> for i32 {
157    fn from(value: VarInt) -> i32 {
158        value.0
159    }
160}
161
162#[cfg(test)]
163mod var_int_tests {
164    use super::*;
165    use std::io::Cursor;
166
167    #[test]
168    fn test_varint_read_write_negative() {
169        let val = VarInt(-1);
170        let mut buf = Vec::new();
171        val.write(&mut buf).expect("write failed");
172
173        // Expected VarInt encoding for -1 (0xFFFFFFFF)
174        assert_eq!(buf, vec![0xff, 0xff, 0xff, 0xff, 0x0f]);
175
176        let mut cursor = Cursor::new(buf.as_slice());
177        let read_val = VarInt::read(&mut cursor).expect("read failed");
178        assert_eq!(read_val, val);
179    }
180}
181
182/// A variable-length long integer.
183#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
184pub struct VarLong(pub i64);
185
186impl VarLong {
187    /// The maximum number of bytes a `VarLong` can be.
188    pub const MAX_SIZE: usize = 10;
189}
190
191impl ReadFrom for VarLong {
192    fn read(read: &mut Cursor<&[u8]>) -> Result<Self, Error> {
193        let mut val = 0i64;
194        for i in 0..Self::MAX_SIZE {
195            let byte = u8::read(read)?;
196            val |= (i64::from(byte) & 0x7F) << (i * 7);
197            if byte & 0x80 == 0 {
198                return Ok(Self(val));
199            }
200        }
201        Err(Error::other("VarLong too long"))
202    }
203}
204
205impl WriteTo for VarLong {
206    #[expect(
207        clippy::inline_always,
208        reason = "codec benchmarks show these fixed-size writes need to inline into packet serializers"
209    )]
210    #[inline(always)]
211    fn write(&self, writer: &mut impl Write) -> Result<(), Error> {
212        let value = self.0 as u64;
213        if value < 0x80 {
214            return writer.write_all(&[value as u8]);
215        }
216
217        if let Ok(value) = u32::try_from(value) {
218            return VarInt(value as i32).write(writer);
219        }
220
221        let packed = pack_payload(value);
222        if value >= (1 << 56) {
223            let low = (packed | 0x8080_8080_8080_8080).to_le_bytes();
224            let bytes = [
225                low[0],
226                low[1],
227                low[2],
228                low[3],
229                low[4],
230                low[5],
231                low[6],
232                low[7],
233                (value >> 56) as u8,
234                1,
235            ];
236
237            return if self.0 >= 0 {
238                writer.write_all(&bytes[..9])
239            } else {
240                writer.write_all(&bytes)
241            };
242        }
243
244        let size = (64 - value.leading_zeros() as usize).div_ceil(7);
245        let continuation = 0x0080_8080_8080_8080_u64 >> ((8 - size) * 8);
246        let bytes = (packed | continuation).to_le_bytes();
247
248        match size {
249            5 => writer.write_all(&bytes[..5]),
250            6 => writer.write_all(&bytes[..6]),
251            7 => writer.write_all(&bytes[..7]),
252            _ => writer.write_all(&bytes),
253        }
254    }
255}
256
257impl From<i64> for VarLong {
258    fn from(value: i64) -> Self {
259        Self(value)
260    }
261}
262
263impl From<VarLong> for i64 {
264    fn from(value: VarLong) -> i64 {
265        value.0
266    }
267}
268
269#[cfg(test)]
270mod var_long_tests {
271    use super::*;
272    use std::io::Cursor;
273
274    #[test]
275    fn test_varlong_read_write() {
276        let test_values = vec![
277            0i64,
278            1i64,
279            127i64,
280            128i64,
281            255i64,
282            2_147_483_647_i64,
283            9_223_372_036_854_775_807_i64,
284            -1i64,
285            -2_147_483_648_i64,
286        ];
287
288        for val in test_values {
289            let var_long = VarLong(val);
290            let mut buf = Vec::new();
291            var_long.write(&mut buf).expect("write failed");
292
293            let mut cursor = Cursor::new(buf.as_slice());
294            let read_val = VarLong::read(&mut cursor).expect("read failed");
295            assert_eq!(read_val, var_long, "Failed for value {val}");
296        }
297    }
298}
299
300/// A variable-length unsigned integer.
301#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
302pub struct VarUint(pub u32);
303
304impl VarUint {
305    const MAX_SIZE: usize = 5;
306
307    /// Returns the exact number of bytes this `VarUInt` will write when
308    /// [`WriteTo::write`] is called, assuming no error occurs.
309    #[must_use]
310    pub const fn written_size(self) -> usize {
311        (32 - self.0.leading_zeros() as usize).max(1).div_ceil(7)
312    }
313
314    /// Writes a `VarUint` to a writer.
315    ///
316    /// # Errors
317    /// - If the writer fails to write.
318    pub fn write(self, writer: &mut impl Write) -> Result<(), Error> {
319        let mut val = self.0;
320        loop {
321            let mut byte = (val & 0x7F) as u8;
322            val >>= 7;
323            if val != 0 {
324                byte |= 0x80;
325            }
326            byte.write(writer)?;
327            if val == 0 {
328                break;
329            }
330        }
331        Ok(())
332    }
333
334    /// Reads a `VarUint` from a cursor.
335    ///
336    /// # Errors
337    /// - If the `VarUint` is too long.
338    pub fn read(read: &mut Cursor<&[u8]>) -> Result<u32, Error> {
339        let mut val = 0;
340        for i in 0..Self::MAX_SIZE {
341            let byte = u8::read(read)?;
342            val |= (u32::from(byte) & 0x7F) << (i * 7);
343            if byte & 0x80 == 0 {
344                return Ok(val);
345            }
346        }
347        Err(Error::other("Malformed VarUint"))
348    }
349}
steel_utils/codec/variable_integer.rs

steel_utils/codec/
variable_integer.rs
