Files
RedBear-OS/src/c_vec.rs
T

295 lines
8.4 KiB
Rust

//! Equivalent of Rust's `Vec<T>`, but using relibc's own allocator.
use crate::{
io::{self, Write},
platform::{self, WriteByte, types::*},
};
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<T> {
ptr: NonNull<T>,
len: usize,
cap: usize,
}
impl<T> CVec<T> {
#[allow(clippy::new_without_default)]
pub fn new() -> Self {
Self {
ptr: NonNull::dangling(),
len: 0,
cap: 0,
}
}
fn check_bounds(i: usize) -> Result<usize, AllocError> {
if i > isize::MAX as usize {
Err(AllocError)
} else {
Ok(i)
}
}
fn check_mul(x: usize, y: usize) -> Result<usize, AllocError> {
x.checked_mul(y)
.ok_or(AllocError)
.and_then(Self::check_bounds)
}
pub fn with_capacity(cap: usize) -> Result<Self, AllocError> {
if cap == 0 {
return Ok(Self::new());
}
let size = Self::check_mul(cap, mem::size_of::<T>())?;
let ptr = NonNull::new(unsafe { platform::alloc(size).cast::<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::<T>())?;
let ptr = if cap == 0 {
NonNull::dangling()
} else if self.cap > 0 {
NonNull::new(
unsafe { platform::realloc(self.ptr.as_ptr().cast::<c_void>(), size) }.cast::<T>(),
)
.ok_or(AllocError)?
} else {
NonNull::new((unsafe { platform::alloc(size) }).cast::<T>()).ok_or(AllocError)?
};
self.ptr = ptr;
self.cap = cap;
Ok(())
}
unsafe fn drop_range(&mut self, start: usize, end: usize) {
let mut start = unsafe { self.ptr.as_ptr().add(start) };
let end = unsafe { self.ptr.as_ptr().add(end) };
while start < end {
unsafe { ptr::drop_in_place(start) };
start = unsafe { 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(), 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<T> {
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<T> Deref for CVec<T> {
type Target = [T];
fn deref(&self) -> &Self::Target {
unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
}
}
impl<T> DerefMut for CVec<T> {
fn deref_mut(&mut self) -> &mut Self::Target {
unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
}
}
impl<T> Drop for CVec<T> {
fn drop(&mut self) {
unsafe {
let len = self.len;
self.drop_range(0, len);
}
}
}
impl<'a, T> IntoIterator for &'a CVec<T> {
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<T> {
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<u8> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
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<u8> {
fn write_str(&mut self, s: &str) -> fmt::Result {
self.write(s.as_bytes()).map_err(|_| fmt::Error)?;
Ok(())
}
}
impl WriteByte for CVec<u8> {
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 crate::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);
}
}