wtx/misc/
filled_buffer.rs

1use crate::{
2  collection::Vector,
3  misc::{Lease, LeaseMut},
4};
5use core::{
6  fmt::Debug,
7  ops::{Deref, DerefMut},
8  slice,
9};
10
11/// A buffer that is always filled with initialized bytes.
12#[derive(Debug, Default)]
13pub struct FilledBuffer {
14  data: Vector<u8>,
15}
16
17impl FilledBuffer {
18  pub(crate) fn from_vector(mut vector: Vector<u8>) -> Self {
19    let prev_init = vector.len();
20    // SAFETY: elements up to `len` are always initialized
21    unsafe {
22      fill_remaining_capacity(&mut vector, prev_init);
23    }
24    Self { data: vector }
25  }
26
27  pub(crate) fn all(&self) -> &[u8] {
28    // SAFETY: allocated elements are always initialized
29    unsafe {
30      let len = self.data.capacity();
31      slice::from_raw_parts(self.data.as_ptr(), len)
32    }
33  }
34
35  pub(crate) fn all_mut(&mut self) -> &mut [u8] {
36    // SAFETY: allocated elements are always initialized
37    unsafe {
38      let len = self.data.capacity();
39      slice::from_raw_parts_mut(self.data.as_ptr_mut(), len)
40    }
41  }
42
43  pub(crate) fn capacity(&self) -> usize {
44    self.data.capacity()
45  }
46
47  pub(crate) fn clear(&mut self) {
48    self.data.clear();
49  }
50
51  #[cfg(test)]
52  pub(crate) fn extend_from_slices<'iter, I>(&mut self, others: I) -> crate::Result<usize>
53  where
54    I: IntoIterator<Item = &'iter [u8]>,
55    I::IntoIter: Clone,
56  {
57    let prev_init = self.data.capacity();
58    let len = self.data.extend_from_copyable_slices(others)?;
59    // SAFETY: inner elements up to `capacity` are always initialized
60    unsafe {
61      fill_remaining_capacity(&mut self.data, prev_init);
62    }
63    Ok(len)
64  }
65
66  #[inline(always)]
67  pub(crate) fn reserve(&mut self, additional: usize) -> crate::Result<()> {
68    let prev_init = self.data.capacity();
69    self.data.reserve(additional)?;
70    // SAFETY: inner elements up to `capacity` are always initialized
71    unsafe {
72      fill_remaining_capacity(&mut self.data, prev_init);
73    }
74    Ok(())
75  }
76
77  pub(crate) fn set_len(&mut self, mut len: usize) {
78    len = len.min(self.data.capacity());
79    // SAFETY: allocated memory is always initialized
80    unsafe { self.data.set_len(len) }
81  }
82}
83
84#[cfg(any(feature = "mysql", feature = "postgres", feature = "web-socket"))]
85impl FilledBuffer {
86  pub(crate) fn with_capacity(capacity: usize) -> crate::Result<Self> {
87    let mut data = Vector::with_capacity(capacity)?;
88    // SAFETY: memory have been allocated
89    unsafe {
90      slice::from_raw_parts_mut(data.as_ptr_mut(), data.capacity()).fill(0);
91    }
92    Ok(Self { data })
93  }
94}
95
96#[cfg(any(feature = "postgres", feature = "web-socket-handshake"))]
97impl FilledBuffer {
98  pub(crate) fn vector_mut(&mut self) -> FilledBufferVectorMut<'_> {
99    FilledBufferVectorMut { prev_init: self.data.capacity(), vector: &mut self.data }
100  }
101}
102
103impl Deref for FilledBuffer {
104  type Target = [u8];
105
106  #[inline]
107  fn deref(&self) -> &Self::Target {
108    self.data.as_slice()
109  }
110}
111
112impl DerefMut for FilledBuffer {
113  #[inline]
114  fn deref_mut(&mut self) -> &mut Self::Target {
115    self.data.as_slice_mut()
116  }
117}
118
119impl From<FilledBuffer> for Vector<u8> {
120  #[inline]
121  fn from(from: FilledBuffer) -> Self {
122    from.data
123  }
124}
125
126impl From<Vector<u8>> for FilledBuffer {
127  #[inline]
128  fn from(from: Vector<u8>) -> Self {
129    Self::from_vector(from)
130  }
131}
132
133#[cfg(feature = "std")]
134impl std::io::Write for FilledBuffer {
135  #[inline]
136  fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
137    self.data.write(buf)
138  }
139
140  #[inline]
141  fn flush(&mut self) -> std::io::Result<()> {
142    self.data.flush()
143  }
144}
145
146/// A wrapper that allows the direct usage of [Vector].
147#[derive(Debug)]
148pub struct FilledBufferVectorMut<'fb> {
149  prev_init: usize,
150  vector: &'fb mut Vector<u8>,
151}
152
153impl Lease<Vector<u8>> for FilledBufferVectorMut<'_> {
154  #[inline]
155  fn lease(&self) -> &Vector<u8> {
156    self.vector
157  }
158}
159
160impl LeaseMut<Vector<u8>> for FilledBufferVectorMut<'_> {
161  #[inline]
162  fn lease_mut(&mut self) -> &mut Vector<u8> {
163    self.vector
164  }
165}
166
167impl Drop for FilledBufferVectorMut<'_> {
168  #[inline]
169  fn drop(&mut self) {
170    // SAFETY: inner elements up to `capacity` are always initialized
171    unsafe {
172      fill_remaining_capacity(self.vector, self.prev_init);
173    }
174  }
175}
176
177unsafe fn fill_remaining_capacity(data: &mut Vector<u8>, prev_init: usize) {
178  let count = data.len().max(prev_init);
179  let Some(diff @ 1..=usize::MAX) = data.capacity().checked_sub(count) else {
180    return;
181  };
182  // SAFETY: caller must ensure `prev_init` elements
183  let ptr = unsafe { data.as_ptr_mut().add(count) };
184  // SAFETY: caller must ensure allocated memory
185  unsafe {
186    slice::from_raw_parts_mut(ptr, diff).fill(0);
187  }
188}
189
190#[cfg(kani)]
191mod kani {
192  use crate::misc::filled_buffer::FilledBuffer;
193
194  #[kani::proof]
195  fn reserve_is_allocation() {
196    let reserve: u8 = kani::any();
197    let mut vec = FilledBuffer::_new();
198    vec._reserve(reserve.into()).unwrap();
199    assert!(vec._capacity() >= reserve.into());
200    assert!(vec._all_mut().len() >= reserve.into());
201    assert_eq!(vec.len(), 0);
202    let len = 16usize.min(reserve.into());
203    vec._set_len(len);
204    assert_eq!(vec.len(), len);
205  }
206}
207
208#[cfg(test)]
209mod tests {
210  use crate::misc::FilledBuffer;
211
212  #[test]
213  fn extend_from_slices_with_increasing_cap() {
214    let mut vec = FilledBuffer::default();
215    let _ = vec.extend_from_slices([&[1, 2, 3][..]]).unwrap();
216    assert_eq!(&*vec, &[1, 2, 3]);
217  }
218}
wtx/misc/filled_buffer.rs

wtx/misc/
filled_buffer.rs
