sans_io_runtime/
buffer.rs

1use std::ops::{Deref, DerefMut};
2
3#[derive(Debug, Clone)]
4enum BufferViewInner<'a> {
5    Ref(&'a [u8]),
6}
7
8impl<'a> Deref for BufferViewInner<'a> {
9    type Target = [u8];
10    fn deref(&self) -> &Self::Target {
11        match self {
12            Self::Ref(r) => r,
13        }
14    }
15}
16
17impl<'a> From<&'a [u8]> for BufferViewInner<'a> {
18    fn from(value: &'a [u8]) -> Self {
19        Self::Ref(value)
20    }
21}
22
23#[derive(Debug, Clone)]
24enum BufferInner {
25    Vec(Vec<u8>),
26}
27
28impl BufferInner {
29    pub fn view(&self) -> BufferViewInner {
30        match self {
31            Self::Vec(v) => v.as_slice().into(),
32        }
33    }
34
35    pub fn allocate_front(&mut self, more: usize) {
36        match self {
37            Self::Vec(v) => {
38                let mut vec2 = vec![0; more + v.len()];
39                vec2[more..].copy_from_slice(v);
40                let _ = std::mem::replace(v, vec2);
41            }
42        }
43    }
44
45    pub fn allocate_back(&mut self, more: usize) {
46        match self {
47            Self::Vec(v) => {
48                v.append(&mut vec![0; more]);
49            }
50        }
51    }
52}
53
54impl Deref for BufferInner {
55    type Target = [u8];
56    fn deref(&self) -> &Self::Target {
57        match self {
58            Self::Vec(v) => v,
59        }
60    }
61}
62
63impl DerefMut for BufferInner {
64    fn deref_mut(&mut self) -> &mut Self::Target {
65        match self {
66            Self::Vec(v) => v,
67        }
68    }
69}
70
71impl From<Vec<u8>> for BufferInner {
72    fn from(value: Vec<u8>) -> Self {
73        Self::Vec(value)
74    }
75}
76
77#[derive(Debug, Clone)]
78pub struct BufferView<'a> {
79    buf: BufferViewInner<'a>,
80    pub range: std::ops::Range<usize>,
81}
82
83impl<'a> BufferView<'a> {
84    pub fn view(&self, range: std::ops::Range<usize>) -> Option<BufferView> {
85        if self.range.end - self.range.start >= range.end {
86            Some(BufferView {
87                buf: self.buf.clone(),
88                range: (self.range.start + range.start..self.range.start + range.end),
89            })
90        } else {
91            None
92        }
93    }
94
95    pub fn pop_back(&mut self, len: usize) -> Option<BufferView> {
96        if self.range.end - self.range.start >= len {
97            self.range.end -= len;
98            Some(BufferView {
99                buf: self.buf.clone(),
100                range: (self.range.end..self.range.end + len),
101            })
102        } else {
103            None
104        }
105    }
106
107    pub fn pop_front(&mut self, len: usize) -> Option<BufferView> {
108        if self.range.end - self.range.start >= len {
109            self.range.start += len;
110            Some(BufferView {
111                buf: self.buf.clone(),
112                range: (self.range.start - len..self.range.start),
113            })
114        } else {
115            None
116        }
117    }
118
119    pub fn to_slice2(&self) -> &'_ [u8] {
120        &self.buf.deref()[self.range.clone()]
121    }
122}
123
124impl<'a> Deref for BufferView<'a> {
125    type Target = [u8];
126
127    fn deref(&self) -> &Self::Target {
128        &self.buf.deref()[self.range.clone()]
129    }
130}
131
132impl<'a> Eq for BufferView<'a> {}
133
134impl<'a> PartialEq for BufferView<'a> {
135    fn eq(&self, other: &Self) -> bool {
136        self.deref() == other.deref()
137    }
138}
139
140impl<'a> From<&'a [u8]> for BufferView<'a> {
141    fn from(value: &'a [u8]) -> Self {
142        BufferView {
143            buf: value.into(),
144            range: (0..value.len()),
145        }
146    }
147}
148
149#[derive(Debug, Clone)]
150pub struct Buffer {
151    buf: BufferInner,
152    range: std::ops::Range<usize>,
153}
154
155impl Buffer {
156    /// Create a buffer with pre-defined size. The left path is useful when we
157    /// working with networking protocol. Sometime we need to append data to front of buffer
158    /// the left-padding with avoiding vector need to re-create
159    pub fn new(left: usize, main: usize) -> Buffer {
160        let v = vec![0; left + main];
161        Buffer {
162            buf: v.into(),
163            range: (left..left),
164        }
165    }
166
167    /// Create a buffer mut with append some bytes at first and some bytes at end
168    pub fn build(data: &[u8], more_left: usize, more_right: usize) -> Buffer {
169        let mut v = vec![0; more_left + data.len() + more_right];
170        v[more_left..more_left + data.len()].copy_from_slice(data);
171        Buffer {
172            buf: v.into(),
173            range: (more_left..more_left + data.len()),
174        }
175    }
176
177    /// Create a new buffer from a vec, we can manually set the start and end of the buffer.
178    pub fn from_vec_raw(data: Vec<u8>, range: std::ops::Range<usize>) -> Buffer {
179        assert!(range.end <= data.len());
180        Buffer {
181            buf: data.into(),
182            range,
183        }
184    }
185
186    /// Getting front buffer for writing
187    pub fn front_mut(&mut self, len: usize) -> &mut [u8] {
188        self.ensure_front(len);
189        &mut self.buf.deref_mut()[self.range.start - len..self.range.start]
190    }
191
192    /// Getting remain buffer for writing, this useful when using with UDP Socket
193    pub fn remain_mut(&mut self) -> &mut [u8] {
194        &mut self.buf.deref_mut()[self.range.start..]
195    }
196
197    /// Getting back buffer for writing
198    pub fn back_mut(&mut self, len: usize) -> &mut [u8] {
199        self.ensure_back(len);
200        &mut self.buf.deref_mut()[self.range.end..self.range.end + len]
201    }
202
203    /// Reverse the front for at least `len` bytes at front
204    pub fn ensure_front(&mut self, more: usize) {
205        if self.range.start < more {
206            self.buf.allocate_front(more - self.range.start);
207            self.range.start += more;
208            self.range.end += more;
209        }
210    }
211
212    /// Reverse the buffer for at least `len` bytes at back.
213    pub fn ensure_back(&mut self, more: usize) {
214        assert!(self.buf.len() >= self.range.end);
215        let remain = self.buf.len() - self.range.end;
216        if remain < more {
217            self.buf.allocate_back(more - remain);
218        }
219    }
220
221    pub fn truncate(&mut self, len: usize) -> Option<()> {
222        if self.range.end - self.range.start >= len {
223            self.range.end = self.range.start + len;
224            Some(())
225        } else {
226            None
227        }
228    }
229
230    pub fn push_front(&mut self, data: &[u8]) {
231        self.ensure_front(data.len());
232        self.range.start -= data.len();
233        self.buf.deref_mut()[self.range.start..self.range.start + data.len()].copy_from_slice(data);
234    }
235
236    pub fn push_back(&mut self, data: &[u8]) {
237        self.ensure_back(data.len());
238        self.buf.deref_mut()[self.range.end..self.range.end + data.len()].copy_from_slice(data);
239        self.range.end += data.len();
240    }
241
242    pub fn pop_back(&mut self, len: usize) -> Option<BufferView<'_>> {
243        if self.range.end - self.range.start >= len {
244            self.range.end -= len;
245            Some(BufferView {
246                buf: self.buf.view(),
247                range: (self.range.end..self.range.end + len),
248            })
249        } else {
250            None
251        }
252    }
253
254    pub fn pop_front(&mut self, len: usize) -> Option<BufferView<'_>> {
255        if self.range.end - self.range.start >= len {
256            self.range.start += len;
257            Some(BufferView {
258                buf: self.buf.view(),
259                range: (self.range.start - len..self.range.start),
260            })
261        } else {
262            None
263        }
264    }
265
266    pub fn move_front_right(&mut self, len: usize) -> Option<()> {
267        if self.range.start + len > self.range.end {
268            return None;
269        }
270        self.range.start += len;
271        Some(())
272    }
273
274    pub fn move_front_left(&mut self, len: usize) -> Option<()> {
275        if self.range.start < len {
276            return None;
277        }
278        self.range.start -= len;
279        Some(())
280    }
281
282    pub fn move_back_right(&mut self, len: usize) -> Option<()> {
283        if self.range.end + len > self.buf.len() {
284            return None;
285        }
286        self.range.end += len;
287        Some(())
288    }
289
290    pub fn move_back_left(&mut self, len: usize) -> Option<()> {
291        if self.range.end < len || self.range.end - len < self.range.start {
292            return None;
293        }
294        self.range.end -= len;
295        Some(())
296    }
297}
298
299impl From<Vec<u8>> for Buffer {
300    fn from(value: Vec<u8>) -> Self {
301        Buffer {
302            range: (0..value.len()),
303            buf: value.into(),
304        }
305    }
306}
307
308impl Deref for Buffer {
309    type Target = [u8];
310
311    fn deref(&self) -> &Self::Target {
312        &self.buf.deref()[self.range.start..self.range.end]
313    }
314}
315
316impl Eq for Buffer {}
317
318impl PartialEq for Buffer {
319    fn eq(&self, other: &Self) -> bool {
320        self.deref() == other.deref()
321    }
322}
323
324impl DerefMut for Buffer {
325    fn deref_mut(&mut self) -> &mut Self::Target {
326        &mut self.buf.deref_mut()[self.range.start..self.range.end]
327    }
328}
329
330#[cfg(test)]
331mod tests {
332    use std::{ops::Deref, vec};
333
334    use super::{Buffer, BufferInner, BufferView};
335
336    #[test]
337    fn simple_buffer_view() {
338        let data = vec![1, 2, 3, 4, 5, 6];
339        let mut buf: BufferView = (data.as_slice()).into();
340        assert_eq!(buf.len(), 6);
341        assert_eq!(buf.pop_back(2).expect("").deref(), &[5, 6]);
342        assert_eq!(buf.pop_front(2).expect("").deref(), &[1, 2]);
343        assert_eq!(buf.deref(), &[3, 4]);
344    }
345
346    #[test]
347    fn simple_buffer_mut() {
348        let mut buf = Buffer::build(&[1, 2, 3, 4, 5, 6], 4, 4);
349        assert_eq!(buf.deref(), &[1, 2, 3, 4, 5, 6]);
350        assert_eq!(buf.to_vec(), vec![1, 2, 3, 4, 5, 6]);
351        println!("{:?}", buf);
352        let res = buf.pop_back(2).expect("");
353        println!("{:?}", res);
354        assert_eq!(res.deref(), &[5, 6]);
355        assert_eq!(res.to_vec(), &[5, 6]);
356        assert_eq!(buf.pop_front(2).expect("").deref(), &[1, 2]);
357    }
358
359    #[test]
360    fn buffer_expand() {
361        let mut buf = Buffer::from(vec![1, 2, 3, 4]);
362        buf.ensure_front(2);
363        buf.ensure_back(2);
364        assert_eq!(buf.len(), 4);
365        assert_eq!(buf.front_mut(2), &[0, 0]);
366        assert_eq!(buf.back_mut(2), &[0, 0]);
367
368        buf.push_front(&[8, 8]);
369        buf.push_back(&[9, 9]);
370
371        assert_eq!(buf.deref(), &[8, 8, 1, 2, 3, 4, 9, 9]);
372    }
373
374    #[test]
375    fn buffer_inner_expand() {
376        let mut buf = BufferInner::from(vec![1, 2, 3, 4]);
377        buf.allocate_front(2);
378        assert_eq!(buf.deref(), &[0, 0, 1, 2, 3, 4]);
379
380        let mut buf = BufferInner::Vec(vec![1, 2, 3, 4]);
381        buf.allocate_back(2);
382        assert_eq!(buf.len(), 6);
383        assert_eq!(buf.deref(), &[1, 2, 3, 4, 0, 0]);
384    }
385}