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