Skip to main content

wasmtime_internal_core/alloc/
vec.rs

1use crate::alloc::{TryClone, try_realloc};
2use crate::error::OutOfMemory;
3use core::{
4    fmt, mem,
5    ops::{Deref, DerefMut, Index, IndexMut},
6};
7use std_alloc::alloc::Layout;
8use std_alloc::boxed::Box;
9use std_alloc::vec::Vec as StdVec;
10
11/// Like `std::vec::Vec` but all methods that allocate force handling allocation
12/// failure.
13#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
14pub struct Vec<T> {
15    inner: StdVec<T>,
16}
17
18impl<T> Default for Vec<T> {
19    fn default() -> Self {
20        Self {
21            inner: Default::default(),
22        }
23    }
24}
25
26impl<T: fmt::Debug> fmt::Debug for Vec<T> {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        fmt::Debug::fmt(&self.inner, f)
29    }
30}
31
32impl<T> TryClone for Vec<T>
33where
34    T: TryClone,
35{
36    fn try_clone(&self) -> Result<Self, OutOfMemory> {
37        let mut v = Vec::with_capacity(self.len())?;
38        for x in self {
39            v.push(x.try_clone()?).expect("reserved capacity");
40        }
41        Ok(v)
42    }
43}
44
45impl<T> Vec<T> {
46    /// Same as [`std::vec::Vec::new`].
47    pub fn new() -> Self {
48        Default::default()
49    }
50
51    /// Same as [`std::vec::Vec::with_capacity`] but returns an error on
52    /// allocation failure.
53    pub fn with_capacity(capacity: usize) -> Result<Self, OutOfMemory> {
54        let mut v = Self::new();
55        v.reserve(capacity)?;
56        Ok(v)
57    }
58
59    /// Same as [`std::vec::Vec::reserve`] but returns an error on allocation
60    /// failure.
61    pub fn reserve(&mut self, additional: usize) -> Result<(), OutOfMemory> {
62        self.inner.try_reserve(additional).map_err(|_| {
63            OutOfMemory::new(
64                self.len()
65                    .saturating_add(additional)
66                    .saturating_mul(mem::size_of::<T>()),
67            )
68        })
69    }
70
71    /// Same as [`std::vec::Vec::reserve_exact`] but returns an error on allocation
72    /// failure.
73    pub fn reserve_exact(&mut self, additional: usize) -> Result<(), OutOfMemory> {
74        self.inner
75            .try_reserve_exact(additional)
76            .map_err(|_| OutOfMemory::new(self.len().saturating_add(additional)))
77    }
78
79    /// Same as [`std::vec::Vec::len`].
80    pub fn len(&self) -> usize {
81        self.inner.len()
82    }
83
84    /// Same as [`std::vec::Vec::capacity`].
85    pub fn capacity(&self) -> usize {
86        self.inner.capacity()
87    }
88
89    /// Same as [`std::vec::Vec::is_empty`].
90    pub fn is_empty(&self) -> bool {
91        self.inner.is_empty()
92    }
93
94    /// Same as [`std::vec::Vec::push`] but returns an error on allocation
95    /// failure.
96    pub fn push(&mut self, value: T) -> Result<(), OutOfMemory> {
97        self.reserve(1)?;
98        self.inner.push(value);
99        Ok(())
100    }
101
102    /// Same as [`std::vec::Vec::pop`].
103    pub fn pop(&mut self) -> Option<T> {
104        self.inner.pop()
105    }
106
107    /// Same as [`std::vec::Vec::into_raw_parts`].
108    pub fn into_raw_parts(mut self) -> (*mut T, usize, usize) {
109        // NB: Can't use `Vec::into_raw_parts` until our MSRV is >= 1.93.
110        #[cfg(not(miri))]
111        {
112            let ptr = self.as_mut_ptr();
113            let len = self.len();
114            let cap = self.capacity();
115            mem::forget(self);
116            (ptr, len, cap)
117        }
118        // NB: Miri requires using `into_raw_parts`, but always run on nightly,
119        // so it's fine to use there.
120        #[cfg(miri)]
121        {
122            let _ = &mut self;
123            self.inner.into_raw_parts()
124        }
125    }
126
127    /// Same as [`std::vec::Vec::from_raw_parts`].
128    pub unsafe fn from_raw_parts(ptr: *mut T, length: usize, capacity: usize) -> Self {
129        Vec {
130            // Safety: Same as our unsafe contract.
131            inner: unsafe { StdVec::from_raw_parts(ptr, length, capacity) },
132        }
133    }
134
135    /// Same as [`std::vec::Vec::drain`].
136    pub fn drain<R>(&mut self, range: R) -> std_alloc::vec::Drain<'_, T>
137    where
138        R: core::ops::RangeBounds<usize>,
139    {
140        self.inner.drain(range)
141    }
142
143    /// Same as [`std::vec::Vec::into_boxed_slice`].
144    pub fn into_boxed_slice(self) -> Result<Box<[T]>, OutOfMemory> {
145        // `realloc` requires a non-zero original layout as well as a non-zero
146        // destination layout, so this guard ensures that the sizes below are
147        // all nonzero. This handles a few case:
148        //
149        // * If `len == cap == 0` then no allocation has ever been made.
150        // * If `len == 0` and `cap != 0` then this function effectively frees
151        //   the memory.
152        // * If `T` is a zero-sized type then nothing's been allocated either.
153        // * If `len == cap` then the allocation doesn't need to be shrunken.
154        //
155        // In all of these cases delegate to the standard library's
156        // `into_boxed_slice` which is guaranteed to not perform a `realloc`.
157        if self.is_empty() || mem::size_of::<T>() == 0 || self.inner.len() == self.inner.capacity()
158        {
159            return Ok(self.inner.into_boxed_slice());
160        }
161
162        let (ptr, len, cap) = self.into_raw_parts();
163        let layout = Layout::array::<T>(cap).unwrap();
164        let new_len = Layout::array::<T>(len).unwrap().size();
165
166        // SAFETY: `ptr` was previously allocated in the global allocator,
167        // `layout` has a nonzero size and matches the current allocation of
168        // `ptr`, `new_size` is nonzero, and `new_size` is a valid array size
169        // for `len` elements given its constructor.
170        let result = unsafe { try_realloc(ptr.cast(), layout, new_len) };
171
172        match result {
173            Ok(ptr) => {
174                // SAFETY: `result` is allocated with the global allocator with
175                // an appropriate size/align to create this `Box` with.
176                unsafe {
177                    Ok(Box::from_raw(core::ptr::slice_from_raw_parts_mut(
178                        ptr.as_ptr().cast(),
179                        len,
180                    )))
181                }
182            }
183            Err(oom) => {
184                // SAFETY: If reallocation fails then it's guaranteed that the
185                // original allocation is not tampered with, so it's safe to
186                // reassemble it back into the original vector.
187                unsafe {
188                    let _ = Vec::from_raw_parts(ptr, len, cap);
189                }
190                Err(oom)
191            }
192        }
193    }
194}
195
196impl<T> Deref for Vec<T> {
197    type Target = [T];
198
199    fn deref(&self) -> &Self::Target {
200        &self.inner
201    }
202}
203
204impl<T> DerefMut for Vec<T> {
205    fn deref_mut(&mut self) -> &mut Self::Target {
206        &mut self.inner
207    }
208}
209
210impl<T> Index<usize> for Vec<T> {
211    type Output = T;
212
213    fn index(&self, index: usize) -> &Self::Output {
214        &self.inner[index]
215    }
216}
217
218impl<T> IndexMut<usize> for Vec<T> {
219    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
220        &mut self.inner[index]
221    }
222}
223
224impl<T> IntoIterator for Vec<T> {
225    type Item = T;
226    type IntoIter = std_alloc::vec::IntoIter<T>;
227
228    fn into_iter(self) -> Self::IntoIter {
229        self.inner.into_iter()
230    }
231}
232
233impl<'a, T> IntoIterator for &'a Vec<T> {
234    type Item = &'a T;
235
236    type IntoIter = core::slice::Iter<'a, T>;
237
238    fn into_iter(self) -> Self::IntoIter {
239        (**self).iter()
240    }
241}
242
243impl<'a, T> IntoIterator for &'a mut Vec<T> {
244    type Item = &'a mut T;
245
246    type IntoIter = core::slice::IterMut<'a, T>;
247
248    fn into_iter(self) -> Self::IntoIter {
249        (**self).iter_mut()
250    }
251}
252
253impl<T> From<Box<[T]>> for Vec<T> {
254    fn from(boxed_slice: Box<[T]>) -> Self {
255        Vec {
256            inner: StdVec::from(boxed_slice),
257        }
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use super::Vec;
264    use crate::error::OutOfMemory;
265
266    #[test]
267    fn test_into_boxed_slice() -> Result<(), OutOfMemory> {
268        assert_eq!(*Vec::<i32>::new().into_boxed_slice()?, []);
269
270        let mut vec = Vec::new();
271        vec.push(1)?;
272        assert_eq!(*vec.into_boxed_slice()?, [1]);
273
274        let mut vec = Vec::with_capacity(2)?;
275        vec.push(1)?;
276        assert_eq!(*vec.into_boxed_slice()?, [1]);
277
278        let mut vec = Vec::with_capacity(2)?;
279        vec.push(1_u128)?;
280        assert_eq!(*vec.into_boxed_slice()?, [1]);
281
282        assert_eq!(*Vec::<()>::new().into_boxed_slice()?, []);
283
284        let mut vec = Vec::new();
285        vec.push(())?;
286        assert_eq!(*vec.into_boxed_slice()?, [()]);
287
288        let vec = Vec::<i32>::with_capacity(2)?;
289        assert_eq!(*vec.into_boxed_slice()?, []);
290        Ok(())
291    }
292}