1 use crate::alloc::{TryClone, try_realloc}; 2 use crate::error::OutOfMemory; 3 use core::{ 4 fmt, mem, 5 ops::{Deref, DerefMut, Index, IndexMut}, 6 }; 7 use std_alloc::alloc::Layout; 8 use std_alloc::boxed::Box; 9 use 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)] 14 pub struct Vec<T> { 15 inner: StdVec<T>, 16 } 17 18 impl<T> Default for Vec<T> { 19 fn default() -> Self { 20 Self { 21 inner: Default::default(), 22 } 23 } 24 } 25 26 impl<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 32 impl<T> TryClone for Vec<T> 33 where 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 45 impl<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 // 154 // In all of these cases delegate to the standard library's 155 // `into_boxed_slice` which is guaranteed to not perform a `realloc`. 156 if self.is_empty() || mem::size_of::<T>() == 0 { 157 return Ok(self.inner.into_boxed_slice()); 158 } 159 160 let (ptr, len, cap) = self.into_raw_parts(); 161 let layout = Layout::array::<T>(cap).unwrap(); 162 let new_len = Layout::array::<T>(len).unwrap().size(); 163 164 // SAFETY: `ptr` was previously allocated in the global allocator, 165 // `layout` has a nonzero size and matches the current allocation of 166 // `ptr`, `new_size` is nonzero, and `new_size` is a valid array size 167 // for `len` elements given its constructor. 168 let result = unsafe { try_realloc(ptr.cast(), layout, new_len) }; 169 170 match result { 171 Ok(ptr) => { 172 // SAFETY: `result` is allocated with the global allocator with 173 // an appropriate size/align to create this `Box` with. 174 unsafe { 175 Ok(Box::from_raw(core::ptr::slice_from_raw_parts_mut( 176 ptr.as_ptr().cast(), 177 len, 178 ))) 179 } 180 } 181 Err(oom) => { 182 // SAFETY: If reallocation fails then it's guaranteed that the 183 // original allocation is not tampered with, so it's safe to 184 // reassemble it back into the original vector. 185 unsafe { 186 let _ = Vec::from_raw_parts(ptr, len, cap); 187 } 188 Err(oom) 189 } 190 } 191 } 192 } 193 194 impl<T> Deref for Vec<T> { 195 type Target = [T]; 196 197 fn deref(&self) -> &Self::Target { 198 &self.inner 199 } 200 } 201 202 impl<T> DerefMut for Vec<T> { 203 fn deref_mut(&mut self) -> &mut Self::Target { 204 &mut self.inner 205 } 206 } 207 208 impl<T> Index<usize> for Vec<T> { 209 type Output = T; 210 211 fn index(&self, index: usize) -> &Self::Output { 212 &self.inner[index] 213 } 214 } 215 216 impl<T> IndexMut<usize> for Vec<T> { 217 fn index_mut(&mut self, index: usize) -> &mut Self::Output { 218 &mut self.inner[index] 219 } 220 } 221 222 impl<T> IntoIterator for Vec<T> { 223 type Item = T; 224 type IntoIter = std_alloc::vec::IntoIter<T>; 225 226 fn into_iter(self) -> Self::IntoIter { 227 self.inner.into_iter() 228 } 229 } 230 231 impl<'a, T> IntoIterator for &'a Vec<T> { 232 type Item = &'a T; 233 234 type IntoIter = core::slice::Iter<'a, T>; 235 236 fn into_iter(self) -> Self::IntoIter { 237 (**self).iter() 238 } 239 } 240 241 impl<'a, T> IntoIterator for &'a mut Vec<T> { 242 type Item = &'a mut T; 243 244 type IntoIter = core::slice::IterMut<'a, T>; 245 246 fn into_iter(self) -> Self::IntoIter { 247 (**self).iter_mut() 248 } 249 } 250 251 impl<T> From<Box<[T]>> for Vec<T> { 252 fn from(boxed_slice: Box<[T]>) -> Self { 253 Vec { 254 inner: StdVec::from(boxed_slice), 255 } 256 } 257 } 258 259 #[cfg(test)] 260 mod tests { 261 use super::Vec; 262 use crate::error::OutOfMemory; 263 264 #[test] 265 fn test_into_boxed_slice() -> Result<(), OutOfMemory> { 266 assert_eq!(*Vec::<i32>::new().into_boxed_slice()?, []); 267 268 let mut vec = Vec::new(); 269 vec.push(1)?; 270 assert_eq!(*vec.into_boxed_slice()?, [1]); 271 272 let mut vec = Vec::with_capacity(2)?; 273 vec.push(1)?; 274 assert_eq!(*vec.into_boxed_slice()?, [1]); 275 276 let mut vec = Vec::with_capacity(2)?; 277 vec.push(1_u128)?; 278 assert_eq!(*vec.into_boxed_slice()?, [1]); 279 280 assert_eq!(*Vec::<()>::new().into_boxed_slice()?, []); 281 282 let mut vec = Vec::new(); 283 vec.push(())?; 284 assert_eq!(*vec.into_boxed_slice()?, [()]); 285 286 let vec = Vec::<i32>::with_capacity(2)?; 287 assert_eq!(*vec.into_boxed_slice()?, []); 288 Ok(()) 289 } 290 } 291