1use std::io;
2use std::io::{Read, Write};
3
4use byteorder::{ByteOrder, LittleEndian};
5
6use super::BinarySerializable;
7
8pub 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
22pub 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#[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#[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
134fn 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
152pub 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}
175pub 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}