steel_utils/
front_vec.rs

1use std::{
2    io::{self, Write},
3    ops::{Deref, DerefMut},
4    pin::Pin,
5    task::{Context, Poll},
6};
7
8use tokio::io::AsyncWrite;
9
10/// Its like a vec but with reserveable front space.
11/// Its meant for our packet serialization,
12/// you can just put the len of the packet in front without reallocating
13/// keep in mind that calling multiple `set_in_front()` sets the data in reverse order compared to `extend_from_slice()`
14#[derive(Debug, Default, Clone, PartialEq, Eq)]
15pub struct FrontVec {
16    buf: Vec<u8>,
17    front_space: usize,
18}
19
20impl FrontVec {
21    /// Creates a new `FrontVec` with the given reserved front space and capacity.
22    #[must_use]
23    pub fn capacity(reserve: usize, capacity: usize) -> Self {
24        let mut buf = Vec::with_capacity(reserve + capacity);
25        buf.resize(reserve, 0);
26
27        Self {
28            buf,
29            front_space: reserve,
30        }
31    }
32
33    /// Creates a new `FrontVec` with the given reserved front space.
34    #[must_use]
35    pub fn new(reserve: usize) -> Self {
36        Self::capacity(reserve, 0)
37    }
38
39    /// Returns the length of the `FrontVec`.
40    #[must_use]
41    pub const fn len(&self) -> usize {
42        self.buf.len() - self.front_space
43    }
44
45    /// Returns whether the `FrontVec` is empty.
46    #[must_use]
47    pub const fn is_empty(&self) -> bool {
48        (self.buf.len() - self.front_space) == 0
49    }
50
51    /// Pushes a value to the back of the `FrontVec`.
52    pub fn push(&mut self, value: u8) {
53        self.buf.push(value);
54    }
55
56    /// Extends the `FrontVec` with a slice.
57    pub fn extend_from_slice(&mut self, other: &[u8]) {
58        self.buf.extend_from_slice(other);
59    }
60
61    /// Sets the data in the front of the `FrontVec`.
62    ///
63    /// # Panics
64    /// - If there is not enough reserved space.
65    #[track_caller]
66    pub fn set_in_front(&mut self, other: &[u8]) {
67        assert!(self.front_space >= other.len(), "Not enough reserved space");
68
69        let new_start = self.front_space - other.len();
70        self.buf[new_start..self.front_space].copy_from_slice(other);
71        self.front_space = new_start;
72    }
73
74    /// Returns a slice of the `FrontVec`.
75    #[must_use]
76    pub fn as_slice(&self) -> &[u8] {
77        &self.buf[self.front_space..self.buf.len()]
78    }
79
80    /// Returns a mutable slice of the `FrontVec`.
81    #[must_use]
82    pub fn as_mut_slice(&mut self) -> &mut [u8] {
83        let len = self.buf.len();
84        &mut self.buf[self.front_space..len]
85    }
86}
87
88impl Write for FrontVec {
89    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
90        self.buf.extend_from_slice(buf);
91        Ok(buf.len())
92    }
93    fn flush(&mut self) -> io::Result<()> {
94        Ok(())
95    }
96}
97
98impl AsyncWrite for FrontVec {
99    fn poll_write(
100        self: Pin<&mut Self>,
101        _: &mut Context<'_>,
102        buf: &[u8],
103    ) -> Poll<io::Result<usize>> {
104        let this = self.get_mut();
105        this.extend_from_slice(buf);
106        Poll::Ready(Ok(buf.len()))
107    }
108
109    fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
110        Poll::Ready(Ok(()))
111    }
112
113    fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
114        Poll::Ready(Ok(()))
115    }
116}
117
118impl Deref for FrontVec {
119    type Target = [u8];
120
121    fn deref(&self) -> &Self::Target {
122        self.as_slice()
123    }
124}
125
126impl DerefMut for FrontVec {
127    fn deref_mut(&mut self) -> &mut Self::Target {
128        self.as_mut_slice()
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn front_space_reservation_and_write_safe() {
138        let mut fv = FrontVec::capacity(4, 8);
139
140        assert_eq!(fv.front_space, 4);
141        assert_eq!(fv.len(), 0);
142        assert_eq!(fv.as_slice(), &[] as &[u8]);
143
144        fv.extend_from_slice(&[1, 2, 3]);
145        assert_eq!(fv.as_slice(), &[1, 2, 3]);
146
147        fv.set_in_front(&[0xAA, 0xBB]);
148        assert_eq!(fv.as_slice(), &[0xAA, 0xBB, 1, 2, 3]);
149
150        fv.set_in_front(&[0xCC]);
151        assert_eq!(fv.as_slice(), &[0xCC, 0xAA, 0xBB, 1, 2, 3]);
152
153        assert_eq!(fv.front_space, 1);
154    }
155
156    #[test]
157    #[should_panic(expected = "Not enough reserved space")]
158    fn set_in_front_panics_if_no_space() {
159        let mut fv = FrontVec::capacity(2, 4);
160        fv.set_in_front(&[1, 2, 3]);
161    }
162}
steel_utils/front_vec.rs

steel_utils/
front_vec.rs
