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, 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 metadata_size() -> usize {
148        let mut size = size_of::<[L; 2]>();
149        if PAD {
150            add_padding::<[L; 2], T>(&mut size);
151        }
152        size
153    }
154
155    #[inline]
156    pub fn data_size(capacity: L) -> usize {
157        let usize_len: usize = u64::from(capacity) as usize;
158        usize_len.saturating_mul(size_of::<T>())
159    }
160
161    #[inline]
162    pub fn required_size_for_capacity(capacity: L) -> usize {
163        Self::metadata_size().saturating_add(Self::data_size(capacity))
164    }
165
166    #[inline]
167    pub fn capacity(&self) -> usize {
168        u64::from(self.get_capacity()) as usize
169    }
170
171    #[inline]
172    pub fn len(&self) -> usize {
173        u64::from(self.get_len()) as usize
174    }
175
176    #[inline]
177    pub fn is_empty(&self) -> bool {
178        self.len() == 0
179    }
180
181    #[inline]
182    pub fn get(&self, index: usize) -> Option<&T> {
183        if index >= self.len() {
184            return None;
185        }
186        Some(&self.slice[index])
187    }
188
189    #[inline]
190    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
191        if index >= self.len() {
192            return None;
193        }
194        Some(&mut self.slice[index])
195    }
196
197    #[inline]
198    pub fn first(&self) -> Option<&T> {
199        self.get(0)
200    }
201
202    #[inline]
203    pub fn first_mut(&mut self) -> Option<&mut T> {
204        self.get_mut(0)
205    }
206
207    #[inline]
208    pub fn last(&self) -> Option<&T> {
209        self.get(self.len().saturating_sub(1))
210    }
211
212    #[inline]
213    pub fn last_mut(&mut self) -> Option<&mut T> {
214        self.get_mut(self.len().saturating_sub(1))
215    }
216
217    #[inline]
218    pub fn as_slice(&self) -> &[T] {
219        &self.slice[..self.len()]
220    }
221
222    #[inline]
223    pub fn as_mut_slice(&mut self) -> &mut [T] {
224        let len = self.len();
225        &mut self.slice[..len]
226    }
227
228    pub fn extend_from_slice(&mut self, slice: &[T]) -> Result<(), ZeroCopyError> {
229        let len = self.len();
230        let new_len = len + slice.len();
231        if new_len > self.capacity() {
232            return Err(ZeroCopyError::InsufficientCapacity);
233        }
234        self.slice[len..].copy_from_slice(slice);
235        *self.get_len_mut() = (new_len as u64)
236            .try_into()
237            .map_err(|_| ZeroCopyError::InvalidConversion)?;
238        Ok(())
239    }
240
241    #[cfg(feature = "std")]
242    #[inline]
243    pub fn to_vec(&self) -> Vec<T> {
244        self.as_slice().to_vec()
245    }
246
247    #[cfg(feature = "std")]
248    pub fn try_into_array<const N: usize>(&self) -> Result<[T; N], ZeroCopyError> {
249        if self.len() != N {
250            return Err(ZeroCopyError::ArraySize(N, self.len()));
251        }
252        Ok(core::array::from_fn(|i| *self.get(i).unwrap()))
253    }
254}
255
256impl<L, T, const PAD: bool> IndexMut<usize> for ZeroCopyVec<'_, L, T, PAD>
257where
258    L: ZeroCopyTraits,
259    T: ZeroCopyTraits,
260    u64: From<L> + TryInto<L>,
261{
262    #[inline]
263    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
264        // Access the underlying mutable slice using as_mut_slice() and index it
265        &mut self.as_mut_slice()[index]
266    }
267}
268
269impl<L, T, const PAD: bool> Index<usize> for ZeroCopyVec<'_, L, T, PAD>
270where
271    L: ZeroCopyTraits,
272    T: ZeroCopyTraits,
273    u64: From<L> + TryInto<L>,
274{
275    type Output = T;
276
277    #[inline]
278    fn index(&self, index: usize) -> &Self::Output {
279        // Access the underlying slice using as_slice() and index it
280        &self.as_slice()[index]
281    }
282}
283
284impl<'a, L, T, const PAD: bool> IntoIterator for &'a ZeroCopyVec<'_, L, T, PAD>
285where
286    L: ZeroCopyTraits,
287    T: ZeroCopyTraits,
288    u64: From<L> + TryInto<L>,
289{
290    type Item = &'a T;
291    type IntoIter = slice::Iter<'a, T>;
292
293    #[inline]
294    fn into_iter(self) -> Self::IntoIter {
295        self.iter()
296    }
297}
298
299impl<'a, L, T, const PAD: bool> IntoIterator for &'a mut ZeroCopyVec<'_, L, T, PAD>
300where
301    L: ZeroCopyTraits,
302    T: ZeroCopyTraits,
303    u64: From<L> + TryInto<L>,
304{
305    type Item = &'a mut T;
306    type IntoIter = slice::IterMut<'a, T>;
307
308    #[inline]
309    fn into_iter(self) -> Self::IntoIter {
310        self.iter_mut()
311    }
312}
313
314impl<'a, L, T, const PAD: bool> ZeroCopyVec<'_, L, T, PAD>
315where
316    L: ZeroCopyTraits,
317    T: ZeroCopyTraits,
318    u64: From<L> + TryInto<L>,
319{
320    #[inline]
321    pub fn iter(&'a self) -> slice::Iter<'a, T> {
322        self.as_slice().iter()
323    }
324
325    #[inline]
326    pub fn iter_mut(&'a mut self) -> slice::IterMut<'a, T> {
327        self.as_mut_slice().iter_mut()
328    }
329}
330
331impl<L, T, const PAD: bool> PartialEq for ZeroCopyVec<'_, L, T, PAD>
332where
333    L: ZeroCopyTraits,
334    T: ZeroCopyTraits + PartialEq,
335    u64: From<L> + TryInto<L>,
336{
337    #[inline]
338    fn eq(&self, other: &Self) -> bool {
339        self.as_slice() == other.as_slice()
340    }
341}
342
343impl<L, T, const PAD: bool> fmt::Debug for ZeroCopyVec<'_, L, T, PAD>
344where
345    L: ZeroCopyTraits,
346    T: ZeroCopyTraits + fmt::Debug,
347    u64: From<L> + TryInto<L>,
348{
349    #[inline]
350    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
351        write!(f, "{:?}", self.as_slice())
352    }
353}
354
355#[test]
356fn test_private_getters() {
357    let mut backing_store = [0u8; 64];
358    let mut zcv = ZeroCopyVec::<u16, u16>::new(5, &mut backing_store).unwrap();
359    assert_eq!(zcv.get_len(), 0);
360    assert_eq!(zcv.get_capacity(), 5);
361    for i in 0..5 {
362        zcv.push(i).unwrap();
363        assert_eq!(zcv.get_len(), i + 1);
364        assert_eq!(zcv.get_len_mut(), &mut (i + 1));
365    }
366}