Skip to main content

orengine_utils/
varint.rs

1//! Variable-length integer (varint) encoding and decoding.
2//!
3//! This module provides efficient serialization routines for integer types.
4//!
5//! Unsigned integers are encoded using a little-endian base-128 (LEB128-like)
6//! variable-length representation. Small values occupy fewer bytes than larger
7//! values.
8//!
9//! Signed integers are encoded using `ZigZag encoding` followed by unsigned
10//! varint encoding. This efficiently stores small negative numbers.
11//!
12//! # Supported integer types
13//!
14//! - Unsigned: `u8`, `u16`, `u32`, `u64`, `u128`
15//! - Signed: `i8`, `i16`, `i32`, `i64`, `i128`
16//!
17//! Convenience extension traits are also provided for all `Read` and `Write`
18//! implementations.
19//!
20//! # Examples
21//!
22//! ```rust
23//! use std::io::Cursor;
24//! use orengine_utils::varint::{ReadVarInt, WriteVarInt};
25//!
26//! let mut buf = Vec::new();
27//!
28//! buf.write_varint(123u32).unwrap();
29//! buf.write_varint(-42i32).unwrap();
30//!
31//! let mut cursor = Cursor::new(buf);
32//!
33//! let a: u32 = cursor.read_varint().unwrap();
34//! let b: i32 = cursor.read_varint().unwrap();
35//!
36//! assert_eq!(a, 123);
37//! assert_eq!(b, -42);
38//! ```
39use std::io::{self, Read, Write};
40
41#[inline]
42pub(crate) fn write_u8_to<W: Write + ?Sized>(n: u8, writer: &mut W) -> io::Result<usize> {
43    writer.write_all(&[n])?;
44
45    Ok(1)
46}
47
48#[inline]
49pub(crate) fn read_u8_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<u8> {
50    let mut result = [0];
51
52    reader.read_exact(&mut result)?;
53
54    Ok(u8::from_le_bytes(result))
55}
56
57#[inline]
58pub(crate) fn write_i8_to<W: Write + ?Sized>(n: i8, writer: &mut W) -> io::Result<usize> {
59    writer.write_all(&n.to_le_bytes())?;
60
61    Ok(1)
62}
63
64#[inline]
65pub(crate) fn read_i8_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<i8> {
66    let mut result = [0];
67
68    reader.read_exact(&mut result)?;
69
70    Ok(result[0].cast_signed())
71}
72
73#[inline]
74#[allow(clippy::cast_possible_truncation, reason = "False positive.")]
75pub(crate) fn write_u128_to<W: Write + ?Sized>(mut n: u128, writer: &mut W) -> io::Result<usize> {
76    let mut buf = [0u8; 19];
77    let mut idx = 0;
78
79    while n >= 0x80 {
80        buf[idx] = (n as u8) | 0x80;
81        n >>= 7;
82        idx += 1;
83    }
84
85    buf[idx] = n as u8;
86    idx += 1;
87
88    writer.write_all(&buf[..idx])?;
89    Ok(idx)
90}
91
92#[inline]
93#[allow(clippy::cast_lossless, reason = "False positive.")]
94pub(crate) fn read_u128_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<u128> {
95    let mut result = 0u128;
96    let mut shift = 0;
97    let mut buf = [0u8; 1];
98
99    loop {
100        reader.read_exact(&mut buf)?;
101        let byte = buf[0];
102
103        if shift >= 126 && (byte & 0x80) != 0 {
104            return Err(io::Error::new(
105                io::ErrorKind::InvalidData,
106                "varint too large for u128",
107            ));
108        }
109
110        result |= ((byte & 0x7f) as u128) << shift;
111
112        if byte & 0x80 == 0 {
113            break;
114        }
115
116        shift += 7;
117    }
118
119    Ok(result)
120}
121
122macro_rules! impl_unsigned {
123    ($ty:ty) => {
124        paste::paste! {
125            #[inline]
126            #[allow(clippy::cast_lossless, reason = "It is generated code.")]
127            #[allow(clippy::cast_possible_truncation, reason = "It is generated code.")]
128            pub(crate) fn [<write_ $ty _to>]<W: Write + ?Sized>(
129                value: $ty,
130                writer: &mut W,
131            ) -> io::Result<usize> {
132                write_u128_to(value as u128, writer)
133            }
134
135            #[inline]
136            #[allow(clippy::cast_lossless, reason = "It is generated code.")]
137            #[allow(clippy::cast_possible_truncation, reason = "It is generated code.")]
138            pub(crate) fn [<read_ $ty _from>]<R: Read + ?Sized>(
139                reader: &mut R,
140            ) -> io::Result<$ty> {
141                let value = read_u128_from(reader)?;
142
143                if value > <$ty>::MAX as u128 {
144                    return Err(io::Error::new(
145                        io::ErrorKind::InvalidData,
146                        concat!("varint too large for ", stringify!($ty)),
147                    ));
148                }
149
150                Ok(value as $ty)
151            }
152
153            pub trait [<Write $ty:camel>]: Write {
154                fn [<write_ $ty>](
155                    &mut self,
156                    value: $ty,
157                ) -> io::Result<usize> {
158                    [<write_ $ty _to>](value, self)
159                }
160            }
161
162            impl<T: Write> [<Write $ty:camel>] for T {}
163
164            pub trait [<Read $ty:camel>]: Read {
165                fn [<read_ $ty>](
166                    &mut self,
167                ) -> io::Result<$ty> {
168                    [<read_ $ty _from>](self)
169                }
170            }
171
172            impl<T: Read> [<Read $ty:camel>] for T {}
173        }
174    };
175}
176
177impl_unsigned!(u16);
178impl_unsigned!(u32);
179impl_unsigned!(u64);
180
181#[inline]
182#[allow(clippy::cast_sign_loss, reason = "It will be restored")]
183fn zigzag_encode(v: i128) -> u128 {
184    ((v << 1) ^ (v >> 127)) as u128
185}
186
187#[inline]
188#[allow(clippy::cast_possible_wrap, reason = "It will be restored")]
189fn zigzag_decode(v: u128) -> i128 {
190    ((v >> 1) as i128) ^ (-((v & 1) as i128))
191}
192
193#[inline]
194pub(crate) fn write_i128_to<W: Write + ?Sized>(n: i128, writer: &mut W) -> io::Result<usize> {
195    write_u128_to(zigzag_encode(n), writer)
196}
197
198#[inline]
199pub(crate) fn read_i128_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<i128> {
200    Ok(zigzag_decode(read_u128_from(reader)?))
201}
202
203macro_rules! impl_signed {
204    ($ty:ty) => {
205        paste::paste! {
206            #[inline]
207            #[allow(clippy::cast_lossless, reason = "It is generated code.")]
208            #[allow(clippy::cast_possible_truncation, reason = "It is generated code.")]
209            pub(crate) fn [<write_ $ty _to>]<W: Write + ?Sized>(
210                value: $ty,
211                writer: &mut W,
212            ) -> io::Result<usize> {
213                write_u128_to(zigzag_encode(value as i128), writer)
214            }
215
216            #[inline]
217            #[allow(clippy::cast_lossless, reason = "It is generated code.")]
218            #[allow(clippy::cast_possible_truncation, reason = "It is generated code.")]
219            pub(crate) fn [<read_ $ty _from>]<R: Read + ?Sized>(
220                reader: &mut R,
221            ) -> io::Result<$ty> {
222                let value = zigzag_decode(read_u128_from(reader)?);
223
224                if value < <$ty>::MIN as i128 || value > <$ty>::MAX as i128 {
225                    return Err(io::Error::new(
226                        io::ErrorKind::InvalidData,
227                        concat!("varint too large for ", stringify!($ty)),
228                    ));
229                }
230
231                Ok(value as $ty)
232            }
233        }
234    };
235}
236
237impl_signed!(i16);
238impl_signed!(i32);
239impl_signed!(i64);
240
241/// A type that can be encoded and decoded as a variable-length integer.
242///
243/// This trait is implemented for all primitive integer types supported by this
244/// module.
245///
246/// It is primarily intended for generic serialization code.
247///
248/// Most users should prefer using the [`ReadVarInt`] and [`WriteVarInt`]
249/// extension traits.
250pub trait VarInt: Sized {
251    /// Writes the number as little-endian base-128 variable-length representation to the writer.
252    fn write_as_varint_to<W: Write + ?Sized>(self, writer: &mut W) -> io::Result<usize>;
253    /// Reads the number as little-endian base-128 variable-length representation from the reader.
254    fn read_varint_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<Self>;
255}
256
257macro_rules! impl_varint {
258    ($($ty:ty),*) => {
259        $(
260            paste::paste! {
261                impl VarInt for $ty {
262                    #[inline]
263                    fn write_as_varint_to<W: Write + ?Sized>(self, writer: &mut W) -> io::Result<usize> {
264                        [<write_ $ty _to>](self, writer)
265                    }
266
267                    #[inline]
268                    fn read_varint_from<R: Read + ?Sized>(reader: &mut R) -> io::Result<Self> {
269                        [<read_ $ty _from>](reader)
270                    }
271                }
272            }
273        )*
274    };
275}
276
277impl_varint!(u8, u16, u32, u64, u128, i8, i16, i32, i64, i128);
278
279/// Extension trait for writing variable-length integers.
280///
281/// Implemented automatically for every type implementing [`Write`].
282///
283/// # Example
284///
285/// ```rust
286/// use std::io::Cursor;
287/// use orengine_utils::varint::WriteVarInt;
288///
289/// let mut writer = Cursor::new(Vec::new());
290///
291/// writer.write_varint(42u64).unwrap();
292/// ```
293pub trait WriteVarInt: Write {
294    /// Writes [`VarInt`] to the writer.
295    ///
296    /// # Example
297    ///
298    /// ```rust
299    /// use std::io::Cursor;
300    /// use orengine_utils::varint::WriteVarInt;
301    ///
302    /// let mut writer = Cursor::new(Vec::new());
303    ///
304    /// writer.write_varint(42u64).unwrap();
305    /// ```
306    fn write_varint<T: VarInt>(&mut self, value: T) -> io::Result<usize> {
307        value.write_as_varint_to(self)
308    }
309}
310
311/// Extension trait for reading variable-length integers.
312///
313/// Implemented automatically for every type implementing [`Read`].
314///
315/// # Example
316///
317/// ```rust
318/// use std::io::Cursor;
319/// use orengine_utils::varint::{ReadVarInt, WriteVarInt};
320///
321/// let mut data = Vec::new();
322/// data.write_varint(500u32).unwrap();
323///
324/// let mut reader = Cursor::new(data);
325///
326/// let value: u32 = reader.read_varint().unwrap();
327///
328/// assert_eq!(value, 500);
329/// ```
330pub trait ReadVarInt: Read {
331    /// Reads a [`VarInt`] from the reader.
332    ///
333    /// # Example
334    ///
335    /// ```rust
336    /// use std::io::Cursor;
337    /// use orengine_utils::varint::{ReadVarInt, WriteVarInt};
338    ///
339    /// let mut data = Vec::new();
340    /// data.write_varint(500u32).unwrap();
341    ///
342    /// let mut reader = Cursor::new(data);
343    ///
344    /// let value: u32 = reader.read_varint().unwrap();
345    ///
346    /// assert_eq!(value, 500);
347    /// ```
348    fn read_varint<T: VarInt>(&mut self) -> io::Result<T> {
349        T::read_varint_from(self)
350    }
351}
352
353impl<W: Write> WriteVarInt for W {}
354impl<R: Read> ReadVarInt for R {}
355
356#[cfg(test)]
357mod tests {
358    use super::*;
359    use std::io::Cursor;
360
361    fn roundtrip<T>(value: T)
362    where
363        T: VarInt + Copy + PartialEq + std::fmt::Debug,
364    {
365        let mut buf = Vec::new();
366
367        value.write_as_varint_to(&mut buf).unwrap();
368
369        let decoded = T::read_varint_from(&mut Cursor::new(buf)).unwrap();
370
371        assert_eq!(value, decoded);
372    }
373
374    #[test]
375    fn roundtrip_unsigned() {
376        roundtrip(0u8);
377        roundtrip(1u8);
378        roundtrip(u8::MAX);
379
380        roundtrip(0u16);
381        roundtrip(127u16);
382        roundtrip(128u16);
383        roundtrip(u16::MAX);
384
385        roundtrip(0u32);
386        roundtrip(u32::MAX);
387
388        roundtrip(0u64);
389        roundtrip(u64::MAX);
390
391        roundtrip(0u128);
392        roundtrip(u128::MAX);
393    }
394
395    #[test]
396    fn roundtrip_signed() {
397        roundtrip(0i8);
398        roundtrip(-1i8);
399        roundtrip(i8::MIN);
400        roundtrip(i8::MAX);
401
402        roundtrip(0i16);
403        roundtrip(-1i16);
404        roundtrip(i16::MIN);
405        roundtrip(i16::MAX);
406
407        roundtrip(0i32);
408        roundtrip(i32::MIN);
409        roundtrip(i32::MAX);
410
411        roundtrip(0i64);
412        roundtrip(i64::MIN);
413        roundtrip(i64::MAX);
414
415        roundtrip(0i128);
416        roundtrip(i128::MIN);
417        roundtrip(i128::MAX);
418    }
419
420    #[test]
421    fn encoding_size() {
422        let mut buf = Vec::new();
423
424        assert_eq!(write_u128_to(0, &mut buf).unwrap(), 1);
425
426        buf.clear();
427        assert_eq!(write_u128_to(127, &mut buf).unwrap(), 1);
428
429        buf.clear();
430        assert_eq!(write_u128_to(128, &mut buf).unwrap(), 2);
431
432        buf.clear();
433        assert_eq!(write_u128_to(16383, &mut buf).unwrap(), 2);
434
435        buf.clear();
436        assert_eq!(write_u128_to(16384, &mut buf).unwrap(), 3);
437    }
438
439    #[test]
440    fn malformed_varint_is_rejected() {
441        let bytes = [0xff; 19];
442
443        let err = read_u128_from(&mut Cursor::new(bytes)).unwrap_err();
444
445        assert_eq!(err.kind(), io::ErrorKind::InvalidData);
446    }
447
448    #[test]
449    fn overflow_for_smaller_type() {
450        let mut buf = Vec::new();
451
452        write_u32_to(70000, &mut buf).unwrap();
453
454        let err = read_u16_from(&mut Cursor::new(buf)).unwrap_err();
455
456        assert_eq!(err.kind(), io::ErrorKind::InvalidData);
457    }
458
459    #[test]
460    fn generic_traits_work() {
461        let mut buf = Vec::new();
462
463        buf.write_varint(12345u32).unwrap();
464        buf.write_varint(-567i32).unwrap();
465
466        let mut cursor = Cursor::new(buf);
467
468        let a: u32 = cursor.read_varint().unwrap();
469        let b: i32 = cursor.read_varint().unwrap();
470
471        assert_eq!(a, 12345);
472        assert_eq!(b, -567);
473    }
474}