Skip to main content

varintbuf/
lib.rs

1use bytes::{Buf, BufMut};
2
3/// Returns how many bytes are needed to encode a value.
4#[inline]
5pub fn length(value: u64) -> usize {
6    let zero_len = 64 - value.leading_zeros();
7    let offset = if zero_len == 0 { 7 } else { 6 };
8    ((offset + zero_len) / 7) as usize
9}
10
11/// Encode a `u64` integer to the [`bytes::BufMut`]. Returns how many bytes were
12/// encoded.
13#[inline]
14pub fn encode(mut val: u64, mut buf: impl BufMut) -> usize {
15    let mut written = 1;
16
17    while val > 127 {
18        buf.put_u8((val as u8) | 128);
19        val >>= 7;
20        written += 1;
21    }
22    buf.put_u8(val as u8);
23    written
24}
25
26/// Read a `u64` from a [`bytes::Buf`].
27#[inline]
28pub fn decode(mut buf: impl Buf) -> u64 {
29    let mut val = 0 as u64;
30    let mut fac = 1 as u64;
31
32    loop {
33        let byte = buf.get_u8();
34        val += fac * u64::from(byte & 127);
35        fac <<= 7;
36        if byte & 128 == 0 {
37            break;
38        }
39    }
40    val
41}
42
43/// Returns how many bytes are needed to encode a value.
44#[inline]
45pub fn signed_length(value: i64) -> usize {
46    length(unsign(value))
47}
48
49/// Encode a `i64` (signed) integer in the [`bytes::BufMut`].
50/// Returns how many bytes were encoded.
51#[inline]
52pub fn signed_encode(value: i64, buf: impl BufMut) -> usize {
53    encode(unsign(value), buf)
54}
55
56/// Decode a single `i64` (signed) integer from a [`bytes::Buf`].
57#[inline]
58pub fn signed_decode(buf: impl Buf) -> i64 {
59    sign(decode(buf))
60}
61
62/// Convert an `i64` into a `u64`.
63#[inline]
64fn unsign(value: i64) -> u64 {
65    if value >= 0 {
66        (value * 2) as u64
67    } else {
68        (value * -2 - 1) as u64
69    }
70}
71
72/// Convert a `u64` into a `i64`.
73#[inline]
74fn sign(value: u64) -> i64 {
75    if value & 1 != 0 {
76        -(((value + 1) / 2) as i64)
77    } else {
78        (value / 2) as i64
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use bytes::{Bytes, BytesMut};
86
87    #[test]
88    fn test_encode() {
89        let mut buf = Vec::new();
90        assert_eq!(encode(100, &mut buf), 1);
91        assert_eq!(buf[0], 100);
92
93        assert_eq!(encode(1000, &mut buf), 2);
94        assert_eq!(buf[1], 232);
95        assert_eq!(buf[2], 7);
96    }
97
98    #[test]
99    fn encoded2() {
100        // encode a `u64` integer into a buffer
101        let mut buf = Vec::new();
102        let mut bytes_encoded = encode(100, &mut buf);
103        assert_eq!(buf[0], 100);
104
105        let mut buf = BytesMut::from(&*buf);
106        bytes_encoded += encode(1000, &mut buf);
107        assert_eq!(buf[1], 232);
108        assert_eq!(buf[2], 7);
109        assert_eq!(bytes_encoded, 3);
110    }
111
112    #[test]
113    fn test_decode() {
114        let mut buf = &[100u8, 232, 7][..];
115        assert_eq!(decode(&mut buf), 100);
116        assert_eq!(buf, &[232, 7][..]);
117
118        assert_eq!(decode(&mut buf), 1000);
119        assert_eq!(buf.remaining(), 0);
120
121        let mut buf = Bytes::from(&[100u8, 232, 7][..]);
122        assert_eq!(decode(&mut buf), 100);
123        assert_eq!(decode(&*buf), 1000);
124    }
125
126    #[test]
127    fn test_length() {
128        assert_eq!(length(100), 1);
129        assert_eq!(length(1000), 2);
130
131        assert_eq!(length(1 << 49), 8);
132        assert_eq!(length((1 << 56) - 1), 8);
133
134        assert_eq!(length(1 << 56), 9);
135        assert_eq!(length((1 << 63) - 1), 9);
136
137        assert_eq!(length(1 << 63), 10);
138    }
139
140    #[test]
141    fn test_signed_encode() {
142        let mut buf = Vec::new();
143        assert_eq!(signed_encode(100, &mut buf), 2);
144        assert_eq!(buf[0], 200);
145        assert_eq!(buf[1], 1);
146
147        assert_eq!(signed_encode(-100, &mut buf), 2);
148        assert_eq!(buf[2], 199);
149        assert_eq!(buf[3], 1);
150        assert_eq!(buf.len(), 4);
151    }
152
153    #[test]
154    fn test_signed_decode() {
155        let mut buf = &[200u8, 1][..];
156        assert_eq!(signed_decode(&mut buf), 100);
157        assert_eq!(buf.remaining(), 0);
158
159        let mut buf = Bytes::from(&[199, 1][..]);
160        assert_eq!(signed_decode(&mut buf), -100);
161        assert_eq!(buf.remaining(), 0);
162    }
163
164    #[test]
165    fn test_signed_length() {
166        assert_eq!(signed_length(100), 2);
167        assert_eq!(signed_length(-100), 2);
168    }
169}