summavy_common/
vint.rs

1use std::io;
2use std::io::{Read, Write};
3
4use byteorder::{ByteOrder, LittleEndian};
5
6use super::BinarySerializable;
7
8/// Variable int serializes a u128 number
9pub fn serialize_vint_u128(mut val: u128, output: &mut Vec<u8>) {
10    loop {
11        let next_byte: u8 = (val % 128u128) as u8;
12        val /= 128u128;
13        if val == 0 {
14            output.push(next_byte | STOP_BIT);
15            return;
16        } else {
17            output.push(next_byte);
18        }
19    }
20}
21
22/// Deserializes a u128 number
23///
24/// Returns the number and the slice after the vint
25pub fn deserialize_vint_u128(data: &[u8]) -> io::Result<(u128, &[u8])> {
26    let mut result = 0u128;
27    let mut shift = 0u64;
28    for i in 0..19 {
29        let b = data[i];
30        result |= u128::from(b % 128u8) << shift;
31        if b >= STOP_BIT {
32            return Ok((result, &data[i + 1..]));
33        }
34        shift += 7;
35    }
36    Err(io::Error::new(
37        io::ErrorKind::InvalidData,
38        "Failed to deserialize u128 vint",
39    ))
40}
41
42///   Wrapper over a `u128` that serializes as a variable int.
43#[derive(Clone, Copy, Debug, Eq, PartialEq)]
44pub struct VIntU128(pub u128);
45
46impl BinarySerializable for VIntU128 {
47    fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
48        let mut buffer = vec![];
49        serialize_vint_u128(self.0, &mut buffer);
50        writer.write_all(&buffer)
51    }
52
53    fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
54        let mut bytes = reader.bytes();
55        let mut result = 0u128;
56        let mut shift = 0u64;
57        loop {
58            match bytes.next() {
59                Some(Ok(b)) => {
60                    result |= u128::from(b % 128u8) << shift;
61                    if b >= STOP_BIT {
62                        return Ok(VIntU128(result));
63                    }
64                    shift += 7;
65                }
66                _ => {
67                    return Err(io::Error::new(
68                        io::ErrorKind::InvalidData,
69                        "Reach end of buffer while reading VInt",
70                    ));
71                }
72            }
73        }
74    }
75}
76
77///   Wrapper over a `u64` that serializes as a variable int.
78#[derive(Clone, Copy, Debug, Eq, PartialEq)]
79pub struct VInt(pub u64);
80
81const STOP_BIT: u8 = 128;
82
83pub fn serialize_vint_u32(val: u32, buf: &mut [u8; 8]) -> &[u8] {
84    const START_2: u64 = 1 << 7;
85    const START_3: u64 = 1 << 14;
86    const START_4: u64 = 1 << 21;
87    const START_5: u64 = 1 << 28;
88
89    const STOP_1: u64 = START_2 - 1;
90    const STOP_2: u64 = START_3 - 1;
91    const STOP_3: u64 = START_4 - 1;
92    const STOP_4: u64 = START_5 - 1;
93
94    const MASK_1: u64 = 127;
95    const MASK_2: u64 = MASK_1 << 7;
96    const MASK_3: u64 = MASK_2 << 7;
97    const MASK_4: u64 = MASK_3 << 7;
98    const MASK_5: u64 = MASK_4 << 7;
99
100    let val = u64::from(val);
101    const STOP_BIT: u64 = 128u64;
102    let (res, num_bytes) = match val {
103        0..=STOP_1 => (val | STOP_BIT, 1),
104        START_2..=STOP_2 => (
105            (val & MASK_1) | ((val & MASK_2) << 1) | (STOP_BIT << (8)),
106            2,
107        ),
108        START_3..=STOP_3 => (
109            (val & MASK_1) | ((val & MASK_2) << 1) | ((val & MASK_3) << 2) | (STOP_BIT << (8 * 2)),
110            3,
111        ),
112        START_4..=STOP_4 => (
113            (val & MASK_1)
114                | ((val & MASK_2) << 1)
115                | ((val & MASK_3) << 2)
116                | ((val & MASK_4) << 3)
117                | (STOP_BIT << (8 * 3)),
118            4,
119        ),
120        _ => (
121            (val & MASK_1)
122                | ((val & MASK_2) << 1)
123                | ((val & MASK_3) << 2)
124                | ((val & MASK_4) << 3)
125                | ((val & MASK_5) << 4)
126                | (STOP_BIT << (8 * 4)),
127            5,
128        ),
129    };
130    LittleEndian::write_u64(&mut buf[..], res);
131    &buf[0..num_bytes]
132}
133
134/// Returns the number of bytes covered by a
135/// serialized vint `u32`.
136///
137/// Expects a buffer data that starts
138/// by the serialized `vint`, scans at most 5 bytes ahead until
139/// it finds the vint final byte.
140///
141/// # May Panic
142/// If the payload does not start by a valid `vint`
143fn vint_len(data: &[u8]) -> usize {
144    for (i, &val) in data.iter().enumerate().take(5) {
145        if val >= STOP_BIT {
146            return i + 1;
147        }
148    }
149    panic!("Corrupted data. Invalid VInt 32");
150}
151
152/// Reads a vint `u32` from a buffer, and
153/// consumes its payload data.
154///
155/// # Panics
156///
157/// If the buffer does not start by a valid
158/// vint payload
159pub fn read_u32_vint(data: &mut &[u8]) -> u32 {
160    let (result, vlen) = read_u32_vint_no_advance(data);
161    *data = &data[vlen..];
162    result
163}
164
165pub fn read_u32_vint_no_advance(data: &[u8]) -> (u32, usize) {
166    let vlen = vint_len(data);
167    let mut result = 0u32;
168    let mut shift = 0u64;
169    for &b in &data[..vlen] {
170        result |= u32::from(b & 127u8) << shift;
171        shift += 7;
172    }
173    (result, vlen)
174}
175/// Write a `u32` as a vint payload.
176pub fn write_u32_vint<W: io::Write>(val: u32, writer: &mut W) -> io::Result<()> {
177    let mut buf = [0u8; 8];
178    let data = serialize_vint_u32(val, &mut buf);
179    writer.write_all(data)
180}
181
182impl VInt {
183    pub fn val(&self) -> u64 {
184        self.0
185    }
186
187    pub fn deserialize_u64<R: Read>(reader: &mut R) -> io::Result<u64> {
188        VInt::deserialize(reader).map(|vint| vint.0)
189    }
190
191    pub fn serialize_into_vec(&self, output: &mut Vec<u8>) {
192        let mut buffer = [0u8; 10];
193        let num_bytes = self.serialize_into(&mut buffer);
194        output.extend(&buffer[0..num_bytes]);
195    }
196
197    pub fn serialize_into(&self, buffer: &mut [u8; 10]) -> usize {
198        let mut remaining = self.0;
199        for (i, b) in buffer.iter_mut().enumerate() {
200            let next_byte: u8 = (remaining % 128u64) as u8;
201            remaining /= 128u64;
202            if remaining == 0u64 {
203                *b = next_byte | STOP_BIT;
204                return i + 1;
205            } else {
206                *b = next_byte;
207            }
208        }
209        unreachable!();
210    }
211}
212
213impl BinarySerializable for VInt {
214    fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
215        let mut buffer = [0u8; 10];
216        let num_bytes = self.serialize_into(&mut buffer);
217        writer.write_all(&buffer[0..num_bytes])
218    }
219
220    fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
221        let mut bytes = reader.bytes();
222        let mut result = 0u64;
223        let mut shift = 0u64;
224        loop {
225            match bytes.next() {
226                Some(Ok(b)) => {
227                    result |= u64::from(b % 128u8) << shift;
228                    if b >= STOP_BIT {
229                        return Ok(VInt(result));
230                    }
231                    shift += 7;
232                }
233                _ => {
234                    return Err(io::Error::new(
235                        io::ErrorKind::InvalidData,
236                        "Reach end of buffer while reading VInt",
237                    ));
238                }
239            }
240        }
241    }
242}
243
244#[cfg(test)]
245mod tests {
246
247    use super::{serialize_vint_u32, BinarySerializable, VInt};
248    use crate::vint::{deserialize_vint_u128, serialize_vint_u128, VIntU128};
249
250    fn aux_test_vint(val: u64) {
251        let mut v = [14u8; 10];
252        let num_bytes = VInt(val).serialize_into(&mut v);
253        for el in &v[num_bytes..10] {
254            assert_eq!(el, &14u8);
255        }
256        assert!(num_bytes > 0);
257        if num_bytes < 10 {
258            assert!(1u64 << (7 * num_bytes) > val);
259        }
260        if num_bytes > 1 {
261            assert!(1u64 << (7 * (num_bytes - 1)) <= val);
262        }
263        let serdeser_val = VInt::deserialize(&mut &v[..]).unwrap();
264        assert_eq!(val, serdeser_val.0);
265    }
266
267    #[test]
268    fn test_vint() {
269        aux_test_vint(0);
270        aux_test_vint(1);
271        aux_test_vint(5);
272        aux_test_vint(u64::MAX);
273        for i in 1..9 {
274            let power_of_128 = 1u64 << (7 * i);
275            aux_test_vint(power_of_128 - 1u64);
276            aux_test_vint(power_of_128);
277            aux_test_vint(power_of_128 + 1u64);
278        }
279        aux_test_vint(10);
280    }
281
282    fn aux_test_serialize_vint_u32(val: u32) {
283        let mut buffer = [0u8; 10];
284        let mut buffer2 = [0u8; 8];
285        let len_vint = VInt(val as u64).serialize_into(&mut buffer);
286        let res2 = serialize_vint_u32(val, &mut buffer2);
287        assert_eq!(&buffer[..len_vint], res2, "array wrong for {}", val);
288    }
289
290    fn aux_test_vint_u128(val: u128) {
291        let mut data = vec![];
292        serialize_vint_u128(val, &mut data);
293        let (deser_val, _data) = deserialize_vint_u128(&data).unwrap();
294        assert_eq!(val, deser_val);
295
296        let mut out = vec![];
297        VIntU128(val).serialize(&mut out).unwrap();
298        let deser_val = VIntU128::deserialize(&mut &out[..]).unwrap();
299        assert_eq!(val, deser_val.0);
300    }
301
302    #[test]
303    fn test_vint_u128() {
304        aux_test_vint_u128(0);
305        aux_test_vint_u128(1);
306        aux_test_vint_u128(u128::MAX / 3);
307        aux_test_vint_u128(u128::MAX);
308    }
309
310    #[test]
311    fn test_vint_u32() {
312        aux_test_serialize_vint_u32(0);
313        aux_test_serialize_vint_u32(1);
314        aux_test_serialize_vint_u32(5);
315        for i in 1..3 {
316            let power_of_128 = 1u32 << (7 * i);
317            aux_test_serialize_vint_u32(power_of_128 - 1u32);
318            aux_test_serialize_vint_u32(power_of_128);
319            aux_test_serialize_vint_u32(power_of_128 + 1u32);
320        }
321        aux_test_serialize_vint_u32(u32::MAX);
322    }
323}