pilota/thrift/
varint_ext.rs1use std::{io::Error, mem::size_of};
2
3use integer_encoding::VarInt;
4
5pub const MSB: u8 = 0b1000_0000;
7
8pub trait VarIntExt {
9 fn varint_max_size() -> usize;
10}
11
12impl<VI: VarInt> VarIntExt for VI {
13 fn varint_max_size() -> usize {
14 (size_of::<VI>() * 8 + 7) / 7
15 }
16}
17
18#[derive(Default)]
19pub struct VarIntProcessor {
20 buf: [u8; 10],
21 maxsize: usize,
22 pub i: usize,
23}
24
25impl VarIntProcessor {
26 pub fn new<VI: VarIntExt>() -> VarIntProcessor {
27 VarIntProcessor {
28 maxsize: VI::varint_max_size(),
29 ..VarIntProcessor::default()
30 }
31 }
32 pub fn push(&mut self, b: u8) -> Result<(), Error> {
33 if self.i >= self.maxsize {
34 return Err(Error::new(
35 std::io::ErrorKind::InvalidData,
36 "Unterminated varint",
37 ));
38 }
39 self.buf[self.i] = b;
40 self.i += 1;
41 Ok(())
42 }
43 pub fn finished(&self) -> bool {
44 self.i > 0 && (self.buf[self.i - 1] & MSB == 0)
45 }
46 pub fn decode<VI: VarInt>(&self) -> Option<VI> {
47 Some(VI::decode_var(&self.buf[0..self.i])?.0)
48 }
49}