pilota/thrift/
varint_ext.rs

1use std::{io::Error, mem::size_of};
2
3use integer_encoding::VarInt;
4
5/// Most-significant byte, == 0x80
6pub 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}