gel_protogen/
writer.rs

1use std::mem::MaybeUninit;
2
3#[derive(Debug)]
4pub struct BufWriter<'a> {
5    buf: &'a mut [MaybeUninit<u8>],
6    size: usize,
7    error: bool,
8}
9
10impl<'a> BufWriter<'a> {
11    /// Create a new `BufWriter` from a slice of bytes.
12    #[inline(always)]
13    pub fn new(buf: &'a mut [u8]) -> Self {
14        Self {
15            // SAFETY: it's safe to go the other way as long as we never
16            // uninitialize bytes
17            buf: unsafe { std::mem::transmute(buf) },
18            size: 0,
19            error: false,
20        }
21    }
22
23    /// Create a new `BufWriter` from a slice of uninitialized bytes.
24    #[inline(always)]
25    pub fn new_uninit(buf: &'a mut [MaybeUninit<u8>]) -> Self {
26        Self {
27            buf,
28            size: 0,
29            error: false,
30        }
31    }
32
33    #[inline]
34    pub fn test(&mut self, size: usize) -> bool {
35        if self.buf.len() < size {
36            self.size += size;
37            self.error = true;
38            false
39        } else {
40            true
41        }
42    }
43
44    #[inline]
45    pub fn size(&self) -> usize {
46        self.size
47    }
48
49    #[inline]
50    pub fn write_rewind(&mut self, offset: usize, buf: &[u8]) {
51        if self.error {
52            return;
53        }
54        self.buf[offset..offset + buf.len()].copy_from_slice(unsafe { std::mem::transmute(buf) });
55    }
56
57    #[inline]
58    pub fn write(&mut self, buf: &[u8]) {
59        let len = buf.len();
60        self.size += len;
61        if self.error {
62            return;
63        }
64        if self.buf.len() < len {
65            self.error = true;
66            return;
67        }
68        self.buf[self.size - len..self.size].copy_from_slice(unsafe { std::mem::transmute(buf) });
69    }
70
71    #[inline]
72    pub fn write_u8(&mut self, value: u8) {
73        self.size += 1;
74        if self.error {
75            return;
76        }
77        if self.buf.is_empty() {
78            self.error = true;
79            return;
80        }
81        self.buf[self.size - 1].write(value);
82    }
83
84    pub const fn finish(self) -> Result<usize, usize> {
85        if self.error {
86            Err(self.size)
87        } else {
88            Ok(self.size)
89        }
90    }
91
92    /// Finish the writer and return a slice of the written bytes.
93    pub const fn finish_buf(self) -> Result<&'a mut [u8], usize> {
94        if self.error {
95            Err(self.size)
96        } else {
97            // SAFETY: we know that the buffer is valid because we've written to
98            // every byte
99            Ok(unsafe { std::slice::from_raw_parts_mut(self.buf.as_mut_ptr().cast(), self.size) })
100        }
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    #[test]
108    fn test_buf_writer() {
109        let mut buf = [0u8; 10];
110        let mut writer = BufWriter::new(&mut buf);
111        writer.write(b"hello");
112        assert_eq!(writer.size(), 5);
113    }
114
115    #[test]
116    fn test_buf_writer_too_large() {
117        let mut buf = [0u8; 10];
118        let mut writer = BufWriter::new(&mut buf);
119        writer.write(b"hello world");
120        assert_eq!(writer.size(), 11);
121        assert!(writer.error);
122    }
123}
gel_protogen/writer.rs

gel_protogen/
writer.rs
