steel_utils/codec/
variable_integer.rs

1use std::io::{Cursor, Error, Write};
2
3use tokio::io::{AsyncRead, AsyncReadExt};
4
5use crate::{
6    FrontVec,
7    serial::{ReadFrom, WriteTo},
8};
9
10/// A variable-length integer.
11#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
12pub struct VarInt(pub i32);
13
14impl VarInt {
15    /// The maximum number of bytes a `VarInt` can be.
16    pub const MAX_SIZE: usize = 5;
17
18    /// Returns the exact number of bytes this `VarInt` will write when
19    /// [`WriteTo::write`] is called, assuming no error occurs.
20    #[must_use]
21    pub const fn written_size(val: i32) -> usize {
22        match val {
23            0 => 1,
24            n => (31 - n.leading_zeros() as usize) / 7 + 1,
25        }
26    }
27
28    /// Reads a `VarInt` from an async reader.
29    ///
30    /// # Errors
31    /// - If the `VarInt` is too long.
32    pub async fn read_async(read: &mut (impl AsyncRead + Unpin)) -> Result<i32, Error> {
33        let mut val = 0;
34        for i in 0..Self::MAX_SIZE {
35            let byte = read
36                .read_u8()
37                .await
38                .map_err(|err| Error::new(err.kind(), "VarInt"))?;
39            val |= (i32::from(byte) & 0x7F) << (i * 7);
40            if byte & 0x80 == 0 {
41                return Ok(val);
42            }
43        }
44        Err(Error::other("VarInt"))
45    }
46
47    // We could just get the written size in place,
48    // but in our use case its already calculated
49    /// Sets the `VarInt` in front of a `FrontVec`.
50    ///
51    /// # Panics
52    /// - If the `VarInt` fails to write to the buffer.
53    pub fn set_in_front(self, vec: &mut FrontVec, varint_size: usize) {
54        // No heap allocation :)
55        let mut buf = [0; Self::MAX_SIZE];
56        self.write(&mut Cursor::new(&mut buf[..]))
57            .expect("writing to a buffer should not fail");
58        vec.set_in_front(&buf[..varint_size]);
59    }
60}
61
62impl ReadFrom for VarInt {
63    fn read(read: &mut Cursor<&[u8]>) -> Result<Self, Error> {
64        let mut val = 0;
65        for i in 0..Self::MAX_SIZE {
66            let byte = u8::read(read)?;
67            val |= (i32::from(byte) & 0x7F) << (i * 7);
68            if byte & 0x80 == 0 {
69                return Ok(Self(val));
70            }
71        }
72        Err(Error::other("VarInt to long"))
73    }
74}
75
76impl WriteTo for VarInt {
77    fn write(&self, writer: &mut impl Write) -> Result<(), Error> {
78        let mut val = self.0 as u32;
79        loop {
80            let b: u8 = val as u8 & 0x7F;
81            val >>= 7;
82            if val == 0 {
83                b.write(writer)?;
84                break;
85            }
86            (b | 0x80).write(writer)?;
87        }
88        Ok(())
89    }
90}
91
92impl From<usize> for VarInt {
93    fn from(value: usize) -> Self {
94        Self(value as _)
95    }
96}
97
98#[expect(
99    clippy::cast_sign_loss,
100    reason = "VarInt values used as lengths are always non-negative"
101)]
102impl From<VarInt> for usize {
103    fn from(value: VarInt) -> usize {
104        value.0 as _
105    }
106}
107
108impl From<i32> for VarInt {
109    fn from(value: i32) -> Self {
110        Self(value as _)
111    }
112}
113
114impl From<VarInt> for i32 {
115    fn from(value: VarInt) -> i32 {
116        value.0
117    }
118}
119
120#[cfg(test)]
121mod var_int_tests {
122    use super::*;
123    use std::io::Cursor;
124
125    #[test]
126    fn test_varint_read_write_negative() {
127        let val = VarInt(-1);
128        let mut buf = Vec::new();
129        val.write(&mut buf).expect("write failed");
130
131        // Expected VarInt encoding for -1 (0xFFFFFFFF)
132        assert_eq!(buf, vec![0xff, 0xff, 0xff, 0xff, 0x0f]);
133
134        let mut cursor = Cursor::new(buf.as_slice());
135        let read_val = VarInt::read(&mut cursor).expect("read failed");
136        assert_eq!(read_val, val);
137    }
138}
139/// A variable-length long integer.
140#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
141pub struct VarLong(pub i64);
142
143impl VarLong {
144    /// The maximum number of bytes a `VarLong` can be.
145    pub const MAX_SIZE: usize = 10;
146}
147
148impl ReadFrom for VarLong {
149    fn read(read: &mut Cursor<&[u8]>) -> Result<Self, Error> {
150        let mut val = 0i64;
151        for i in 0..Self::MAX_SIZE {
152            let byte = u8::read(read)?;
153            val |= (i64::from(byte) & 0x7F) << (i * 7);
154            if byte & 0x80 == 0 {
155                return Ok(Self(val));
156            }
157        }
158        Err(Error::other("VarLong too long"))
159    }
160}
161
162impl WriteTo for VarLong {
163    fn write(&self, writer: &mut impl Write) -> Result<(), Error> {
164        let mut val = self.0 as u64;
165        loop {
166            let b: u8 = val as u8 & 0x7F;
167            val >>= 7;
168            if val == 0 {
169                b.write(writer)?;
170                break;
171            }
172            (b | 0x80).write(writer)?;
173        }
174        Ok(())
175    }
176}
177
178impl From<i64> for VarLong {
179    fn from(value: i64) -> Self {
180        Self(value)
181    }
182}
183
184impl From<VarLong> for i64 {
185    fn from(value: VarLong) -> i64 {
186        value.0
187    }
188}
189
190#[cfg(test)]
191mod var_long_tests {
192    use super::*;
193    use std::io::Cursor;
194
195    #[test]
196    fn test_varlong_read_write() {
197        let test_values = vec![
198            0i64,
199            1i64,
200            127i64,
201            128i64,
202            255i64,
203            2_147_483_647_i64,
204            9_223_372_036_854_775_807_i64,
205            -1i64,
206            -2_147_483_648_i64,
207        ];
208
209        for val in test_values {
210            let var_long = VarLong(val);
211            let mut buf = Vec::new();
212            var_long.write(&mut buf).expect("write failed");
213
214            let mut cursor = Cursor::new(buf.as_slice());
215            let read_val = VarLong::read(&mut cursor).expect("read failed");
216            assert_eq!(read_val, var_long, "Failed for value {val}");
217        }
218    }
219}
220/// A variable-length unsigned integer.
221#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
222pub struct VarUint(pub u32);
223
224impl VarUint {
225    const MAX_SIZE: usize = 5;
226
227    /// Returns the exact number of bytes this `VarUInt` will write when
228    /// [`WriteTo::write`] is called, assuming no error occurs.
229    #[must_use]
230    pub const fn written_size(self) -> usize {
231        (32 - self.0.leading_zeros() as usize).max(1).div_ceil(7)
232    }
233
234    /// Writes a `VarUint` to a writer.
235    ///
236    /// # Errors
237    /// - If the writer fails to write.
238    pub fn write(self, writer: &mut impl Write) -> Result<(), Error> {
239        let mut val = self.0;
240        loop {
241            let mut byte = (val & 0x7F) as u8;
242            val >>= 7;
243            if val != 0 {
244                byte |= 0x80;
245            }
246            byte.write(writer)?;
247            if val == 0 {
248                break;
249            }
250        }
251        Ok(())
252    }
253
254    /// Reads a `VarUint` from a cursor.
255    ///
256    /// # Errors
257    /// - If the `VarUint` is too long.
258    pub fn read(read: &mut Cursor<&[u8]>) -> Result<u32, Error> {
259        let mut val = 0;
260        for i in 0..Self::MAX_SIZE {
261            let byte = u8::read(read)?;
262            val |= (u32::from(byte) & 0x7F) << (i * 7);
263            if byte & 0x80 == 0 {
264                return Ok(val);
265            }
266        }
267        Err(Error::other("Malformed VarUint"))
268    }
269}
steel_utils/codec/variable_integer.rs

steel_utils/codec/
variable_integer.rs
