1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
#![feature(int_bits_const)]
mod overflow;
use std::{
io::{self, Read, Write},
mem::size_of,
};
pub use crate::overflow::OverflowError;
pub fn read_size_from_prefix(byte: u8) -> Result<usize, OverflowError> {
let size = byte.trailing_zeros() as usize;
if size > size_of::<u64>() {
Err(OverflowError::new(size))
} else {
Ok(size)
}
}
pub fn read_prefix(byte: u8) -> Result<(u64, usize), OverflowError> {
let size = read_size_from_prefix(byte)?;
Ok(((u64::from(byte) >> (size + 1)), size))
}
pub fn read_varint<R>(mut rdr: R) -> io::Result<u64>
where
R: Read,
{
let prefix_byte = {
let mut buf = [u8::MIN];
rdr.read_exact(&mut buf)?;
buf[0]
};
let (mut int, size) = read_prefix(prefix_byte)?;
let rest = {
let mut buf = vec![u8::MIN; size];
rdr.read_exact(&mut buf)?;
buf
};
let pad = 8_u8.min(size as u8 + 1);
for (i, byte) in (1_u8..).zip(rest) {
int |= u64::from(byte) << usize::from((i * 8) - pad)
}
Ok(int)
}
pub fn calc_varint_size(int: u64) -> usize {
let bits = match int.checked_next_power_of_two() {
Some(pow) => match pow {
pow if pow > int => pow,
pow => pow + pow,
}
.trailing_zeros(),
None => u64::BITS - 1,
} as usize;
1.max(bits / 7 + usize::from(bits % 7 > 0))
}
pub fn write_varint<W: Write>(mut int: u64, mut f: W) -> io::Result<usize> {
let size = calc_varint_size(int) as u8;
let packed_bits = 8_u8.saturating_sub(size);
let prefix = {
let int = if packed_bits == 0 {
0
} else {
((int & ((2_u64.pow(u32::from(packed_bits + 1))) - 1)) << size) as u8
};
int | (1_u8.checked_shl(u32::from(size) - 1).unwrap_or(0))
};
int >>= packed_bits;
f.write_all(&[prefix])?;
for _ in 0..(size - 1) {
let byte: u8 = (int & u64::from(u8::MAX)) as u8;
f.write_all(&[byte])?;
int >>= 8;
}
Ok(size.into())
}
#[cfg(test)]
mod tests;