light_zero_copy/
vec.rs

1use core::{
2    fmt,
3    mem::size_of,
4    ops::{Index, IndexMut},
5    slice,
6};
7#[cfg(feature = "std")]
8use std::vec::Vec;
9
10use zerocopy::{little_endian::U32, IntoBytes, Ref};
11
12use crate::{add_padding, errors::ZeroCopyError, ZeroCopyTraits};
13
14pub type ZeroCopyVecU64<'a, T> = ZeroCopyVec<'a, u64, T>;
15pub type ZeroCopyVecU32<'a, T> = ZeroCopyVec<'a, u32, T>;
16pub type ZeroCopyVecU16<'a, T> = ZeroCopyVec<'a, u16, T>;
17pub type ZeroCopyVecU8<'a, T> = ZeroCopyVec<'a, u8, T>;
18pub type ZeroCopyVecBorsh<'a, T> = ZeroCopyVec<'a, U32, T, false>;
19
20/// `ZeroCopyVec` is a custom vector implementation which forbids
21/// post-initialization reallocations. The size is not known during compile
22/// time (that makes it different from arrays), but can be defined only once
23/// (that makes it different from [`Vec`](std::vec::Vec)).
24pub struct ZeroCopyVec<'a, L, T, const PAD: bool = true>
25where
26    L: ZeroCopyTraits,
27    T: ZeroCopyTraits,
28{
29    /// [length, capacity]
30    metadata: Ref<&'a mut [u8], [L; 2]>,
31    slice: Ref<&'a mut [u8], [T]>,
32}
33
34const LENGTH_INDEX: usize = 0;
35const CAPACITY_INDEX: usize = 1;
36
37impl<'a, L, T, const PAD: bool> ZeroCopyVec<'a, L, T, PAD>
38where
39    L: ZeroCopyTraits,
40    T: ZeroCopyTraits,
41    u64: From<L> + TryInto<L>,
42{
43    pub fn new(capacity: L, bytes: &'a mut [u8]) -> Result<Self, ZeroCopyError> {
44        Ok(Self::new_at(capacity, bytes)?.0)
45    }
46
47    pub fn new_at(capacity: L, bytes: &'a mut [u8]) -> Result<(Self, &'a mut [u8]), ZeroCopyError> {
48        let metadata_size = Self::metadata_size();
49        if bytes.len() < metadata_size {
50            return Err(ZeroCopyError::InsufficientMemoryAllocated(
51                bytes.len(),
52                metadata_size,
53            ));
54        }
55        let (meta_data, bytes) = bytes.split_at_mut(metadata_size);
56
57        let (mut metadata, _padding) = Ref::<&mut [u8], [L; 2]>::from_prefix(meta_data)?;
58        if u64::from(metadata[LENGTH_INDEX]) != 0 || u64::from(metadata[CAPACITY_INDEX]) != 0 {
59            return Err(ZeroCopyError::MemoryNotZeroed);
60        }
61        metadata[CAPACITY_INDEX] = capacity;
62        let capacity_usize: usize = u64::from(metadata[CAPACITY_INDEX]) as usize;
63
64        let (slice, remaining_bytes) =
65            Ref::<&mut [u8], [T]>::from_prefix_with_elems(bytes, capacity_usize)?;
66        Ok((Self { metadata, slice }, remaining_bytes))
67    }
68
69    #[inline]
70    pub fn from_bytes(bytes: &'a mut [u8]) -> Result<Self, ZeroCopyError> {
71        Ok(Self::from_bytes_at(bytes)?.0)
72    }
73
74    #[inline]
75    pub fn from_bytes_at(bytes: &'a mut [u8]) -> Result<(Self, &'a mut [u8]), ZeroCopyError> {
76        let metadata_size = Self::metadata_size();
77        if bytes.len() < metadata_size {
78            return Err(ZeroCopyError::InsufficientMemoryAllocated(
79                bytes.len(),
80                metadata_size,
81            ));
82        }
83
84        let (meta_data, bytes) = bytes.split_at_mut(metadata_size);
85        let (metadata, _padding) = Ref::<&mut [u8], [L; 2]>::from_prefix(meta_data)?;
86        let usize_capacity: usize = u64::from(metadata[CAPACITY_INDEX]) as usize;
87        let usize_len: usize = u64::from(metadata[LENGTH_INDEX]) as usize;
88
89        if usize_len > usize_capacity {
90            return Err(ZeroCopyError::LengthGreaterThanCapacity);
91        }
92
93        let full_vector_size = Self::data_size(metadata[CAPACITY_INDEX]);
94        if bytes.len() < full_vector_size {
95            return Err(ZeroCopyError::InsufficientMemoryAllocated(
96                bytes.len() + metadata_size,
97                full_vector_size + metadata_size,
98            ));
99        }
100        let (slice, remaining_bytes) =
101            Ref::<&mut [u8], [T]>::from_prefix_with_elems(bytes, usize_capacity)?;
102        Ok((Self { metadata, slice }, remaining_bytes))
103    }
104
105    /// Convenience method to get the length of the vector.
106    #[inline]
107    fn get_len(&self) -> L {
108        self.metadata[LENGTH_INDEX]
109    }
110
111    /// Convenience method to get the length of the vector.
112    #[inline]
113    fn get_len_mut(&mut self) -> &mut L {
114        &mut self.metadata[LENGTH_INDEX]
115    }
116
117    /// Convenience method to get the capacity of the vector.
118    #[inline]
119    fn get_capacity(&self) -> L {
120        self.metadata[CAPACITY_INDEX]
121    }
122
123    #[inline]
124    pub fn push(&mut self, value: T) -> Result<(), ZeroCopyError> {
125        if self.len() == self.capacity() {
126            return Err(ZeroCopyError::Full);
127        }
128
129        let len = self.len();
130        self.slice[len] = value;
131        *self.get_len_mut() = (len as u64 + 1u64)
132            .try_into()
133            .map_err(|_| ZeroCopyError::InvalidConversion)?;
134
135        Ok(())
136    }
137
138    #[inline]
139    pub fn clear(&mut self) {
140        *self.get_len_mut() = 0
141            .try_into()
142            .map_err(|_| ZeroCopyError::InvalidConversion)
143            .unwrap();
144    }
145
146    #[inline]
147    pub fn zero_out(&mut self) {
148        *self.get_len_mut() = 0
149            .try_into()
150            .map_err(|_| ZeroCopyError::InvalidConversion)
151            .unwrap();
152        self.slice.as_mut_bytes().fill(0);
153    }
154
155    #[inline]
156    pub fn metadata_size() -> usize {
157        let mut size = size_of::<[L; 2]>();
158        if PAD {
159            add_padding::<[L; 2], T>(&mut size);
160        }
161        size
162    }
163
164    #[inline]
165    pub fn data_size(capacity: L) -> usize {
166        let usize_len: usize = u64::from(capacity) as usize;
167        usize_len.saturating_mul(size_of::<T>())
168    }
169
170    #[inline]
171    pub fn required_size_for_capacity(capacity: L) -> usize {
172        Self::metadata_size().saturating_add(Self::data_size(capacity))
173    }
174
175    #[inline]
176    pub fn capacity(&self) -> usize {
177        u64::from(self.get_capacity()) as usize
178    }
179
180    #[inline]
181    pub fn len(&self) -> usize {
182        u64::from(self.get_len()) as usize
183    }
184
185    #[inline]
186    pub fn is_empty(&self) -> bool {
187        self.len() == 0
188    }
189
190    #[inline]
191    pub fn get(&self, index: usize) -> Option<&T> {
192        if index >= self.len() {
193            return None;
194        }
195        Some(&self.slice[index])
196    }
197
198    #[inline]
199    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
200        if index >= self.len() {
201            return None;
202        }
203        Some(&mut self.slice[index])
204    }
205
206    #[inline]
207    pub fn first(&self) -> Option<&T> {
208        self.get(0)
209    }
210
211    #[inline]
212    pub fn first_mut(&mut self) -> Option<&mut T> {
213        self.get_mut(0)
214    }
215
216    #[inline]
217    pub fn last(&self) -> Option<&T> {
218        self.get(self.len().saturating_sub(1))
219    }
220
221    #[inline]
222    pub fn last_mut(&mut self) -> Option<&mut T> {
223        self.get_mut(self.len().saturating_sub(1))
224    }
225
226    #[inline]
227    pub fn as_slice(&self) -> &[T] {
228        &self.slice[..self.len()]
229    }
230
231    #[inline]
232    pub fn as_mut_slice(&mut self) -> &mut [T] {
233        let len = self.len();
234        &mut self.slice[..len]
235    }
236
237    pub fn extend_from_slice(&mut self, slice: &[T]) -> Result<(), ZeroCopyError> {
238        let len = self.len();
239        let new_len = len + slice.len();
240        if new_len > self.capacity() {
241            return Err(ZeroCopyError::InsufficientCapacity);
242        }
243        self.slice[len..].copy_from_slice(slice);
244        *self.get_len_mut() = (new_len as u64)
245            .try_into()
246            .map_err(|_| ZeroCopyError::InvalidConversion)?;
247        Ok(())
248    }
249
250    #[cfg(feature = "std")]
251    #[inline]
252    pub fn to_vec(&self) -> Vec<T> {
253        self.as_slice().to_vec()
254    }
255
256    #[cfg(feature = "std")]
257    pub fn try_into_array<const N: usize>(&self) -> Result<[T; N], ZeroCopyError> {
258        if self.len() != N {
259            return Err(ZeroCopyError::ArraySize(N, self.len()));
260        }
261        Ok(core::array::from_fn(|i| *self.get(i).unwrap()))
262    }
263}
264
265impl<L, T, const PAD: bool> IndexMut<usize> for ZeroCopyVec<'_, L, T, PAD>
266where
267    L: ZeroCopyTraits,
268    T: ZeroCopyTraits,
269    u64: From<L> + TryInto<L>,
270{
271    #[inline]
272    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
273        // Access the underlying mutable slice using as_mut_slice() and index it
274        &mut self.as_mut_slice()[index]
275    }
276}
277
278impl<L, T, const PAD: bool> Index<usize> for ZeroCopyVec<'_, L, T, PAD>
279where
280    L: ZeroCopyTraits,
281    T: ZeroCopyTraits,
282    u64: From<L> + TryInto<L>,
283{
284    type Output = T;
285
286    #[inline]
287    fn index(&self, index: usize) -> &Self::Output {
288        // Access the underlying slice using as_slice() and index it
289        &self.as_slice()[index]
290    }
291}
292
293impl<'a, L, T, const PAD: bool> IntoIterator for &'a ZeroCopyVec<'_, L, T, PAD>
294where
295    L: ZeroCopyTraits,
296    T: ZeroCopyTraits,
297    u64: From<L> + TryInto<L>,
298{
299    type Item = &'a T;
300    type IntoIter = slice::Iter<'a, T>;
301
302    #[inline]
303    fn into_iter(self) -> Self::IntoIter {
304        self.iter()
305    }
306}
307
308impl<'a, L, T, const PAD: bool> IntoIterator for &'a mut ZeroCopyVec<'_, L, T, PAD>
309where
310    L: ZeroCopyTraits,
311    T: ZeroCopyTraits,
312    u64: From<L> + TryInto<L>,
313{
314    type Item = &'a mut T;
315    type IntoIter = slice::IterMut<'a, T>;
316
317    #[inline]
318    fn into_iter(self) -> Self::IntoIter {
319        self.iter_mut()
320    }
321}
322
323impl<'a, L, T, const PAD: bool> ZeroCopyVec<'_, L, T, PAD>
324where
325    L: ZeroCopyTraits,
326    T: ZeroCopyTraits,
327    u64: From<L> + TryInto<L>,
328{
329    #[inline]
330    pub fn iter(&'a self) -> slice::Iter<'a, T> {
331        self.as_slice().iter()
332    }
333
334    #[inline]
335    pub fn iter_mut(&'a mut self) -> slice::IterMut<'a, T> {
336        self.as_mut_slice().iter_mut()
337    }
338}
339
340impl<L, T, const PAD: bool> PartialEq for ZeroCopyVec<'_, L, T, PAD>
341where
342    L: ZeroCopyTraits,
343    T: ZeroCopyTraits + PartialEq,
344    u64: From<L> + TryInto<L>,
345{
346    #[inline]
347    fn eq(&self, other: &Self) -> bool {
348        self.as_slice() == other.as_slice()
349    }
350}
351
352impl<L, T, const PAD: bool> fmt::Debug for ZeroCopyVec<'_, L, T, PAD>
353where
354    L: ZeroCopyTraits,
355    T: ZeroCopyTraits + fmt::Debug,
356    u64: From<L> + TryInto<L>,
357{
358    #[inline]
359    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
360        write!(f, "{:?}", self.as_slice())
361    }
362}
363
364#[test]
365fn test_private_getters() {
366    let mut backing_store = [0u8; 64];
367    let mut zcv = ZeroCopyVec::<u16, u16>::new(5, &mut backing_store).unwrap();
368    assert_eq!(zcv.get_len(), 0);
369    assert_eq!(zcv.get_capacity(), 5);
370    for i in 0..5 {
371        zcv.push(i).unwrap();
372        assert_eq!(zcv.get_len(), i + 1);
373        assert_eq!(zcv.get_len_mut(), &mut (i + 1));
374    }
375}