hipool 0.3.4

RUST Memory Pool
Documentation
use crate::{Allocator, Boxed, GenericAlloc, PoolAlloc, Result};
use core::alloc::Layout;
use core::any::Any;
use core::cell::Cell;
use core::cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd};
use core::fmt::{self, Debug, Display, Formatter};
use core::hash::{Hash, Hasher};
use core::marker::PhantomData;
use core::mem::ManuallyDrop;
use core::ops::Deref;
use core::ptr::{self, NonNull};

#[repr(C)]
struct RcInner<'a, A: Allocator> {
    cnt: Cell<isize>,
    data_layout: Layout,
    data_alloc: &'a A,
    alloc: &'a A,
}

#[repr(C)]
pub struct Rc<'a, T: ?Sized + 'a, A: Allocator + 'a = PoolAlloc> {
    inner: NonNull<RcInner<'a, A>>,
    data: NonNull<T>,
    mark: PhantomData<(RcInner<'a, A>, T)>,
}

impl<T: ?Sized> Rc<'static, T, PoolAlloc> {
    pub fn from_boxed(data: Boxed<'static, T, PoolAlloc>) -> Result<Self> {
        Self::from_boxed_in(&PoolAlloc, data)
    }
}

impl<T: Sized + 'static> Rc<'static, T, PoolAlloc> {
    pub fn new(val: T) -> Result<Self> {
        Rc::new_in(&PoolAlloc, val)
    }
}

impl<T: Clone + 'static> Rc<'static, [T], PoolAlloc> {
    pub fn new_slice(len: usize, val: T) -> Result<Self> {
        Rc::new_slice_in(&PoolAlloc, len, val)
    }
}

impl Rc<'static, (), PoolAlloc> {
    pub fn new_then<T, F>(f: F) -> Result<Rc<'static, T, PoolAlloc>>
    where
        F: FnOnce() -> Result<T>,
    {
        Rc::new_then_in(&PoolAlloc, f)
    }
    pub fn new_slice_then<T, F>(len: usize, f: F) -> Result<Rc<'static, [T], PoolAlloc>>
    where
        F: FnMut(usize) -> Result<T>,
    {
        Rc::new_slice_then_in(&PoolAlloc, len, f)
    }
}

impl<'a, T: ?Sized, A: Allocator + 'a> Rc<'a, T, A> {
    /// 不能使用data内部的alloc,无法保证A支持Sync. 因此要求显示传递alloc.
    pub fn from_boxed_in(alloc: &'a A, data: Boxed<'a, T, A>) -> Result<Self> {
        let rc = Boxed::new_in(
            alloc,
            RcInner {
                cnt: Cell::new(1),
                data_layout: data.layout(),
                data_alloc: data.allocator(),
                alloc,
            },
        )?;
        Ok(Self {
            inner: rc.leak().0.into(),
            data: data.leak().0.into(),
            mark: PhantomData,
        })
    }

    pub fn as_ptr(&self) -> *const T {
        self.data.as_ptr()
    }

    pub fn get_mut(&mut self) -> Option<&mut T> {
        if self.get_inner().cnt.get() == 1 {
            Some(unsafe { self.get_mut_unchecked() })
        } else {
            None
        }
    }

    /// # Safety
    /// Rc使用者保证可写借用和只读借用并存,或者多个可写借用是安全的
    /// Rc不支持Send/Sync,因此无需考虑并发冲突的风险
    pub unsafe fn get_mut_unchecked(&mut self) -> &mut T {
        self.data.as_mut()
    }

    fn get_inner(&self) -> &RcInner<'a, A> {
        unsafe { self.inner.as_ref() }
    }
}

impl<'a, A: Allocator> Rc<'a, dyn Any, A> {
    pub fn downcast<T: Any>(self) -> core::result::Result<Rc<'a, T, A>, Self> {
        if <dyn Any>::is::<T>(self.as_ref()) {
            let data = self.data.cast::<T>();
            let this = ManuallyDrop::new(self);
            Ok(Rc {
                inner: this.inner,
                data,
                mark: PhantomData,
            })
        } else {
            Err(self)
        }
    }
}

impl<'a, T: Any, A: Allocator> Rc<'a, T, A> {
    pub fn to_any(self) -> Rc<'a, dyn Any, A> {
        let any: &dyn Any = self.as_ref();
        let data = NonNull::from(any);
        let other = ManuallyDrop::new(self);
        Rc {
            inner: other.inner,
            data,
            mark: PhantomData,
        }
    }
}

impl<'a, T: ?Sized, A: Allocator> Rc<'a, T, A> {
    /// 用于到某个trait的转换.
    pub fn upcast<U: ?Sized>(
        self,
        f: impl FnOnce(&T) -> &U,
    ) -> core::result::Result<Rc<'a, U, A>, Self> {
        let data = NonNull::from(f(unsafe { self.data.as_ref() }));
        if !ptr::eq(data.cast::<u8>().as_ptr(), self.data.cast::<u8>().as_ptr()) {
            return Err(self);
        }
        let this = ManuallyDrop::new(self);
        Ok(Rc {
            inner: this.inner,
            data,
            mark: PhantomData,
        })
    }

    /// 用于trait到具体类的转换.
    pub unsafe fn cast_unchecked<U>(self) -> Rc<'a, U, A> {
        let this = ManuallyDrop::new(self);
        Rc {
            inner: this.inner,
            data: this.data.cast::<U>(),
            mark: PhantomData,
        }
    }
}

