wasmtime_internal_core/alloc/
try_collect.rs

1use crate::alloc::Vec;
2use crate::error::OutOfMemory;
3use std_alloc::boxed::Box;
4
5/// Extension trait for an `Iterator` to fallibly collect into a container.
6pub trait TryCollect: Iterator {
7    /// Attempts to collect the iterator `self` into `B`.
8    ///
9    /// Same as [`Iterator::collect`] except returns OOM instead of aborting.
10    fn try_collect<B, E>(self) -> Result<B, E>
11    where
12        B: TryFromIterator<Self::Item, E>,
13        Self: Sized,
14    {
15        B::try_from_iter(self)
16    }
17}
18
19impl<I: Iterator> TryCollect for I {}
20
21/// Analogue of [`FromIterator`] in the standard library, but used with
22/// [`TryCollect::try_collect`] instead.
23pub trait TryFromIterator<T, E>: Sized {
24    /// Creates an intance of this collection from the `iter` provided.
25    ///
26    /// Does not abort on OOM but insteads return an error.
27    fn try_from_iter<I>(iter: I) -> Result<Self, E>
28    where
29        I: Iterator<Item = T>;
30}
31
32impl<T> TryFromIterator<T, OutOfMemory> for Vec<T> {
33    fn try_from_iter<I>(iter: I) -> Result<Self, OutOfMemory>
34    where
35        I: Iterator<Item = T>,
36    {
37        let mut result = Vec::with_capacity(iter.size_hint().0)?;
38        for item in iter {
39            result.push(item)?;
40        }
41        Ok(result)
42    }
43}
44
45impl<T> TryFromIterator<T, OutOfMemory> for Box<[T]> {
46    fn try_from_iter<I>(iter: I) -> Result<Self, OutOfMemory>
47    where
48        I: Iterator<Item = T>,
49    {
50        let vec = Vec::try_from_iter(iter)?;
51        vec.into_boxed_slice()
52    }
53}
54
55impl<T, E> TryFromIterator<Result<T, E>, E> for Vec<T>
56where
57    E: From<OutOfMemory>,
58{
59    fn try_from_iter<I>(iter: I) -> Result<Self, E>
60    where
61        I: Iterator<Item = Result<T, E>>,
62    {
63        let mut result = Vec::with_capacity(iter.size_hint().0)?;
64        for item in iter {
65            result.push(item?)?;
66        }
67        Ok(result)
68    }
69}
70
71impl<T, E> TryFromIterator<Result<T, E>, E> for Box<[T]>
72where
73    E: From<OutOfMemory>,
74{
75    fn try_from_iter<I>(iter: I) -> Result<Self, E>
76    where
77        I: Iterator<Item = Result<T, E>>,
78    {
79        let vec = iter.try_collect::<Vec<_>, E>()?;
80        Ok(vec.into_boxed_slice()?)
81    }
82}
83
84/// Analogue of [`Extend`] except handles OOM conditions.
85pub trait TryExtend<T> {
86    /// Extends `self` with the items from `iter`.
87    ///
88    /// Returns an error if allocation fails while adding items to `self`. If an
89    /// OOM happens then some items from `iter` may have been added to `self`
90    /// already. On OOM no further items from the iterator will be consumed.
91    fn try_extend<I>(&mut self, iter: I) -> Result<(), OutOfMemory>
92    where
93        I: IntoIterator<Item = T>;
94}
95
96impl<T> TryExtend<T> for Vec<T> {
97    fn try_extend<I>(&mut self, iter: I) -> Result<(), OutOfMemory>
98    where
99        I: IntoIterator<Item = T>,
100    {
101        let iter = iter.into_iter();
102        self.reserve(iter.size_hint().0)?;
103        for item in iter {
104            self.push(item)?;
105        }
106        Ok(())
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::{Box, TryCollect, TryExtend, Vec};
113    use crate::error::{OutOfMemory, Result};
114
115    #[test]
116    fn test_vec_collect() -> Result<(), OutOfMemory> {
117        let v: Vec<i32> = (0..10).try_collect()?;
118        assert_eq!(&*v, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
119        Ok(())
120    }
121
122    #[test]
123    fn test_box_collect() -> Result<(), OutOfMemory> {
124        let v: Box<[i32]> = (0..10).try_collect()?;
125        assert_eq!(&*v, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
126        Ok(())
127    }
128
129    #[test]
130    fn test_vec_result_collect() -> Result<()> {
131        let v: Result<Vec<i32>> = [].into_iter().try_collect();
132        assert!(v?.is_empty());
133
134        let v: Result<Vec<i32>> = [Ok(1), Ok(2)].into_iter().try_collect();
135        assert_eq!(&*v?, &[1, 2]);
136
137        let v: Result<Vec<i32>> = [Ok(1), Err(crate::format_err!("hi"))]
138            .into_iter()
139            .try_collect();
140        assert!(v.is_err());
141
142        let v: Result<Vec<i32>> = [Err(crate::format_err!("hi")), Ok(1)]
143            .into_iter()
144            .try_collect();
145        assert!(v.is_err());
146        Ok(())
147    }
148
149    #[test]
150    fn test_box_result_collect() -> Result<()> {
151        let v: Result<Box<[i32]>> = [].into_iter().try_collect();
152        assert!(v?.is_empty());
153
154        let v: Result<Box<[i32]>> = [Ok(1), Ok(2)].into_iter().try_collect();
155        assert_eq!(&*v?, &[1, 2]);
156
157        let v: Result<Box<[i32]>> = [Ok(1), Err(crate::format_err!("hi"))]
158            .into_iter()
159            .try_collect();
160        assert!(v.is_err());
161
162        let v: Result<Box<[i32]>> = [Err(crate::format_err!("hi")), Ok(1)]
163            .into_iter()
164            .try_collect();
165        assert!(v.is_err());
166        Ok(())
167    }
168
169    #[test]
170    fn test_try_extend() -> Result<(), OutOfMemory> {
171        let mut vec = Vec::new();
172        vec.try_extend([1, 2, 3].iter().cloned())?;
173        assert_eq!(&*vec, &[1, 2, 3]);
174
175        vec.try_extend([])?;
176        assert_eq!(&*vec, &[1, 2, 3]);
177        Ok(())
178    }
179}
wasmtime_internal_core/alloc/try_collect.rs

wasmtime_internal_core/alloc/
try_collect.rs
