better-refcell 0.1.1

//! A drop-in replacement for `RefCell` with safe *unborrow* and *reborrow* capabilities.
//!
//!
//! # Features
//!
//! - Allows putting back a borrowed value and reborrowing it within a closure.
//! - Fully compatible with the `RefCell` interface of the stable Rust standard library.
//! - Zero dependencies.
//!
//!
//! # Usage
//!
//! ```rust
//! use better_refcell::BetterRefCell;
//! use std::cell::*;
//!
//! let cell = BetterRefCell::new(42);
//!
//! let mut guard: RefMut<i32> = cell.borrow_mut();
//! let mut_reference: &mut i32 = &mut *guard;
//!
//! let ret = cell.unborrow(mut_reference, || {
//!     let mut guard = cell.borrow_mut();
//!     assert_eq!(*guard, 42);
//!     *guard += 1;
//!     format!("Returns {guard}")
//! });
//!
//! assert_eq!(*guard, 43);
//! assert_eq!(ret, "Returns 43");
//! ```

#![no_std]
#![allow(clippy::missing_safety_doc, clippy::missing_errors_doc)]
use core::cell::{BorrowError, BorrowMutError, Cell, Ref, RefCell, RefMut, UnsafeCell};
use core::ptr::{self, NonNull};
use core::{cmp, mem};

/// An enhanced `RefCell` with double-borrowing support.
pub struct BetterRefCell<T: ?Sized> {
    ptr: Cell<Option<NonNull<T>>>,
    state: Cell<RefCell<()>>,
    value: UnsafeCell<T>,
}

impl<T> BetterRefCell<T> {
    /// Equivalent to [`RefCell::new`].
    #[inline]
    pub const fn new(value: T) -> Self {
        Self {
            ptr: Cell::new(None),
            state: Cell::new(RefCell::new(())),
            value: UnsafeCell::new(value),
        }
    }

    /// Equivalent to [`RefCell::into_inner`].
    #[inline]
    pub fn into_inner(self) -> T {
        self.value.into_inner()
    }

    /// Equivalent to [`RefCell::replace`].
    #[inline]
    #[track_caller]
    pub fn replace(&self, t: T) -> T {
        mem::replace(&mut *self.borrow_mut(), t)
    }

    /// Equivalent to [`RefCell::replace_with`].
    #[inline]
    #[track_caller]
    pub fn replace_with<F: FnOnce(&mut T) -> T>(&self, f: F) -> T {
        let mut_borrow = &mut *self.borrow_mut();
        let replacement = f(mut_borrow);
        mem::replace(mut_borrow, replacement)
    }

    /// Equivalent to [`RefCell::swap`].
    #[inline]
    pub fn swap(&self, other: &Self) {
        mem::swap(&mut *self.borrow_mut(), &mut *other.borrow_mut());
    }
}

impl<T: ?Sized> BetterRefCell<T> {
    /// Equivalent to [`RefCell::borrow`].
    #[inline]
    #[track_caller]
    pub fn borrow(&self) -> Ref<'_, T> {
        let state = unsafe { &*self.state.as_ptr() };
        Ref::map(state.borrow(), |()| unsafe {
            self.as_ptr().as_ref().unwrap_unchecked()
        })
    }

    /// Equivalent to [`RefCell::try_borrow`].
    #[inline]
    pub fn try_borrow(&self) -> Result<Ref<'_, T>, BorrowError> {
        let state = unsafe { &*self.state.as_ptr() };
        Ok(Ref::map(state.try_borrow()?, |()| unsafe {
            self.as_ptr().as_ref().unwrap_unchecked()
        }))
    }

    /// Equivalent to [`RefCell::borrow_mut`].
    #[inline]
    #[track_caller]
    pub fn borrow_mut(&self) -> RefMut<'_, T> {
        let state = unsafe { &*self.state.as_ptr() };
        RefMut::map(state.borrow_mut(), |()| unsafe {
            self.as_ptr().as_mut().unwrap_unchecked()
        })
    }

    /// Equivalent to [`RefCell::try_borrow_mut`].
    #[inline]
    pub fn try_borrow_mut(&self) -> Result<RefMut<'_, T>, BorrowMutError> {
        let state = unsafe { &*self.state.as_ptr() };
        Ok(RefMut::map(state.try_borrow_mut()?, |()| unsafe {
            self.as_ptr().as_mut().unwrap_unchecked()
        }))
    }

    /// Equivalent to [`RefCell::as_ptr`].
    #[inline]
    pub fn as_ptr(&self) -> *mut T {
        if let Some(ptr) = self.ptr.get() {
            ptr.as_ptr()
        } else {
            self.value.get()
        }
    }

    /// Equivalent to [`RefCell::get_mut`].
    #[inline]
    pub fn get_mut(&mut self) -> &mut T {
        self.value.get_mut()
    }

    /// Equivalent to [`RefCell::try_borrow_unguarded`].
    #[inline]
    pub unsafe fn try_borrow_unguarded(&self) -> Result<&T, BorrowError> {
        let state = unsafe { &*self.state.as_ptr() };
        unsafe { state.try_borrow_unguarded()? };
        Ok(unsafe { self.as_ptr().as_ref().unwrap_unchecked() })
    }
}

