Skip to main content

steel_utils/codec/
var_int.rs

1use std::io::{Cursor, Error, Write};
2
3use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
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    /// Writes a `VarInt` to an async writer.
48    ///
49    /// # Errors
50    /// - If the writer fails to write.
51    pub async fn write_async(self, write: &mut (impl AsyncWrite + Unpin)) -> Result<(), Error> {
52        let mut val = self.0 as u32;
53        loop {
54            let b: u8 = (val as u8) & 0b0111_1111;
55            val >>= 7;
56            write
57                .write_u8(if val == 0 { b } else { b | 0b1000_0000 })
58                .await?;
59            if val == 0 {
60                break;
61            }
62        }
63        Ok(())
64    }
65
66    // We could just get the written size in place,
67    // but in our use case its already calculated
68    /// Sets the `VarInt` in front of a `FrontVec`.
69    ///
70    /// # Panics
71    /// - If the `VarInt` fails to write to the buffer.
72    pub fn set_in_front(self, vec: &mut FrontVec, varint_size: usize) {
73        // No heap allocation :)
74        let mut buf = [0; Self::MAX_SIZE];
75        self.write(&mut Cursor::new(&mut buf[..]))
76            .expect("writing to a buffer should not fail");
77        vec.set_in_front(&buf[..varint_size]);
78    }
79}
80
81impl ReadFrom for VarInt {
82    fn read(read: &mut Cursor<&[u8]>) -> Result<Self, Error> {
83        let mut val = 0;
84        for i in 0..Self::MAX_SIZE {
85            let byte = u8::read(read)?;
86            val |= (i32::from(byte) & 0x7F) << (i * 7);
87            if byte & 0x80 == 0 {
88                return Ok(Self(val));
89            }
90        }
91        Err(Error::other("VarInt to long"))
92    }
93}
94
95impl WriteTo for VarInt {
96    fn write(&self, writer: &mut impl Write) -> Result<(), Error> {
97        let mut val = self.0 as u32;
98        loop {
99            let b: u8 = val as u8 & 0x7F;
100            val >>= 7;
101            if val == 0 {
102                b.write(writer)?;
103                break;
104            }
105            (b | 0x80).write(writer)?;
106        }
107        Ok(())
108    }
109}
110
111impl From<usize> for VarInt {
112    fn from(value: usize) -> Self {
113        Self(value as _)
114    }
115}
116
117#[expect(
118    clippy::cast_sign_loss,
119    reason = "VarInt values used as lengths are always non-negative"
120)]
121impl From<VarInt> for usize {
122    fn from(value: VarInt) -> usize {
123        value.0 as _
124    }
125}
126
127impl From<i32> for VarInt {
128    fn from(value: i32) -> Self {
129        Self(value as _)
130    }
131}
132
133impl From<VarInt> for i32 {
134    fn from(value: VarInt) -> i32 {
135        value.0
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    use std::io::Cursor;
143
144    #[test]
145    fn test_varint_read_write_negative() {
146        let val = VarInt(-1);
147        let mut buf = Vec::new();
148        val.write(&mut buf).expect("write failed");
149
150        // Expected VarInt encoding for -1 (0xFFFFFFFF)
151        assert_eq!(buf, vec![0xff, 0xff, 0xff, 0xff, 0x0f]);
152
153        let mut cursor = Cursor::new(buf.as_slice());
154        let read_val = VarInt::read(&mut cursor).expect("read failed");
155        assert_eq!(read_val, val);
156    }
157}