ax_codec_core/
varint.rs

1use crate::{BufferReader, BufferWriter, DecodeError, EncodeError};
2
3#[inline(always)]
4#[cfg_attr(feature = "simd", allow(dead_code))]
5pub(crate) fn encode_uvarint_slow<W: BufferWriter>(
6    mut value: u64,
7    writer: &mut W,
8) -> Result<(), EncodeError> {
9    loop {
10        let mut byte = (value & 0x7f) as u8;
11        value >>= 7;
12        if value != 0 {
13            byte |= 0x80;
14        }
15        writer.write_all(&[byte])?;
16        if value == 0 {
17            break;
18        }
19    }
20    Ok(())
21}
22
23#[inline(always)]
24pub(crate) fn decode_uvarint_slow<'a, R: BufferReader<'a>>(
25    reader: &mut R,
26) -> Result<u64, DecodeError> {
27    let mut result = 0u64;
28    let mut shift = 0;
29    let mut count = 0;
30    const MAX_VARINT_BYTES: usize = 10; // u64 max
31    while count < MAX_VARINT_BYTES {
32        count += 1;
33        let byte = reader.next().ok_or(DecodeError::UnexpectedEOF)?;
34        let val = (byte & 0x7f) as u64;
35        if shift == 63 && (byte & 0x7f) > 1 {
36            return Err(DecodeError::InvalidVarint);
37        }
38        result |= val.checked_shl(shift).ok_or(DecodeError::InvalidVarint)?;
39        if (byte & 0x80) == 0 {
40            return Ok(result);
41        }
42        shift += 7;
43        if shift > 63 {
44            return Err(DecodeError::InvalidVarint);
45        }
46    }
47    Err(DecodeError::InvalidVarint)
48}
49
50#[inline]
51pub fn encode_uvarint<W: BufferWriter>(value: u64, writer: &mut W) -> Result<(), EncodeError> {
52    #[cfg(feature = "simd")]
53    {
54        crate::simd::encode_uvarint_fast(value, writer)
55    }
56    #[cfg(not(feature = "simd"))]
57    {
58        encode_uvarint_slow(value, writer)
59    }
60}
61
62#[inline]
63pub fn decode_uvarint<'a, R: BufferReader<'a>>(reader: &mut R) -> Result<u64, DecodeError> {
64    #[cfg(feature = "simd")]
65    {
66        crate::simd::decode_uvarint_fast(reader)
67    }
68    #[cfg(not(feature = "simd"))]
69    {
70        decode_uvarint_slow(reader)
71    }
72}
73
74#[inline]
75pub fn encode_svarint<W: BufferWriter>(value: i64, writer: &mut W) -> Result<(), EncodeError> {
76    let zigzag = ((value << 1) ^ (value >> 63)) as u64;
77    encode_uvarint(zigzag, writer)
78}
79
80#[inline]
81pub fn decode_svarint<'a, R: BufferReader<'a>>(reader: &mut R) -> Result<i64, DecodeError> {
82    let zigzag = decode_uvarint(reader)?;
83    Ok(((zigzag >> 1) as i64) ^ -((zigzag & 1) as i64))
84}
85
86#[inline]
87pub fn encode_uvarint128<W: BufferWriter>(
88    mut value: u128,
89    writer: &mut W,
90) -> Result<(), EncodeError> {
91    loop {
92        let mut byte = (value & 0x7f) as u8;
93        value >>= 7;
94        if value != 0 {
95            byte |= 0x80;
96        }
97        writer.write_all(&[byte])?;
98        if value == 0 {
99            break;
100        }
101    }
102    Ok(())
103}
104
105#[inline]
106pub fn decode_uvarint128<'a, R: BufferReader<'a>>(reader: &mut R) -> Result<u128, DecodeError> {
107    let mut result: u128 = 0;
108    let mut shift: u32 = 0;
109    loop {
110        let byte = reader.next().ok_or(DecodeError::UnexpectedEOF)?;
111        let val = (byte & 0x7f) as u128;
112        if shift == 127 && (byte & 0x7f) > 1 {
113            return Err(DecodeError::InvalidVarint);
114        }
115        result |= val.checked_shl(shift).ok_or(DecodeError::InvalidVarint)?;
116        if (byte & 0x80) == 0 {
117            return Ok(result);
118        }
119        shift += 7;
120        if shift > 127 {
121            return Err(DecodeError::InvalidVarint);
122        }
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use crate::buffer::{SliceReader, VecWriter};
130
131    #[test]
132    fn uvarint_roundtrip() {
133        let values = [0u64, 1, 127, 128, 255, 256, 16383, 16384, u64::MAX];
134        for &val in &values {
135            let mut w = VecWriter::new();
136            encode_uvarint(val, &mut w).unwrap();
137            let mut r = SliceReader::new(w.as_slice());
138            let decoded = decode_uvarint(&mut r).unwrap();
139            assert_eq!(val, decoded, "mismatch for {}", val);
140        }
141    }
142
143    #[test]
144    fn svarint_roundtrip() {
145        let values = [0i64, 1, -1, 2, -2, i64::MAX, i64::MIN];
146        for &val in &values {
147            let mut w = VecWriter::new();
148            encode_svarint(val, &mut w).unwrap();
149            let mut r = SliceReader::new(w.as_slice());
150            let decoded = decode_svarint(&mut r).unwrap();
151            assert_eq!(val, decoded, "mismatch for {}", val);
152        }
153    }
154
155    #[test]
156    fn uvarint_known_encodings() {
157        let mut w = VecWriter::new();
158        encode_uvarint(1, &mut w).unwrap();
159        assert_eq!(w.as_slice(), &[0x01]);
160
161        let mut w = VecWriter::new();
162        encode_uvarint(127, &mut w).unwrap();
163        assert_eq!(w.as_slice(), &[0x7f]);
164
165        let mut w = VecWriter::new();
166        encode_uvarint(128, &mut w).unwrap();
167        assert_eq!(w.as_slice(), &[0x80, 0x01]);
168    }
169
170    #[test]
171    fn invalid_varint_too_long() {
172        let buf = [0xff; 10];
173        let mut r = SliceReader::new(&buf);
174        assert!(matches!(
175            decode_uvarint(&mut r),
176            Err(DecodeError::InvalidVarint)
177        ));
178    }
179}
ax_codec_core/varint.rs

ax_codec_core/
varint.rs