impl<T: Default> BetterRefCell<T> {
    /// Equivalent to [`RefCell::take`].
    pub fn take(&self) -> T {
        self.replace(Default::default())
    }
}

impl<T: Clone> Clone for BetterRefCell<T> {
    /// Equivalent to [`RefCell::clone`].
    #[inline]
    #[track_caller]
    fn clone(&self) -> Self {
        Self::new(self.borrow().clone())
    }

    /// Equivalent to [`RefCell::clone_from`].
    #[inline]
    #[track_caller]
    fn clone_from(&mut self, source: &Self) {
        self.get_mut().clone_from(&source.borrow());
    }
}

impl<T: Default> Default for BetterRefCell<T> {
    /// Equivalent to [`RefCell::default`].
    #[inline]
    fn default() -> Self {
        Self::new(Default::default())
    }
}

impl<T: ?Sized + PartialEq> PartialEq for BetterRefCell<T> {
    /// Equivalent to [`RefCell::eq`].
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        *self.borrow() == *other.borrow()
    }
}

impl<T: ?Sized + Eq> Eq for BetterRefCell<T> {}

impl<T: ?Sized + PartialOrd> PartialOrd for BetterRefCell<T> {
    /// Equivalent to [`RefCell::partial_cmp`].
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        self.borrow().partial_cmp(&*other.borrow())
    }

    /// Equivalent to [`RefCell::lt`].
    #[inline]
    fn lt(&self, other: &Self) -> bool {
        *self.borrow() < *other.borrow()
    }

    /// Equivalent to [`RefCell::le`].
    #[inline]
    fn le(&self, other: &Self) -> bool {
        *self.borrow() <= *other.borrow()
    }

    /// Equivalent to [`RefCell::gt`].
    #[inline]
    fn gt(&self, other: &Self) -> bool {
        *self.borrow() > *other.borrow()
    }

    /// Equivalent to [`RefCell::ge`].
    #[inline]
    fn ge(&self, other: &Self) -> bool {
        *self.borrow() >= *other.borrow()
    }
}

impl<T: ?Sized + Ord> Ord for BetterRefCell<T> {
    /// Equivalent to [`RefCell::cmp`].
    #[inline]
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        self.borrow().cmp(&*other.borrow())
    }
}

impl<T> From<T> for BetterRefCell<T> {
    /// Equivalent to [`RefCell::from`].
    fn from(t: T) -> Self {
        Self::new(t)
    }
}

impl<T: ?Sized> BetterRefCell<T> {
    /// Temporarily re-grant exclusive access to the internal value, allowing `borrow_mut`/`try_borrow_mut` again in `f`.
    ///
    /// # Panics
    ///
    /// This function panics only if any of the following conditions are violated:
    /// 1. `borrowed` must point to the internal value stored in `self`.
    /// 2. All borrow guards obtained from `self` within `f` must be dropped before `f` returns.
    #[track_caller]
    pub fn unborrow<R>(&self, borrowed: &mut T, f: impl FnOnce() -> R) -> R {
        if !ptr::eq(self.value.get(), borrowed) {
            panic_different_address()
        }
        let ptr = self.ptr.replace(Some(borrowed.into()));
        let state = self.state.take();
        let result = f();
        let state = self.state.replace(state);
        if state.try_borrow_mut().is_ok() {
            self.ptr.set(ptr);
            result
        } else {
            self.state.set(state);
            panic_borrow_guard_leaked()
        }
    }