impl<'a, T: ?Sized, A: Allocator + 'a> Drop for Rc<'a, T, A> {
    fn drop(&mut self) {
        let inner = self.get_inner();
        let cnt = inner.cnt.get();
        if cnt > 1 {
            inner.cnt.set(cnt - 1);
        } else {
            unsafe { self.inner.as_ptr().drop_in_place() };
            unsafe { self.data.as_ptr().drop_in_place() };
            unsafe { inner.data_alloc.release_with(self.data, inner.data_layout) };
            let layout = Layout::new::<RcInner<'a, A>>();
            unsafe { inner.alloc.release_with(self.inner, layout) };
        }
    }
}

impl<'a, T: Sized + 'a, A: Allocator + 'a> Rc<'a, T, A> {
    pub fn new_in(alloc: &'a A, val: T) -> Result<Self> {
        let data = Boxed::new_in(alloc, val)?;
        Rc::from_boxed_in(alloc, data)
    }
}

impl<'a, T: Clone + 'a, A: Allocator + 'a> Rc<'a, [T], A> {
    pub fn new_slice_in(alloc: &'a A, len: usize, val: T) -> Result<Self> {
        let data = Boxed::new_slice_in(alloc, len, val)?;
        Rc::from_boxed_in(alloc, data)
    }
}

impl<'a, A: Allocator + 'a> Rc<'a, (), A> {
    pub fn new_then_in<T, F>(alloc: &'a A, f: F) -> Result<Rc<'a, T, A>>
    where
        F: FnOnce() -> Result<T>,
    {
        let data = Boxed::new_then_in(alloc, f)?;
        Rc::from_boxed_in(alloc, data)
    }
    pub fn new_slice_then_in<T, F>(alloc: &'a A, len: usize, f: F) -> Result<Rc<'a, [T], A>>
    where
        F: FnMut(usize) -> Result<T>,
    {
        let data = Boxed::new_slice_then_in(alloc, len, f)?;
        Rc::from_boxed_in(alloc, data)
    }
}

impl<T: ?Sized, A: Allocator> AsRef<T> for Rc<'_, T, A> {
    #[inline(always)]
    fn as_ref(&self) -> &T {
        unsafe { self.data.as_ref() }
    }
}

impl<'a, T: ?Sized, A: Allocator> Clone for Rc<'a, T, A> {
    fn clone(&self) -> Self {
        let inner = self.get_inner();
        inner.cnt.set(inner.cnt.get() + 1);
        Self {
            inner: self.inner,
            data: self.data,
            mark: PhantomData,
        }
    }
}

impl<T: Unpin + ?Sized, A: Allocator> Unpin for Rc<'_, T, A> {}

impl<T: ?Sized, A: Allocator> Deref for Rc<'_, T, A> {
    type Target = T;
    #[inline]
    fn deref(&self) -> &Self::Target {
        self.as_ref()
    }
}

impl<T: Display + ?Sized, A: Allocator> Display for Rc<'_, T, A> {
    #[inline]
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        Display::fmt(&**self, f)
    }
}

impl<T: Debug + ?Sized, A: Allocator> Debug for Rc<'_, T, A> {
    #[inline]
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        Debug::fmt(&**self, f)
    }
}

impl<T: ?Sized, A: Allocator> fmt::Pointer for Rc<'_, T, A> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let ptr = self.as_ref() as *const _;
        fmt::Pointer::fmt(&ptr, f)
    }
}

impl<T: PartialEq + ?Sized, A: Allocator> PartialEq for Rc<'_, T, A> {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        PartialEq::eq(&**self, &**other)
    }
}

impl<T: PartialOrd + ?Sized, A: Allocator> PartialOrd for Rc<'_, T, A> {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        PartialOrd::partial_cmp(&**self, &**other)
    }
    #[inline]
    fn lt(&self, other: &Self) -> bool {
        PartialOrd::lt(&**self, &**other)
    }
    #[inline]
    fn le(&self, other: &Self) -> bool {
        PartialOrd::le(&**self, &**other)
    }
    #[inline]
    fn gt(&self, other: &Self) -> bool {
        PartialOrd::gt(&**self, &**other)
    }
    #[inline]
    fn ge(&self, other: &Self) -> bool {
        PartialOrd::ge(&**self, &**other)
    }
}

impl<T: Ord + ?Sized, A: Allocator> Ord for Rc<'_, T, A> {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        Ord::cmp(&**self, &**other)
    }
}

impl<T: Eq + ?Sized, A: Allocator> Eq for Rc<'_, T, A> {}

impl<T: Hash + ?Sized, A: Allocator> Hash for Rc<'_, T, A> {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        Hash::hash(&**self, state)
    }
}

#[cfg(test)]
mod test {
    use crate::{MemPool, Rc};
    #[test]
    fn test() {
        let rc = Rc::new(100).unwrap();

        let cnt = *rc + 100;
        assert_eq!(200, cnt);

        let mut rc2 = rc.clone();
        let ret = rc2.get_mut();
        assert_eq!(ret, None);
    }

    #[test]
    fn test_pool() {
        static mut DROP: usize = 0;
        struct Foo;
        impl Drop for Foo {
            fn drop(&mut self) {
                unsafe { DROP += 1 };
            }
        }
        let pool = MemPool::new_boxed(0).unwrap();
        unsafe { DROP = 0 };
        {
            let rc = Rc::new_in(pool.as_ref(), Foo).unwrap();
            {
                let _rc1 = rc.clone();
            }
            unsafe { assert_eq!(DROP, 0) };
        }
        unsafe { assert_eq!(DROP, 1) };
    }

    #[test]
    fn test_cast() {
        trait Trait {}
        impl Trait for i8 {}

        let rc = Rc::new(100_i8).unwrap();
        let rc = rc.to_any();
        let rc = rc.downcast::<i8>().unwrap();
        let rc = rc.upcast::<dyn Trait>(|val| val);
        assert!(rc.is_ok());
        let _ = unsafe { rc.unwrap().cast_unchecked::<i8>() };
    }
}