Skip to main content

pylon_replication/
varint.rs

1//! LEB128 varints and zigzag signed integers.
2
3/// Append `v` as an unsigned LEB128 varint (1 to 10 bytes).
4pub fn write_u64(out: &mut Vec<u8>, mut v: u64) {
5    loop {
6        let byte = (v & 0x7f) as u8;
7        v >>= 7;
8        if v == 0 {
9            out.push(byte);
10            return;
11        }
12        out.push(byte | 0x80);
13    }
14}
15
16/// Bytes `write_u64` uses for `v`.
17pub fn len_u64(v: u64) -> usize {
18    (64 - (v | 1).leading_zeros() as usize).div_ceil(7)
19}
20
21/// Read an unsigned LEB128 varint and advance `bytes` past it. None when
22/// it is truncated or longer than 10 bytes.
23pub fn read_u64(bytes: &mut &[u8]) -> Option<u64> {
24    let mut v: u64 = 0;
25    for (i, &b) in bytes.iter().enumerate().take(10) {
26        let part = (b & 0x7f) as u64;
27        if i == 9 && part > 1 {
28            return None; // past 64 bits
29        }
30        v |= part << (7 * i);
31        if b & 0x80 == 0 {
32            *bytes = &bytes[i + 1..];
33            return Some(v);
34        }
35    }
36    None
37}
38
39/// Append a signed integer as a zigzag varint (small magnitudes are short).
40pub fn write_i64(out: &mut Vec<u8>, v: i64) {
41    write_u64(out, ((v << 1) ^ (v >> 63)) as u64);
42}
43
44pub fn read_i64(bytes: &mut &[u8]) -> Option<i64> {
45    let z = read_u64(bytes)?;
46    Some(((z >> 1) as i64) ^ -((z & 1) as i64))
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52
53    #[test]
54    fn round_trips_and_sizes() {
55        for v in [0u64, 1, 127, 128, 300, u32::MAX as u64, u64::MAX] {
56            let mut b = Vec::new();
57            write_u64(&mut b, v);
58            let mut s = &b[..];
59            assert_eq!(read_u64(&mut s), Some(v));
60            assert!(s.is_empty());
61        }
62        for v in [0i64, -1, 1, -64, 63, -65, i64::MIN, i64::MAX] {
63            let mut b = Vec::new();
64            write_i64(&mut b, v);
65            let mut s = &b[..];
66            assert_eq!(read_i64(&mut s), Some(v));
67        }
68        let mut b = Vec::new();
69        write_i64(&mut b, -63);
70        assert_eq!(b.len(), 1);
71    }
72
73    #[test]
74    fn len_matches_the_encoding() {
75        for v in [0u64, 1, 127, 128, 16383, 16384, u64::MAX] {
76            let mut b = Vec::new();
77            write_u64(&mut b, v);
78            assert_eq!(len_u64(v), b.len(), "{v}");
79        }
80    }
81
82    #[test]
83    fn truncated_and_overlong_varints_are_refused() {
84        assert_eq!(read_u64(&mut &[0x80u8][..]), None);
85        assert_eq!(read_u64(&mut &[0xffu8; 11][..]), None);
86        assert_eq!(
87            read_u64(&mut &[0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x02][..]),
88            None
89        );
90    }
91}