1#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum WireType {
12 Varint,
14 I64,
16 Len,
19 I32,
21}
22
23impl WireType {
24 #[must_use]
27 pub fn from_bits(v: u32) -> Option<Self> {
28 Some(match v {
29 0 => Self::Varint,
30 1 => Self::I64,
31 2 => Self::Len,
32 5 => Self::I32,
33 _ => return None,
34 })
35 }
36}
37
38pub struct Reader<'a> {
41 buf: &'a [u8],
42 pos: usize,
43}
44
45impl<'a> Reader<'a> {
46 #[must_use]
48 pub fn new(buf: &'a [u8]) -> Self {
49 Self { buf, pos: 0 }
50 }
51
52 #[must_use]
54 pub fn at_end(&self) -> bool {
55 self.pos >= self.buf.len()
56 }
57
58 pub fn read_varint(&mut self) -> Option<u64> {
61 let mut result: u64 = 0;
62 let mut shift: u32 = 0;
63 loop {
64 let byte = *self.buf.get(self.pos)?;
65 self.pos += 1;
66 result |= u64::from(byte & 0x7f).checked_shl(shift)?;
67 if byte & 0x80 == 0 {
68 return Some(result);
69 }
70 shift += 7;
71 if shift >= 64 {
72 return None;
73 }
74 }
75 }
76
77 pub fn read_tag(&mut self) -> Option<(u32, WireType)> {
79 let tag = self.read_varint()?;
80 let field = u32::try_from(tag >> 3).ok()?;
81 if field == 0 {
82 return None;
83 }
84 let wt = WireType::from_bits((tag & 0x7) as u32)?;
85 Some((field, wt))
86 }
87
88 pub fn read_len_delimited(&mut self) -> Option<&'a [u8]> {
90 let len = usize::try_from(self.read_varint()?).ok()?;
91 let end = self.pos.checked_add(len)?;
92 let slice = self.buf.get(self.pos..end)?;
93 self.pos = end;
94 Some(slice)
95 }
96
97 pub fn read_i32(&mut self) -> Option<u32> {
99 let end = self.pos.checked_add(4)?;
100 let bytes = self.buf.get(self.pos..end)?;
101 self.pos = end;
102 Some(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
103 }
104
105 pub fn read_i64(&mut self) -> Option<u64> {
107 let end = self.pos.checked_add(8)?;
108 let bytes = self.buf.get(self.pos..end)?;
109 self.pos = end;
110 let mut arr = [0u8; 8];
111 arr.copy_from_slice(bytes);
112 Some(u64::from_le_bytes(arr))
113 }
114
115 pub fn skip(&mut self, wt: WireType) -> Option<()> {
118 match wt {
119 WireType::Varint => {
120 self.read_varint()?;
121 }
122 WireType::I64 => {
123 self.read_i64()?;
124 }
125 WireType::Len => {
126 self.read_len_delimited()?;
127 }
128 WireType::I32 => {
129 self.read_i32()?;
130 }
131 }
132 Some(())
133 }
134}
135
136#[cfg(test)]
137mod tests {
138 #![allow(clippy::unwrap_used)]
139 use super::*;
140
141 #[test]
142 fn varint_single_and_multi_byte() {
143 assert_eq!(Reader::new(&[0x00]).read_varint(), Some(0));
144 assert_eq!(Reader::new(&[0x7f]).read_varint(), Some(127));
145 assert_eq!(Reader::new(&[0xac, 0x02]).read_varint(), Some(300));
147 assert_eq!(Reader::new(&[0x80]).read_varint(), None);
149 }
150
151 #[test]
152 fn tag_splits_field_and_wire_type() {
153 assert_eq!(Reader::new(&[0x0a]).read_tag(), Some((1, WireType::Len)));
155 assert_eq!(Reader::new(&[0x18]).read_tag(), Some((3, WireType::Varint)));
157 assert_eq!(Reader::new(&[0x0b]).read_tag(), None);
159 }
160
161 #[test]
162 fn len_delimited_reads_exactly() {
163 let mut r = Reader::new(&[0x03, b'a', b'b', b'c']);
165 assert_eq!(r.read_len_delimited(), Some(&b"abc"[..]));
166 assert!(r.at_end());
167 assert_eq!(Reader::new(&[0x05, b'a']).read_len_delimited(), None);
169 }
170
171 #[test]
172 fn skip_advances_past_each_wire_type() {
173 let mut r = Reader::new(&[0x08, 0x96, 0x01, 0x12, 0x02, b'h', b'i']);
175 let (f, wt) = r.read_tag().unwrap();
176 assert_eq!((f, wt), (1, WireType::Varint));
177 r.skip(wt).unwrap();
178 let (f2, wt2) = r.read_tag().unwrap();
179 assert_eq!((f2, wt2), (2, WireType::Len));
180 assert_eq!(r.read_len_delimited(), Some(&b"hi"[..]));
181 }
182}