wasmtime_internal_core/alloc/
vec.rs

1use crate::alloc::{TryClone, try_realloc};
2use crate::error::OutOfMemory;
3use core::borrow::Borrow;
4use core::{
5    cmp::Ordering,
6    fmt, mem,
7    num::NonZeroUsize,
8    ops::{Deref, DerefMut, Index, IndexMut},
9    slice::SliceIndex,
10};
11#[cfg(feature = "serde")]
12use serde::ser::SerializeSeq;
13use std_alloc::alloc::Layout;
14use std_alloc::boxed::Box;
15use std_alloc::vec::Vec as StdVec;
16
17/// Same as the [`std::vec!`] macro but returns an error on allocation failure.
18#[macro_export]
19macro_rules! try_vec {
20    ( $( $elem:expr ),* ) => {{
21        let len = $crate::private_len!( $( $elem ),* );
22        $crate::alloc::TryVec::with_capacity(len).and_then(|mut v| {
23            $( v.push($elem)?; )*
24            let _ = &mut v;
25            Ok(v)
26        })
27    }};
28
29    ( $elem:expr; $len:expr ) => {{
30        let len: usize = $len;
31        if let Some(len) = ::core::num::NonZeroUsize::new(len) {
32            let elem = $elem;
33            $crate::alloc::TryVec::from_elem(elem, len)
34        } else {
35            Ok($crate::alloc::TryVec::new())
36        }
37    }};
38
39}
40
41// Only for use by the `vec!` macro.
42#[doc(hidden)]
43#[macro_export]
44macro_rules! private_len {
45    ( ) => { 0 };
46    ( $e:expr $( , $es:expr )* ) => { 1 + $crate::private_len!( $( $es ),* ) };
47}
48
49/// Like `std::vec::Vec` but all methods that allocate force handling allocation
50/// failure.
51#[derive(PartialEq, Eq, PartialOrd, Ord, Hash)]
52pub struct TryVec<T> {
53    inner: StdVec<T>,
54}
55
56impl<T> Default for TryVec<T> {
57    fn default() -> Self {
58        Self {
59            inner: Default::default(),
60        }
61    }
62}
63
64impl<T: fmt::Debug> fmt::Debug for TryVec<T> {
65    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
66        fmt::Debug::fmt(&self.inner, f)
67    }
68}
69
70impl<T> TryClone for TryVec<T>
71where
72    T: TryClone,
73{
74    fn try_clone(&self) -> Result<Self, OutOfMemory> {
75        let mut v = TryVec::with_capacity(self.len())?;
76        for x in self {
77            v.push(x.try_clone()?).expect("reserved capacity");
78        }
79        Ok(v)
80    }
81}
82
83impl<T> TryVec<T> {
84    /// Same as [`std::vec::Vec::new`].
85    pub const fn new() -> Self {
86        Self {
87            inner: StdVec::new(),
88        }
89    }
90
91    /// Same as [`std::vec::Vec::with_capacity`] but returns an error on
92    /// allocation failure.
93    pub fn with_capacity(capacity: usize) -> Result<Self, OutOfMemory> {
94        let mut v = Self::new();
95        v.reserve(capacity)?;
96        Ok(v)
97    }
98
99    // For use with the `vec!` macro.
100    #[doc(hidden)]
101    #[inline]
102    pub fn from_elem(elem: T, len: NonZeroUsize) -> Result<Self, OutOfMemory>
103    where
104        T: TryClone,
105    {
106        let mut v = Self::with_capacity(len.get())?;
107
108        // Minimize calls to `TryClone` by always pushing `elem` itself as the
109        // last element.
110        for _ in 0..len.get() - 1 {
111            v.push(elem.try_clone()?)?;
112        }
113        v.push(elem)?;
114
115        Ok(v)
116    }
117
118    /// Same as [`std::vec::Vec::reserve`] but returns an error on allocation
119    /// failure.
120    pub fn reserve(&mut self, additional: usize) -> Result<(), OutOfMemory> {
121        self.inner.try_reserve(additional).map_err(|_| {
122            OutOfMemory::new(
123                self.len()
124                    .saturating_add(additional)
125                    .saturating_mul(mem::size_of::<T>()),
126            )
127        })
128    }
129
130    /// Same as [`std::vec::Vec::reserve_exact`] but returns an error on allocation
131    /// failure.
132    pub fn reserve_exact(&mut self, additional: usize) -> Result<(), OutOfMemory> {
133        self.inner
134            .try_reserve_exact(additional)
135            .map_err(|_| OutOfMemory::new(self.len().saturating_add(additional)))
136    }
137
138    /// Same as [`std::vec::Vec::len`].
139    pub fn len(&self) -> usize {
140        self.inner.len()
141    }
142
143    /// Same as [`std::vec::Vec::capacity`].
144    pub fn capacity(&self) -> usize {
145        self.inner.capacity()
146    }
147
148    /// Same as [`std::vec::Vec::is_empty`].
149    pub fn is_empty(&self) -> bool {
150        self.inner.is_empty()
151    }
152
153    /// Same as [`std::vec::Vec::push`] but returns an error on allocation
154    /// failure.
155    pub fn push(&mut self, value: T) -> Result<(), OutOfMemory> {
156        self.reserve(1)?;
157        self.inner.push(value);
158        Ok(())
159    }
160
161    /// Same as [`std::vec::Vec::pop`].
162    pub fn pop(&mut self) -> Option<T> {
163        self.inner.pop()
164    }
165
166    /// Same as [`std::vec::Vec::truncate`].
167    pub fn truncate(&mut self, len: usize) {
168        self.inner.truncate(len);
169    }
170
171    /// Same as [`std::vec::Vec::resize`] but returns an error on allocation
172    /// failure.
173    pub fn resize(&mut self, new_len: usize, value: T) -> Result<(), OutOfMemory>
174    where
175        T: TryClone,
176    {
177        match new_len.cmp(&self.len()) {
178            Ordering::Less => self.truncate(new_len),
179            Ordering::Equal => {}
180            Ordering::Greater => {
181                let delta = new_len - self.len();
182                self.reserve(delta)?;
183                // Minimize `try_clone` calls by always pushing `value` directly
184                // as the last element.
185                for _ in 0..delta - 1 {
186                    self.push(value.try_clone()?)?;
187                }
188                self.push(value)?;
189            }
190        }
191        Ok(())
192    }
193
194    /// Same as [`std::vec::Vec::resize_with`] but returns an error on
195    /// allocation failure.
196    pub fn resize_with<F>(&mut self, new_len: usize, f: F) -> Result<(), OutOfMemory>
197    where
198        F: FnMut() -> T,
199    {
200        let len = self.len();
201        if new_len > len {
202            self.reserve(new_len - len)?;
203        }
204        self.inner.resize_with(new_len, f);
205        Ok(())
206    }
207
208    /// Same as [`std::vec::Vec::retain`].
209    pub fn retain<F>(&mut self, f: F)
210    where
211        F: FnMut(&T) -> bool,
212    {
213        self.inner.retain(f);
214    }
215
216    /// Same as [`std::vec::Vec::retain_mut`].
217    pub fn retain_mut<F>(&mut self, f: F)
218    where
219        F: FnMut(&mut T) -> bool,
220    {
221        self.inner.retain_mut(f);
222    }
223
224    /// Same as [`std::vec::Vec::into_raw_parts`].
225    pub fn into_raw_parts(mut self) -> (*mut T, usize, usize) {
226        // NB: Can't use `Vec::into_raw_parts` until our MSRV is >= 1.93.
227        #[cfg(not(miri))]
228        {
229            let ptr = self.as_mut_ptr();
230            let len = self.len();
231            let cap = self.capacity();
232            mem::forget(self);
233            (ptr, len, cap)
234        }
235        // NB: Miri requires using `into_raw_parts`, but always run on nightly,
236        // so it's fine to use there.
237        #[cfg(miri)]
238        {
239            let _ = &mut self;
240            self.inner.into_raw_parts()
241        }
242    }
243
244    /// Same as [`std::vec::Vec::from_raw_parts`].
245    pub unsafe fn from_raw_parts(ptr: *mut T, length: usize, capacity: usize) -> Self {
246        TryVec {
247            // Safety: Same as our unsafe contract.
248            inner: unsafe { StdVec::from_raw_parts(ptr, length, capacity) },
249        }
250    }
251
252    /// Same as [`std::vec::Vec::drain`].
253    pub fn drain<R>(&mut self, range: R) -> std_alloc::vec::Drain<'_, T>
254    where
255        R: core::ops::RangeBounds<usize>,
256    {
257        self.inner.drain(range)
258    }
259
260    /// Same as [`std::vec::Vec::shrink_to_fit`] but returns an error on
261    /// allocation failure.
262    pub fn shrink_to_fit(&mut self) -> Result<(), OutOfMemory> {
263        // If our length is already equal to our capacity, then there is nothing
264        // to shrink.
265        if self.len() == self.capacity() {
266            return Ok(());
267        }
268
269        // `realloc` requires a non-zero original layout as well as a non-zero
270        // destination layout, so this guard ensures that the sizes below are
271        // all nonzero. This handles a few cases:
272        //
273        // * If `len == cap == 0` then no allocation has ever been made.
274        // * If `len == 0` and `cap != 0` then this function effectively frees
275        //   the memory.
276        // * If `T` is a zero-sized type then nothing's been allocated either.
277        //
278        // In all of these cases delegate to the standard library's
279        // `shrink_to_fit` which is guaranteed to not perform a `realloc`.
280        if self.is_empty() || mem::size_of::<T>() == 0 {
281            self.inner.shrink_to_fit();
282            return Ok(());
283        }
284
285        let (ptr, len, cap) = mem::take(self).into_raw_parts();
286        let layout = Layout::array::<T>(cap).unwrap();
287        let new_size = Layout::array::<T>(len).unwrap().size();
288
289        // SAFETY: `ptr` was previously allocated in the global allocator,
290        // `layout` has a nonzero size and matches the current allocation of
291        // `ptr`, `new_size` is nonzero, and `new_size` is a valid array size
292        // for `len` elements given its constructor.
293        let result = unsafe { try_realloc(ptr.cast(), layout, new_size) };
294
295        match result {
296            Ok(ptr) => {
297                // SAFETY: `result` is allocated with the global allocator and
298                // has room for exactly `[T; len]`.
299                *self = unsafe { Self::from_raw_parts(ptr.cast::<T>().as_ptr(), len, len) };
300                Ok(())
301            }
302            Err(oom) => {
303                // SAFETY: If reallocation fails then it's guaranteed that the
304                // original allocation is not tampered with, so it's safe to
305                // reassemble the original vector.
306                *self = unsafe { TryVec::from_raw_parts(ptr, len, cap) };
307                Err(oom)
308            }
309        }
310    }
311
312    /// Same as [`std::vec::Vec::into_boxed_slice`] but returns an error on
313    /// allocation failure.
314    pub fn into_boxed_slice(mut self) -> Result<Box<[T]>, OutOfMemory> {
315        self.shrink_to_fit()?;
316
317        // Once we've shrunken the allocation to just the actual length, we can
318        // use `std`'s `into_boxed_slice` without fear of `realloc`.
319        Ok(self.inner.into_boxed_slice())
320    }
321
322    /// Same as [`std::vec::Vec::clear`].
323    pub fn clear(&mut self) {
324        self.inner.clear();
325    }
326
327    /// Same as [`std::vec::Vec::as_mut_ptr`].
328    //
329    // Note that this is technically inherited through the `DerefMut` impl but
330    // that converts `&mut Self` to `&mut [T]` which invalidates all previously
331    // derived pointers. This causes problems in Miri so by having an inherent
332    // method here it means that the borrow scope matches what we want with
333    // Miri.
334    pub fn as_mut_ptr(&mut self) -> *mut T {
335        self.inner.as_mut_ptr()
336    }
337}
338
339impl<T> Deref for TryVec<T> {
340    type Target = [T];
341
342    fn deref(&self) -> &Self::Target {
343        &self.inner
344    }
345}
346
347impl<T> DerefMut for TryVec<T> {
348    fn deref_mut(&mut self) -> &mut Self::Target {
349        &mut self.inner
350    }
351}
352
353impl<T> AsRef<[T]> for TryVec<T> {
354    fn as_ref(&self) -> &[T] {
355        self
356    }
357}
358
359impl<T> Borrow<[T]> for TryVec<T> {
360    fn borrow(&self) -> &[T] {
361        self
362    }
363}
364
365impl<T, I> Index<I> for TryVec<T>
366where
367    I: SliceIndex<[T]>,
368{
369    type Output = <I as SliceIndex<[T]>>::Output;
370
371    fn index(&self, index: I) -> &Self::Output {
372        &self.inner[index]
373    }
374}
375
376impl<T, I> IndexMut<I> for TryVec<T>
377where
378    I: SliceIndex<[T]>,
379{
380    fn index_mut(&mut self, index: I) -> &mut Self::Output {
381        &mut self.inner[index]
382    }
383}
384
385impl<T> IntoIterator for TryVec<T> {
386    type Item = T;
387    type IntoIter = std_alloc::vec::IntoIter<T>;
388
389    fn into_iter(self) -> Self::IntoIter {
390        self.inner.into_iter()
391    }
392}
393
394impl<'a, T> IntoIterator for &'a TryVec<T> {
395    type Item = &'a T;
396
397    type IntoIter = core::slice::Iter<'a, T>;
398
399    fn into_iter(self) -> Self::IntoIter {
400        (**self).iter()
401    }
402}
403
404impl<'a, T> IntoIterator for &'a mut TryVec<T> {
405    type Item = &'a mut T;
406
407    type IntoIter = core::slice::IterMut<'a, T>;
408
409    fn into_iter(self) -> Self::IntoIter {
410        (**self).iter_mut()
411    }
412}
413
414impl<T> From<TryVec<T>> for StdVec<T> {
415    fn from(v: TryVec<T>) -> Self {
416        v.inner
417    }
418}
419
420impl<T> From<StdVec<T>> for TryVec<T> {
421    fn from(inner: StdVec<T>) -> Self {
422        Self { inner }
423    }
424}
425
426impl<T> From<Box<[T]>> for TryVec<T> {
427    fn from(boxed_slice: Box<[T]>) -> Self {
428        Self::from(StdVec::from(boxed_slice))
429    }
430}
431
432#[cfg(feature = "serde")]
433impl<T> serde::ser::Serialize for TryVec<T>
434where
435    T: serde::ser::Serialize,
436{
437    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
438    where
439        S: serde::Serializer,
440    {
441        let mut seq = serializer.serialize_seq(Some(self.len()))?;
442        for elem in self {
443            seq.serialize_element(elem)?;
444        }
445        seq.end()
446    }
447}
448
449#[cfg(feature = "serde")]
450impl<'de, T> serde::de::Deserialize<'de> for TryVec<T>
451where
452    T: serde::de::Deserialize<'de>,
453{
454    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
455    where
456        D: serde::Deserializer<'de>,
457    {
458        use core::marker::PhantomData;
459
460        struct Visitor<T>(PhantomData<fn() -> TryVec<T>>);
461
462        impl<'de, T> serde::de::Visitor<'de> for Visitor<T>
463        where
464            T: serde::de::Deserialize<'de>,
465        {
466            type Value = TryVec<T>;
467
468            fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
469                f.write_str("a `wasmtime_core::alloc::Vec` sequence")
470            }
471
472            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
473            where
474                A: serde::de::SeqAccess<'de>,
475            {
476                use serde::de::Error as _;
477
478                let mut v = TryVec::new();
479
480                if let Some(len) = seq.size_hint() {
481                    v.reserve_exact(len).map_err(|oom| A::Error::custom(oom))?;
482                }
483
484                while let Some(elem) = seq.next_element()? {
485                    v.push(elem).map_err(|oom| A::Error::custom(oom))?;
486                }
487
488                Ok(v)
489            }
490        }
491
492        deserializer.deserialize_seq(Visitor(PhantomData))
493    }
494}
495
496#[cfg(test)]
497mod tests {
498    use super::TryVec;
499    use crate::error::OutOfMemory;
500
501    #[test]
502    fn test_into_boxed_slice() -> Result<(), OutOfMemory> {
503        assert_eq!(*TryVec::<i32>::new().into_boxed_slice()?, []);
504
505        let mut vec = TryVec::new();
506        vec.push(1)?;
507        assert_eq!(*vec.into_boxed_slice()?, [1]);
508
509        let mut vec = TryVec::with_capacity(2)?;
510        vec.push(1)?;
511        assert_eq!(*vec.into_boxed_slice()?, [1]);
512
513        let mut vec = TryVec::with_capacity(2)?;
514        vec.push(1_u128)?;
515        assert_eq!(*vec.into_boxed_slice()?, [1]);
516
517        assert_eq!(*TryVec::<()>::new().into_boxed_slice()?, []);
518
519        let mut vec = TryVec::new();
520        vec.push(())?;
521        assert_eq!(*vec.into_boxed_slice()?, [()]);
522
523        let vec = TryVec::<i32>::with_capacity(2)?;
524        assert_eq!(*vec.into_boxed_slice()?, []);
525        Ok(())
526    }
527}
wasmtime_internal_core/alloc/vec.rs

wasmtime_internal_core/alloc/
vec.rs
