memkit-async 0.2.0-beta.1

Async-aware memory allocators for memkit
//! Async-safe container types.

use std::ops::{Deref, DerefMut};
use std::ptr::NonNull;

use crate::allocator::MkAsyncFrameAlloc;

/// Async-safe box type.
pub struct MkAsyncBox<T> {
    ptr: NonNull<T>,
    alloc: MkAsyncFrameAlloc,
}

impl<T> MkAsyncBox<T> {
    /// Create a new async box.
    pub async fn new(alloc: MkAsyncFrameAlloc, value: T) -> Option<Self> {
        let ptr = alloc.alloc::<T>().await?;
        unsafe {
            ptr.write(value);
        }
        Some(Self {
            ptr: NonNull::new(ptr)?,
            alloc,
        })
    }

    /// Get the raw pointer.
    pub fn as_ptr(&self) -> *const T {
        self.ptr.as_ptr()
    }

    /// Get the mutable raw pointer.
    pub fn as_mut_ptr(&mut self) -> *mut T {
        self.ptr.as_ptr()
    }
}

impl<T> Deref for MkAsyncBox<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        unsafe { self.ptr.as_ref() }
    }
}

impl<T> DerefMut for MkAsyncBox<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { self.ptr.as_mut() }
    }
}

impl<T> Drop for MkAsyncBox<T> {
    fn drop(&mut self) {
        unsafe {
            std::ptr::drop_in_place(self.ptr.as_ptr());
        }
        self.alloc.release();
    }
}

// Safety: MkAsyncBox is Send/Sync if T is
unsafe impl<T: Send> Send for MkAsyncBox<T> {}
unsafe impl<T: Sync> Sync for MkAsyncBox<T> {}

/// Async-safe vector type.
pub struct MkAsyncVec<T> {
    ptr: *mut T,
    len: usize,
    capacity: usize,
    alloc: MkAsyncFrameAlloc,
}

impl<T> MkAsyncVec<T> {
    /// Create a new async vec with capacity.
    pub async fn with_capacity(alloc: MkAsyncFrameAlloc, capacity: usize) -> Option<Self> {
        if capacity == 0 {
            return Some(Self {
                ptr: std::ptr::null_mut(),
                len: 0,
                capacity: 0,
                alloc,
            });
        }

        let layout = std::alloc::Layout::array::<T>(capacity).ok()?;
        let ptr = alloc.alloc_layout(layout).await? as *mut T;
        
        Some(Self {
            ptr,
            len: 0,
            capacity,
            alloc,
        })
    }

    /// Push a value onto the vec.
    pub fn push(&mut self, value: T) -> Result<(), T> {
        if self.len >= self.capacity {
            return Err(value);
        }
        unsafe {
            self.ptr.add(self.len).write(value);
        }
        self.len += 1;
        Ok(())
    }

    /// Get the length.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Check if empty.
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Get the capacity.
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// Clear the vec.
    pub fn clear(&mut self) {
        for i in 0..self.len {
            unsafe {
                std::ptr::drop_in_place(self.ptr.add(i));
            }
        }
        self.len = 0;
    }
}

impl<T> Deref for MkAsyncVec<T> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        if self.ptr.is_null() {
            &[]
        } else {
            unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
        }
    }
}

impl<T> DerefMut for MkAsyncVec<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        if self.ptr.is_null() {
            &mut []
        } else {
            unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
        }
    }
}

impl<T> Drop for MkAsyncVec<T> {
    fn drop(&mut self) {
        self.clear();
        if self.capacity > 0 {
            self.alloc.release();
        }
    }
}

unsafe impl<T: Send> Send for MkAsyncVec<T> {}
unsafe impl<T: Sync> Sync for MkAsyncVec<T> {}