use crate::{ io::{self, Write}, platform::{self, types::*, WriteByte}, }; use core::{ cmp, fmt, iter::IntoIterator, mem, ops::{Deref, DerefMut}, ptr::{self, NonNull}, slice, }; /// Error that occurs when an allocation fails #[derive(Debug, Default, Hash, PartialEq, Eq, Clone, Copy)] pub struct AllocError; /// A normal vector allocated in Rust needs to be dropped from Rust /// too, in order to avoid UB. This CVec is an abstraction that works /// using only C allocations functions and can therefore be dropped /// from C. Just like the Rust Vec, this does bounds checks to assure /// you never reach isize::MAX. Unless you need to drop something from /// C, prefer Rust's builtin Vec. pub struct CVec { ptr: NonNull, len: usize, cap: usize, } impl CVec { pub fn new() -> Self { Self { ptr: NonNull::dangling(), len: 0, cap: 0, } } fn check_bounds(i: usize) -> Result { if i > core::isize::MAX as usize { Err(AllocError) } else { Ok(i) } } fn check_mul(x: usize, y: usize) -> Result { x.checked_mul(y) .ok_or(AllocError) .and_then(Self::check_bounds) } pub fn with_capacity(cap: usize) -> Result { if cap == 0 { return Ok(Self::new()); } let size = Self::check_mul(cap, mem::size_of::())?; let ptr = NonNull::new(unsafe { platform::alloc(size) as *mut T }).ok_or(AllocError)?; Ok(Self { ptr, len: 0, cap }) } unsafe fn resize(&mut self, cap: usize) -> Result<(), AllocError> { let size = Self::check_mul(cap, mem::size_of::())?; let ptr = if cap == 0 { NonNull::dangling() } else if self.cap > 0 { NonNull::new(platform::realloc(self.ptr.as_ptr() as *mut c_void, size) as *mut T) .ok_or(AllocError)? } else { NonNull::new((platform::alloc(size)) as *mut T).ok_or(AllocError)? }; self.ptr = ptr; self.cap = cap; Ok(()) } unsafe fn drop_range(&mut self, start: usize, end: usize) { let mut start = self.ptr.as_ptr().add(start); let end = self.ptr.as_ptr().add(end); while start < end { ptr::drop_in_place(start); start = start.add(1); } } // Push stuff pub fn reserve(&mut self, required: usize) -> Result<(), AllocError> { let required_len = self .len .checked_add(required) .ok_or(AllocError) .and_then(Self::check_bounds)?; if required_len > self.cap { let new_cap = cmp::min(required_len.next_power_of_two(), core::isize::MAX as usize); unsafe { self.resize(new_cap)?; } } Ok(()) } pub fn push(&mut self, elem: T) -> Result<(), AllocError> { self.reserve(1)?; unsafe { ptr::write(self.ptr.as_ptr().add(self.len), elem); } self.len += 1; // no need to bounds check, as new len <= cap Ok(()) } pub fn extend_from_slice(&mut self, elems: &[T]) -> Result<(), AllocError> where T: Copy, { self.reserve(elems.len())?; unsafe { ptr::copy_nonoverlapping(elems.as_ptr(), self.ptr.as_ptr().add(self.len), elems.len()); } self.len += elems.len(); // no need to bounds check, as new len <= cap Ok(()) } pub fn append(&mut self, other: &mut Self) -> Result<(), AllocError> { let len = other.len; other.len = 0; // move self.reserve(len)?; unsafe { ptr::copy_nonoverlapping(other.as_ptr(), self.ptr.as_ptr().add(self.len), len); } self.len += other.len(); // no need to bounds check, as new len <= cap Ok(()) } // Pop stuff pub fn truncate(&mut self, len: usize) { if len < self.len { unsafe { let old_len = self.len; self.drop_range(len, old_len); } self.len = len; } } pub fn shrink_to_fit(&mut self) -> Result<(), AllocError> { if self.len < self.cap { unsafe { let new_cap = self.len; self.resize(new_cap)?; } } Ok(()) } pub fn pop(&mut self) -> Option { if self.is_empty() { None } else { let elem = unsafe { ptr::read(self.as_ptr().add(self.len - 1)) }; self.len -= 1; Some(elem) } } // Misc stuff pub fn capacity(&self) -> usize { self.cap } pub fn as_ptr(&self) -> *const T { self.ptr.as_ptr() } pub fn as_mut_ptr(&mut self) -> *mut T { self.ptr.as_ptr() } /// Leaks the inner data. This is safe to drop from C! pub fn leak(mut self) -> *mut T { let ptr = self.as_mut_ptr(); mem::forget(self); ptr } } impl Deref for CVec { type Target = [T]; fn deref(&self) -> &Self::Target { unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) } } } impl DerefMut for CVec { fn deref_mut(&mut self) -> &mut Self::Target { unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) } } } impl Drop for CVec { fn drop(&mut self) { unsafe { let len = self.len; self.drop_range(0, len); } } } impl<'a, T> IntoIterator for &'a CVec { type Item = <&'a [T] as IntoIterator>::Item; type IntoIter = <&'a [T] as IntoIterator>::IntoIter; fn into_iter(self) -> Self::IntoIter { <&[T]>::into_iter(&*self) } } impl<'a, T> IntoIterator for &'a mut CVec { type Item = <&'a mut [T] as IntoIterator>::Item; type IntoIter = <&'a mut [T] as IntoIterator>::IntoIter; fn into_iter(self) -> Self::IntoIter { <&mut [T]>::into_iter(&mut *self) } } impl Write for CVec { fn write(&mut self, buf: &[u8]) -> io::Result { self.extend_from_slice(buf).map_err(|err| { io::Error::new( io::ErrorKind::Other, "AllocStringWriter::write failed to allocate", ) })?; Ok(buf.len()) } fn flush(&mut self) -> io::Result<()> { Ok(()) } } impl fmt::Write for CVec { fn write_str(&mut self, s: &str) -> fmt::Result { self.write(s.as_bytes()).map_err(|_| fmt::Error)?; Ok(()) } } impl WriteByte for CVec { fn write_u8(&mut self, byte: u8) -> fmt::Result { self.write(&[byte]).map_err(|_| fmt::Error)?; Ok(()) } } #[cfg(test)] mod tests { use super::CVec; #[test] fn push_pop() { let mut vec = CVec::new(); vec.push(1).unwrap(); vec.push(2).unwrap(); vec.push(3).unwrap(); assert_eq!(&vec[..], &[1, 2, 3]); assert_eq!(vec.pop().unwrap(), 3); assert_eq!(&vec[..], &[1, 2]); } #[test] fn extend_from_slice() { use core_io::Write; let mut vec = CVec::new(); vec.extend_from_slice(&[1, 2, 3]).unwrap(); vec.extend_from_slice(&[4, 5, 6]).unwrap(); assert_eq!(&vec[..], &[1, 2, 3, 4, 5, 6]); assert_eq!(vec.write(&[7, 8, 9]).unwrap(), 3); assert_eq!(&vec[..], &[1, 2, 3, 4, 5, 6, 7, 8, 9]); } #[test] fn dropped() { use alloc::rc::Rc; let counter = Rc::new(()); let mut vec = CVec::with_capacity(3).unwrap(); vec.push(Rc::clone(&counter)).unwrap(); vec.push(Rc::clone(&counter)).unwrap(); vec.push(Rc::clone(&counter)).unwrap(); assert_eq!(Rc::strong_count(&counter), 4); let popped = vec.pop().unwrap(); assert_eq!(Rc::strong_count(&counter), 4); drop(popped); assert_eq!(Rc::strong_count(&counter), 3); vec.push(Rc::clone(&counter)).unwrap(); vec.push(Rc::clone(&counter)).unwrap(); vec.push(Rc::clone(&counter)).unwrap(); assert_eq!(vec.len(), 5); assert_eq!(Rc::strong_count(&counter), 6); vec.truncate(1); assert_eq!(Rc::strong_count(&counter), 2); drop(vec); assert_eq!(Rc::strong_count(&counter), 1); } }