#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WireType {
Varint,
I64,
Len,
I32,
}
impl WireType {
#[must_use]
pub fn from_bits(v: u32) -> Option<Self> {
Some(match v {
0 => Self::Varint,
1 => Self::I64,
2 => Self::Len,
5 => Self::I32,
_ => return None,
})
}
}
pub struct Reader<'a> {
buf: &'a [u8],
pos: usize,
}
impl<'a> Reader<'a> {
#[must_use]
pub fn new(buf: &'a [u8]) -> Self {
Self { buf, pos: 0 }
}
#[must_use]
pub fn at_end(&self) -> bool {
self.pos >= self.buf.len()
}
pub fn read_varint(&mut self) -> Option<u64> {
let mut result: u64 = 0;
let mut shift: u32 = 0;
loop {
let byte = *self.buf.get(self.pos)?;
self.pos += 1;
result |= u64::from(byte & 0x7f).checked_shl(shift)?;
if byte & 0x80 == 0 {
return Some(result);
}
shift += 7;
if shift >= 64 {
return None;
}
}
}
pub fn read_tag(&mut self) -> Option<(u32, WireType)> {
let tag = self.read_varint()?;
let field = u32::try_from(tag >> 3).ok()?;
if field == 0 {
return None;
}
let wt = WireType::from_bits((tag & 0x7) as u32)?;
Some((field, wt))
}
pub fn read_len_delimited(&mut self) -> Option<&'a [u8]> {
let len = usize::try_from(self.read_varint()?).ok()?;
let end = self.pos.checked_add(len)?;
let slice = self.buf.get(self.pos..end)?;
self.pos = end;
Some(slice)
}
pub fn read_i32(&mut self) -> Option<u32> {
let end = self.pos.checked_add(4)?;
let bytes = self.buf.get(self.pos..end)?;
self.pos = end;
Some(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
}
pub fn read_i64(&mut self) -> Option<u64> {
let end = self.pos.checked_add(8)?;
let bytes = self.buf.get(self.pos..end)?;
self.pos = end;
let mut arr = [0u8; 8];
arr.copy_from_slice(bytes);
Some(u64::from_le_bytes(arr))
}
pub fn skip(&mut self, wt: WireType) -> Option<()> {
match wt {
WireType::Varint => {
self.read_varint()?;
}
WireType::I64 => {
self.read_i64()?;
}
WireType::Len => {
self.read_len_delimited()?;
}
WireType::I32 => {
self.read_i32()?;
}
}
Some(())
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
#[test]
fn varint_single_and_multi_byte() {
assert_eq!(Reader::new(&[0x00]).read_varint(), Some(0));
assert_eq!(Reader::new(&[0x7f]).read_varint(), Some(127));
assert_eq!(Reader::new(&[0xac, 0x02]).read_varint(), Some(300));
assert_eq!(Reader::new(&[0x80]).read_varint(), None);
}
#[test]
fn tag_splits_field_and_wire_type() {
assert_eq!(Reader::new(&[0x0a]).read_tag(), Some((1, WireType::Len)));
assert_eq!(Reader::new(&[0x18]).read_tag(), Some((3, WireType::Varint)));
assert_eq!(Reader::new(&[0x0b]).read_tag(), None);
}
#[test]
fn len_delimited_reads_exactly() {
let mut r = Reader::new(&[0x03, b'a', b'b', b'c']);
assert_eq!(r.read_len_delimited(), Some(&b"abc"[..]));
assert!(r.at_end());
assert_eq!(Reader::new(&[0x05, b'a']).read_len_delimited(), None);
}
#[test]
fn skip_advances_past_each_wire_type() {
let mut r = Reader::new(&[0x08, 0x96, 0x01, 0x12, 0x02, b'h', b'i']);
let (f, wt) = r.read_tag().unwrap();
assert_eq!((f, wt), (1, WireType::Varint));
r.skip(wt).unwrap();
let (f2, wt2) = r.read_tag().unwrap();
assert_eq!((f2, wt2), (2, WireType::Len));
assert_eq!(r.read_len_delimited(), Some(&b"hi"[..]));
}
}