    /// Temporarily re-grant exclusive access to the internal value, allowing `borrow_mut`/`try_borrow_mut` again in `f`.
    ///
    /// # Safety
    ///
    /// This function is safe to call if all of the following conditions are met:
    /// 1. `borrowed` must point to the internal value stored in `self`.
    /// 2. All borrow guards obtained from `self` within `f` must be dropped before `f` returns.
    /// 3. Any borrow guards obtained from `self` outside of `f` must not be used within `f`.
    #[track_caller]
    pub unsafe fn unborrow_unchecked<R>(&self, borrowed: &mut T, f: impl FnOnce() -> R) -> R {
        let ptr = self.ptr.replace(Some(borrowed.into()));
        let state = self.state.take();
        let result = f();
        self.state.set(state);
        self.ptr.set(ptr);
        result
    }

    /// Temporarily re-grant shared access to the internal value, allowing `borrow`/`try_borrow` again in `f`.
    ///
    /// # Panics
    ///
    /// This function panics only if any of the following conditions are violated:
    /// 1. `borrowed` must point to the internal value stored in `self`.
    /// 2. If `borrowed` is (transitively) originated from `RefMut`, then all `Ref`s obtained from `self` within `f` must be dropped before `f` returns.
    #[track_caller]
    pub fn unborrow_ref<R>(&self, borrowed: &T, f: impl FnOnce() -> R) -> R {
        if !ptr::eq(self.value.get(), borrowed) {
            panic_different_address()
        }
        if unsafe { &mut *self.state.as_ptr() }.try_borrow().is_ok() {
            return f();
        }
        let ptr = self.ptr.replace(Some(borrowed.into()));
        let state = self.state.take();
        let guard = unsafe { (*self.state.as_ptr()).try_borrow().unwrap_unchecked() };
        let result = f();
        drop(guard);
        let state = self.state.replace(state);
        if state.try_borrow_mut().is_ok() {
            self.ptr.set(ptr);
            result
        } else {
            self.state.set(state);
            panic_borrow_guard_leaked()
        }
    }

    /// Temporarily re-grant shared access to the internal value, allowing `borrow`/`try_borrow` again in `f`.
    ///
    /// # Safety
    ///
    /// This function is safe to call if all of the following conditions are met:
    /// 1. `borrowed` must point to the internal value stored in `self`.
    /// 2. If `borrowed` is (transitively) originated from `RefMut`, then all `Ref`s obtained from `self` within `f` must be dropped before `f` returns.
    /// 3. Any borrow guards obtained from `self` outside of `f` must not be used within `f`.
    #[track_caller]
    pub unsafe fn unborrow_ref_unchecked<R>(&self, borrowed: &T, f: impl FnOnce() -> R) -> R {
        if unsafe { &mut *self.state.as_ptr() }.try_borrow().is_ok() {
            return f();
        }
        let ptr = self.ptr.replace(Some(borrowed.into()));
        let state = self.state.take();
        let result = f();
        self.state.set(state);
        self.ptr.set(ptr);
        result
    }
}

#[inline(never)]
#[track_caller]
#[cold]
fn panic_different_address() -> ! {
    panic!("reference is pointing to a different address")
}

#[inline(never)]
#[track_caller]
#[cold]
fn panic_borrow_guard_leaked() -> ! {
    panic!("borrow guard leaked in closure")
}

