Skip to main content

razor_stream/
buffer.rs

1/// A trait to adapt various type of buffer
2pub trait AllocateBuf: 'static + Sized + Send {
3    /// Alloc buffer or reserve space to fit blob_len inside the Buffer.
4    ///
5    /// When size is not enough, return None
6    fn reserve(&mut self, _blob_len: i32) -> Option<&mut [u8]>;
7}
8
9/// If Option is None, create a new `Vec<u8>` on call, otherwise grow to fit the requirement
10impl AllocateBuf for Option<Vec<u8>> {
11    #[inline]
12    #[allow(clippy::uninit_vec)]
13    fn reserve(&mut self, blob_len: i32) -> Option<&mut [u8]> {
14        let blob_len = blob_len as usize;
15        if let Some(buf) = self.as_mut() {
16            if buf.len() != blob_len {
17                if buf.capacity() < blob_len {
18                    buf.reserve(blob_len - buf.capacity());
19                }
20                unsafe { buf.set_len(blob_len) };
21            }
22        } else {
23            let mut v = Vec::with_capacity(blob_len);
24            unsafe { v.set_len(blob_len) };
25            self.replace(v);
26        }
27        return self.as_deref_mut();
28    }
29}
30
31/// Grow to fit the requirement
32impl AllocateBuf for Vec<u8> {
33    #[inline]
34    fn reserve(&mut self, blob_len: i32) -> Option<&mut [u8]> {
35        let blob_len = blob_len as usize;
36        if self.len() != blob_len {
37            if self.capacity() < blob_len {
38                self.reserve(blob_len - self.capacity());
39            }
40            unsafe { self.set_len(blob_len) };
41        }
42        return Some(self);
43    }
44}
45
46/// If Option is None, create a new [io_buffer::Buffer](https://docs.rs/io_buffer) on call.
47/// Otherwise will check the pre-allocated buffer.
48///
49/// RPC will return encode error or decode error when the size is not enough.
50impl AllocateBuf for Option<io_buffer::Buffer> {
51    #[inline]
52    fn reserve(&mut self, blob_len: i32) -> Option<&mut [u8]> {
53        if let Some(buf) = self.as_mut() {
54            let blob_len = blob_len as usize;
55            if buf.len() != blob_len {
56                if buf.capacity() < blob_len {
57                    return None;
58                }
59                buf.set_len(blob_len);
60            }
61        } else if let Ok(v) = io_buffer::Buffer::alloc(blob_len) {
62            self.replace(v);
63        } else {
64            // alloc failed
65            return None;
66        }
67        return self.as_deref_mut();
68    }
69}
70
71/// Check an pre-allocated [io_buffer::Buffer](https://docs.rs/io_buffer).
72///
73/// RPC will return encode error or decode error when the size is not enough.
74impl AllocateBuf for io_buffer::Buffer {
75    #[inline]
76    fn reserve(&mut self, blob_len: i32) -> Option<&mut [u8]> {
77        let blob_len = blob_len as usize;
78        if self.len() != blob_len {
79            if self.capacity() < blob_len {
80                return None;
81            }
82            self.set_len(blob_len);
83        }
84        Some(self)
85    }
86}