1 use crate::alloc::TryVec; 2 use crate::error::OutOfMemory; 3 use std_alloc::boxed::Box; 4 5 /// Extension trait for an `Iterator` to fallibly collect into a container. 6 pub trait TryCollect: Iterator { 7 /// Attempts to collect the iterator `self` into `B`. 8 /// 9 /// Same as [`Iterator::collect`] except returns OOM instead of aborting. try_collect<B, E>(self) -> Result<B, E> where B: TryFromIterator<Self::Item, E>, Self: Sized,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 19 impl<I: Iterator> TryCollect for I {} 20 21 /// Analogue of [`FromIterator`] in the standard library, but used with 22 /// [`TryCollect::try_collect`] instead. 23 pub trait TryFromIterator<T, E>: Sized { 24 /// Creates an instance of this collection from the `iter` provided. 25 /// 26 /// Does not abort on OOM but instead returns an error. try_from_iter<I>(iter: I) -> Result<Self, E> where I: Iterator<Item = T>27 fn try_from_iter<I>(iter: I) -> Result<Self, E> 28 where 29 I: Iterator<Item = T>; 30 } 31 32 impl<T> TryFromIterator<T, OutOfMemory> for TryVec<T> { try_from_iter<I>(iter: I) -> Result<Self, OutOfMemory> where I: Iterator<Item = T>,33 fn try_from_iter<I>(iter: I) -> Result<Self, OutOfMemory> 34 where 35 I: Iterator<Item = T>, 36 { 37 let mut result = TryVec::with_capacity(iter.size_hint().0)?; 38 for item in iter { 39 result.push(item)?; 40 } 41 Ok(result) 42 } 43 } 44 45 impl<T> TryFromIterator<T, OutOfMemory> for Box<[T]> { try_from_iter<I>(iter: I) -> Result<Self, OutOfMemory> where I: Iterator<Item = T>,46 fn try_from_iter<I>(iter: I) -> Result<Self, OutOfMemory> 47 where 48 I: Iterator<Item = T>, 49 { 50 let vec = TryVec::try_from_iter(iter)?; 51 vec.into_boxed_slice() 52 } 53 } 54 55 impl<T, E> TryFromIterator<Result<T, E>, E> for TryVec<T> 56 where 57 E: From<OutOfMemory>, 58 { try_from_iter<I>(iter: I) -> Result<Self, E> where I: Iterator<Item = Result<T, E>>,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 = TryVec::with_capacity(iter.size_hint().0)?; 64 for item in iter { 65 result.push(item?)?; 66 } 67 Ok(result) 68 } 69 } 70 71 impl<T, E> TryFromIterator<Result<T, E>, E> for Box<[T]> 72 where 73 E: From<OutOfMemory>, 74 { try_from_iter<I>(iter: I) -> Result<Self, E> where I: Iterator<Item = Result<T, E>>,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::<TryVec<_>, E>()?; 80 Ok(vec.into_boxed_slice()?) 81 } 82 } 83 84 /// Analogue of [`Extend`] except handles OOM conditions. 85 pub 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. try_extend<I>(&mut self, iter: I) -> Result<(), OutOfMemory> where I: IntoIterator<Item = T>91 fn try_extend<I>(&mut self, iter: I) -> Result<(), OutOfMemory> 92 where 93 I: IntoIterator<Item = T>; 94 } 95 96 impl<T> TryExtend<T> for TryVec<T> { try_extend<I>(&mut self, iter: I) -> Result<(), OutOfMemory> where I: IntoIterator<Item = 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)] 111 mod tests { 112 use super::{Box, TryCollect, TryExtend, TryVec}; 113 use crate::error::{OutOfMemory, Result}; 114 115 #[test] test_vec_collect() -> Result<(), OutOfMemory>116 fn test_vec_collect() -> Result<(), OutOfMemory> { 117 let v: TryVec<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] test_box_collect() -> Result<(), OutOfMemory>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] test_vec_result_collect() -> Result<()>130 fn test_vec_result_collect() -> Result<()> { 131 let v: Result<TryVec<i32>> = [].into_iter().try_collect(); 132 assert!(v?.is_empty()); 133 134 let v: Result<TryVec<i32>> = [Ok(1), Ok(2)].into_iter().try_collect(); 135 assert_eq!(&*v?, &[1, 2]); 136 137 let v: Result<TryVec<i32>> = [Ok(1), Err(crate::format_err!("hi"))] 138 .into_iter() 139 .try_collect(); 140 assert!(v.is_err()); 141 142 let v: Result<TryVec<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] test_box_result_collect() -> Result<()>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] test_try_extend() -> Result<(), OutOfMemory>170 fn test_try_extend() -> Result<(), OutOfMemory> { 171 let mut vec = TryVec::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 } 180