/// A macro to unborrow multiple instances at once, using either [`BetterRefCell::unborrow`] or [`BetterRefCell::unborrow_ref`].
///
/// The references to be unborrowed must be explicitly prefixed with `&` or `&mut` and should align with the order of their respective sources.
///
/// # Examples
/// ```rust
/// use better_refcell::{BetterRefCell, unborrow_all};
///
/// let cell_1 = BetterRefCell::new(1);
/// let cell_2 = BetterRefCell::new(2);
///
/// let mut var_1 = cell_1.borrow_mut();
/// let mut var_2 = cell_2.borrow_mut();
///
/// unborrow_all!((), (), move || {});
///
/// unborrow_all!(&cell_1, &var_1, || {});
/// unborrow_all!(&cell_1, &mut var_1, || {});
///
/// unborrow_all!((&cell_1,), (&var_1), || {});
/// unborrow_all!((&cell_1,), (&mut var_1), || {});
///
/// unborrow_all!((&cell_1, &cell_2), (&var_1, &mut var_2), || {});
/// unborrow_all!((&cell_1, &cell_2), (&mut var_1, &var_2), || {});
/// ```
#[macro_export]
macro_rules! unborrow_all {
    ((), (), $closure:expr $(,)?) => {
        $closure()
    };

    (($cell:expr, $($cells:expr),+ $(,)?), (&$var:expr, $($vars:tt)+), $closure:expr $(,)?) => {
        $crate::unborrow_all!($cell, &$var, || {
            $crate::unborrow_all!(($($cells),+), ($($vars)+), $closure)
        })
    };
    (($cell:expr, $($cells:expr),+ $(,)?), (&mut $var:expr, $($vars:tt)+), $closure:expr $(,)?) => {
        $crate::unborrow_all!($cell, &mut $var, || {
            $crate::unborrow_all!(($($cells),+), ($($vars)+), $closure)
        })
    };
    (($($cells:expr),+ $(,)?), ($var:expr, $($vars:tt)+), $closure:expr $(,)?) => {
        $crate::unborrow_all!(@error_no_prefix $var)
    };

    (($cell:expr $(,)?), (&$var:expr $(,)?), $closure:expr $(,)?) => {
        $crate::unborrow_all!($cell, &$var, $closure)
    };
    (($cell:expr $(,)?), (&mut $var:expr $(,)?), $closure:expr $(,)?) => {
        $crate::unborrow_all!($cell, &mut $var, $closure)
    };
    (($cell:expr $(,)?), ($var:expr $(,)?), $closure:expr $(,)?) => {
        $crate::unborrow_all!(@error_no_prefix $var)
    };

    ($cell:expr, &$var:expr, $closure:expr $(,)?) => {
        $crate::BetterRefCell::unborrow_ref($cell, &$var, $closure)
    };
    ($cell:expr, &mut $var:expr, $closure:expr $(,)?) => {
        $crate::BetterRefCell::unborrow($cell, &mut $var, $closure)
    };
    ($cell:expr, $var:expr, $closure:expr $(,)?) => {
        $crate::unborrow_all!(@error_no_prefix $var)
    };

    (@error_no_prefix $var:expr) => {
        ::core::compile_error!(::core::concat!(
            "consider prefixing `",
            ::core::stringify!($var),
            "` with `&` or `&mut`",
        ))
    }
}

