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 // * 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 196 impl<T> Deref for Vec<T> { 197 type Target = [T]; 198 199 fn deref(&self) -> &Self::Target { 200 &self.inner 201 } 202 } 203 204 impl<T> DerefMut for Vec<T> { 205 fn deref_mut(&mut self) -> &mut Self::Target { 206 &mut self.inner 207 } 208 } 209 210 impl<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 218 impl<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 224 impl<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 233 impl<'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 243 impl<'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 253 impl<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)] 262 mod 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 } 293