protobuf/coded_input_stream/
input_buf.rs

1use std::cmp;
2use std::mem::MaybeUninit;
3
4use crate::misc::maybe_uninit_write_slice;
5
6#[derive(Debug)]
7pub(crate) struct InputBuf<'a> {
8    // Invariants: `0 <= pos_within_buf <= limit_within_buf <= buf.len()`.
9    buf: &'a [u8],
10    pos_within_buf: usize,
11    limit_within_buf: usize,
12}
13
14impl<'a> InputBuf<'a> {
15    #[inline]
16    pub(crate) fn assertions(&self) {
17        debug_assert!(self.pos_within_buf <= self.limit_within_buf);
18        debug_assert!(self.limit_within_buf <= self.buf.len());
19    }
20
21    pub(crate) fn empty() -> InputBuf<'a> {
22        InputBuf {
23            buf: &[],
24            pos_within_buf: 0,
25            limit_within_buf: 0,
26        }
27    }
28
29    pub(crate) fn from_bytes(buf: &'a [u8]) -> InputBuf<'a> {
30        InputBuf {
31            buf,
32            pos_within_buf: 0,
33            limit_within_buf: buf.len(),
34        }
35    }
36
37    pub(crate) unsafe fn from_bytes_ignore_lifetime(buf: &[u8]) -> InputBuf<'a> {
38        let buf = &*(buf as *const [u8]);
39        Self::from_bytes(buf)
40    }
41
42    pub(crate) fn update_limit(&mut self, limit: u64) {
43        let limit_within_buf = cmp::min(self.buf.len() as u64, limit);
44        assert!(limit_within_buf >= self.pos_within_buf as u64);
45        self.limit_within_buf = limit_within_buf as usize;
46    }
47
48    pub(crate) fn pos_within_buf(&self) -> usize {
49        self.pos_within_buf
50    }
51
52    #[inline(always)]
53    pub(crate) fn remaining_in_buf(&self) -> &'a [u8] {
54        // SAFETY: Invariants.
55        unsafe {
56            self.buf
57                .get_unchecked(self.pos_within_buf..self.limit_within_buf)
58        }
59    }
60
61    #[inline(always)]
62    pub(crate) fn consume(&mut self, amt: usize) {
63        assert!(amt <= self.remaining_in_buf().len());
64        self.pos_within_buf += amt;
65    }
66
67    #[inline(always)]
68    pub(crate) fn read_byte(&mut self) -> Option<u8> {
69        let r = self.remaining_in_buf().first().copied();
70        if let Some(..) = r {
71            self.pos_within_buf += 1;
72        }
73        r
74    }
75
76    pub(crate) fn read_bytes<'b>(&mut self, dest: &'b mut [MaybeUninit<u8>]) -> &'b mut [u8] {
77        // This panics if this has not enough data.
78        let r = maybe_uninit_write_slice(dest, &self.remaining_in_buf()[..dest.len()]);
79        self.pos_within_buf += r.len();
80        r
81    }
82}