/// A macro to unborrow multiple instances at once, using either [`BetterRefCell::unborrow_unchecked`] or [`BetterRefCell::unborrow_ref_unchecked`].
///
/// The references to be unborrowed must be explicitly prefixed with `&` or `&mut` and should align with the order of their respective sources.
///
/// # Examples
/// ```rust
/// use better_refcell::{BetterRefCell, unborrow_all_unchecked};
///
/// let cell_1 = BetterRefCell::new(1);
/// let cell_2 = BetterRefCell::new(2);
///
/// let mut var_1 = cell_1.borrow_mut();
/// let mut var_2 = cell_2.borrow_mut();
///
/// unsafe {
///     unborrow_all_unchecked!((), (), move || {});
///
///     unborrow_all_unchecked!(&cell_1, &var_1, || {});
///     unborrow_all_unchecked!(&cell_1, &mut var_1, || {});
///
///     unborrow_all_unchecked!((&cell_1,), (&var_1), || {});
///     unborrow_all_unchecked!((&cell_1,), (&mut var_1), || {});
///
///     unborrow_all_unchecked!((&cell_1, &cell_2), (&var_1, &mut var_2), || {});
///     unborrow_all_unchecked!((&cell_1, &cell_2), (&mut var_1, &var_2), || {});
/// }
/// ```
#[macro_export]
macro_rules! unborrow_all_unchecked {
    ((), (), $closure:expr $(,)?) => {
        $closure()
    };

    (($cell:expr, $($cells:expr),+ $(,)?), (&$var:expr, $($vars:tt)+), $closure:expr $(,)?) => {
        $crate::unborrow_all_unchecked!($cell, &$var, || {
            $crate::unborrow_all_unchecked!(($($cells),+), ($($vars)+), $closure)
        })
    };
    (($cell:expr, $($cells:expr),+ $(,)?), (&mut $var:expr, $($vars:tt)+), $closure:expr $(,)?) => {
        $crate::unborrow_all_unchecked!($cell, &mut $var, || {
            $crate::unborrow_all_unchecked!(($($cells),+), ($($vars)+), $closure)
        })
    };
    (($($cells:expr),+ $(,)?), ($var:expr, $($vars:tt)+), $closure:expr $(,)?) => {
        $crate::unborrow_all_unchecked!(@error_no_prefix $var)
    };

    (($cell:expr $(,)?), (&$var:expr $(,)?), $closure:expr $(,)?) => {
        $crate::unborrow_all_unchecked!($cell, &$var, $closure)
    };
    (($cell:expr $(,)?), (&mut $var:expr $(,)?), $closure:expr $(,)?) => {
        $crate::unborrow_all_unchecked!($cell, &mut $var, $closure)
    };
    (($cell:expr $(,)?), ($var:expr $(,)?), $closure:expr $(,)?) => {
        $crate::unborrow_all_unchecked!(@error_no_prefix $var)
    };

    ($cell:expr, &$var:expr, $closure:expr $(,)?) => {
        $crate::BetterRefCell::unborrow_ref_unchecked($cell, &$var, $closure)
    };
    ($cell:expr, &mut $var:expr, $closure:expr $(,)?) => {
        $crate::BetterRefCell::unborrow_unchecked($cell, &mut $var, $closure)
    };
    ($cell:expr, $var:expr, $closure:expr $(,)?) => {
        $crate::unborrow_all!(@error_no_prefix $var)
    };

    (@error_no_prefix $var:expr) => {
        ::core::compile_error!(::core::concat!(
            "consider prefixing `",
            ::core::stringify!($var),
            "` with `&` or `&mut`",
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn unborrow() {
        let cell = BetterRefCell::new(0);
        let mut guard = cell.borrow_mut();
        cell.unborrow(&mut guard, || {
            let mut guard = cell.borrow_mut();
            assert_eq!(*guard, 0);
            *guard += 1;
        });
        assert_eq!(*guard, 1);
    }

    #[test]
    #[should_panic = "reference is pointing to a different address"]
    fn unborrow_by_invalid_reference() {
        let cell = BetterRefCell::new(0);
        cell.unborrow(&mut 0, || {});
    }

    #[test]
    #[should_panic = "borrow guard leaked"]
    fn unborrow_and_leak_ref_guard() {
        let cell = BetterRefCell::new(0);
        cell.unborrow(&mut cell.borrow_mut(), || cell.borrow());
    }

    #[test]
    #[should_panic = "borrow guard leaked"]
    fn unborrow_and_leak_mut_guard() {
        let cell = BetterRefCell::new(0);
        cell.unborrow(&mut cell.borrow_mut(), || cell.borrow_mut());
    }

    #[test]
    fn unborrow_ref_by_ref() {
        let cell = BetterRefCell::new(0);
        let guard = cell.borrow();
        cell.unborrow_ref(&guard, || {
            let guard = cell.borrow();
            assert_eq!(*guard, 0);
        });
        assert_eq!(*guard, 0);
    }

    #[test]
    fn unborrow_ref_by_mut() {
        let cell = BetterRefCell::new(0);
        let guard = cell.borrow_mut();
        cell.unborrow_ref(&guard, || {
            let guard = cell.borrow();
            assert_eq!(*guard, 0);
        });
        assert_eq!(*guard, 0);
    }

    #[test]
    #[should_panic = "reference is pointing to a different address"]
    fn unborrow_ref_by_invalid_reference() {
        let cell = BetterRefCell::new(0);
        cell.unborrow_ref(&0, || {});
    }

    #[test]
    fn unborrow_ref_by_ref_and_leak_guard() {
        let cell = BetterRefCell::new(0);
        cell.unborrow_ref(&cell.borrow(), || cell.borrow());
    }

    #[test]
    #[should_panic = "borrow guard leaked"]
    fn unborrow_ref_by_mut_and_leak_guard() {
        let cell = BetterRefCell::new(0);
        cell.unborrow_ref(&cell.borrow_mut(), || cell.borrow());
    }

    #[test]
    fn unborrow_ref_and_drop_outer_guard() {
        let cell = BetterRefCell::new(0);
        let outer_guard = cell.borrow();
        cell.unborrow_ref(&cell.borrow(), || {
            drop(outer_guard);
            assert!(cell.try_borrow_mut().is_err());
        });
        assert!(cell.try_borrow_mut().is_ok());
    }
}