swym 0.1.0-preview

Transactional memory for rust
use std::{
    alloc::{Alloc, AllocErr, CannotReallocInPlace, Excess, GlobalAlloc, Layout},
    ptr::NonNull,
    sync::atomic::{AtomicUsize, Ordering::Relaxed},
};

static ALLOC_COUNT: AtomicUsize = AtomicUsize::new(0);

#[derive(Default, Debug)]
pub struct DebugAlloc<T>(pub T);

#[cfg(test)]
pub fn alloc_count() -> usize {
    ALLOC_COUNT.load(Relaxed)
}

unsafe impl<T: GlobalAlloc> GlobalAlloc for DebugAlloc<T> {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let res = self.0.alloc(layout);
        if likely!(!res.is_null()) {
            ALLOC_COUNT.fetch_add(1, Relaxed);
        }
        res
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        ALLOC_COUNT.fetch_sub(1, Relaxed);
        self.0.dealloc(ptr, layout)
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        self.0.realloc(ptr, layout, new_size)
    }

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let res = self.0.alloc_zeroed(layout);
        if likely!(!res.is_null()) {
            ALLOC_COUNT.fetch_add(1, Relaxed);
        }
        res
    }
}

unsafe impl<'a, T: Alloc> Alloc for DebugAlloc<T> {
    #[inline]
    unsafe fn alloc(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
        let res = self.0.alloc(layout)?;
        ALLOC_COUNT.fetch_add(1, Relaxed);
        Ok(res)
    }

    #[inline]
    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        ALLOC_COUNT.fetch_sub(1, Relaxed);
        self.0.dealloc(ptr, layout)
    }

    #[inline]
    fn usable_size(&self, layout: &Layout) -> (usize, usize) {
        self.0.usable_size(layout)
    }

    #[inline]
    unsafe fn realloc(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
    ) -> Result<NonNull<u8>, AllocErr> {
        self.0.realloc(ptr, layout, new_size)
    }

    #[inline]
    unsafe fn alloc_zeroed(&mut self, layout: Layout) -> Result<NonNull<u8>, AllocErr> {
        let res = self.0.alloc_zeroed(layout)?;
        ALLOC_COUNT.fetch_add(1, Relaxed);
        Ok(res)
    }

    #[inline]
    unsafe fn alloc_excess(&mut self, layout: Layout) -> Result<Excess, AllocErr> {
        let res = self.0.alloc_excess(layout)?;
        ALLOC_COUNT.fetch_add(1, Relaxed);
        Ok(res)
    }

    #[inline]
    unsafe fn realloc_excess(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
    ) -> Result<Excess, AllocErr> {
        self.0.realloc_excess(ptr, layout, new_size)
    }

    #[inline]
    unsafe fn grow_in_place(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
    ) -> Result<(), CannotReallocInPlace> {
        self.0.grow_in_place(ptr, layout, new_size)
    }

    #[inline]
    unsafe fn shrink_in_place(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
    ) -> Result<(), CannotReallocInPlace> {
        self.0.shrink_in_place(ptr, layout, new_size)
    }

    #[inline]
    fn alloc_one<U>(&mut self) -> Result<NonNull<U>, AllocErr> {
        let res = self.0.alloc_one::<U>()?;
        ALLOC_COUNT.fetch_add(1, Relaxed);
        Ok(res)
    }

    #[inline]
    unsafe fn dealloc_one<U>(&mut self, ptr: NonNull<U>) {
        ALLOC_COUNT.fetch_sub(1, Relaxed);
        self.0.dealloc_one(ptr)
    }

    #[inline]
    fn alloc_array<U>(&mut self, n: usize) -> Result<NonNull<U>, AllocErr> {
        let res = self.0.alloc_array::<U>(n)?;
        ALLOC_COUNT.fetch_add(1, Relaxed);
        Ok(res)
    }

    #[inline]
    unsafe fn realloc_array<U>(
        &mut self,
        ptr: NonNull<U>,
        n_old: usize,
        n_new: usize,
    ) -> Result<NonNull<U>, AllocErr> {
        self.0.realloc_array::<U>(ptr, n_old, n_new)
    }

    #[inline]
    unsafe fn dealloc_array<U>(&mut self, ptr: NonNull<U>, n: usize) -> Result<(), AllocErr> {
        ALLOC_COUNT.fetch_sub(1, Relaxed);
        self.0.dealloc_array::<U>(ptr, n)
    }
}