better_refcell/
lib.rs

1//! A drop-in replacement for `RefCell` with safe *unborrow* and *reborrow* capabilities.
2//!
3//!
4//! # Features
5//!
6//! - Allows putting back a borrowed value and reborrowing it within a closure.
7//! - Fully compatible with the `RefCell` interface of the stable Rust standard library.
8//! - Zero dependencies.
9//!
10//!
11//! # Usage
12//!
13//! ```rust
14//! use better_refcell::BetterRefCell;
15//! use std::cell::*;
16//!
17//! let cell = BetterRefCell::new(42);
18//!
19//! let mut guard: RefMut<i32> = cell.borrow_mut();
20//! let mut_reference: &mut i32 = &mut *guard;
21//!
22//! let ret = cell.unborrow(mut_reference, || {
23//!     let mut guard = cell.borrow_mut();
24//!     assert_eq!(*guard, 42);
25//!     *guard += 1;
26//!     format!("Returns {guard}")
27//! });
28//!
29//! assert_eq!(*guard, 43);
30//! assert_eq!(ret, "Returns 43");
31//! ```
32
33#![no_std]
34#![allow(clippy::missing_safety_doc, clippy::missing_errors_doc)]
35use core::cell::{BorrowError, BorrowMutError, Cell, Ref, RefCell, RefMut, UnsafeCell};
36use core::ptr::{self, NonNull};
37use core::{cmp, mem};
38
39/// An enhanced `RefCell` with double-borrowing support.
40pub struct BetterRefCell<T: ?Sized> {
41    ptr: Cell<Option<NonNull<T>>>,
42    state: Cell<RefCell<()>>,
43    value: UnsafeCell<T>,
44}
45
46impl<T> BetterRefCell<T> {
47    /// Equivalent to [`RefCell::new`].
48    #[inline]
49    pub const fn new(value: T) -> Self {
50        Self {
51            ptr: Cell::new(None),
52            state: Cell::new(RefCell::new(())),
53            value: UnsafeCell::new(value),
54        }
55    }
56
57    /// Equivalent to [`RefCell::into_inner`].
58    #[inline]
59    pub fn into_inner(self) -> T {
60        self.value.into_inner()
61    }
62
63    /// Equivalent to [`RefCell::replace`].
64    #[inline]
65    #[track_caller]
66    pub fn replace(&self, t: T) -> T {
67        mem::replace(&mut *self.borrow_mut(), t)
68    }
69
70    /// Equivalent to [`RefCell::replace_with`].
71    #[inline]
72    #[track_caller]
73    pub fn replace_with<F: FnOnce(&mut T) -> T>(&self, f: F) -> T {
74        let mut_borrow = &mut *self.borrow_mut();
75        let replacement = f(mut_borrow);
76        mem::replace(mut_borrow, replacement)
77    }
78
79    /// Equivalent to [`RefCell::swap`].
80    #[inline]
81    pub fn swap(&self, other: &Self) {
82        mem::swap(&mut *self.borrow_mut(), &mut *other.borrow_mut());
83    }
84}
85
86impl<T: ?Sized> BetterRefCell<T> {
87    /// Equivalent to [`RefCell::borrow`].
88    #[inline]
89    #[track_caller]
90    pub fn borrow(&self) -> Ref<'_, T> {
91        let state = unsafe { &*self.state.as_ptr() };
92        Ref::map(state.borrow(), |()| unsafe {
93            self.as_ptr().as_ref().unwrap_unchecked()
94        })
95    }
96
97    /// Equivalent to [`RefCell::try_borrow`].
98    #[inline]
99    pub fn try_borrow(&self) -> Result<Ref<'_, T>, BorrowError> {
100        let state = unsafe { &*self.state.as_ptr() };
101        Ok(Ref::map(state.try_borrow()?, |()| unsafe {
102            self.as_ptr().as_ref().unwrap_unchecked()
103        }))
104    }
105
106    /// Equivalent to [`RefCell::borrow_mut`].
107    #[inline]
108    #[track_caller]
109    pub fn borrow_mut(&self) -> RefMut<'_, T> {
110        let state = unsafe { &*self.state.as_ptr() };
111        RefMut::map(state.borrow_mut(), |()| unsafe {
112            self.as_ptr().as_mut().unwrap_unchecked()
113        })
114    }
115
116    /// Equivalent to [`RefCell::try_borrow_mut`].
117    #[inline]
118    pub fn try_borrow_mut(&self) -> Result<RefMut<'_, T>, BorrowMutError> {
119        let state = unsafe { &*self.state.as_ptr() };
120        Ok(RefMut::map(state.try_borrow_mut()?, |()| unsafe {
121            self.as_ptr().as_mut().unwrap_unchecked()
122        }))
123    }
124
125    /// Equivalent to [`RefCell::as_ptr`].
126    #[inline]
127    pub fn as_ptr(&self) -> *mut T {
128        if let Some(ptr) = self.ptr.get() {
129            ptr.as_ptr()
130        } else {
131            self.value.get()
132        }
133    }
134
135    /// Equivalent to [`RefCell::get_mut`].
136    #[inline]
137    pub fn get_mut(&mut self) -> &mut T {
138        self.value.get_mut()
139    }
140
141    /// Equivalent to [`RefCell::try_borrow_unguarded`].
142    #[inline]
143    pub unsafe fn try_borrow_unguarded(&self) -> Result<&T, BorrowError> {
144        let state = unsafe { &*self.state.as_ptr() };
145        unsafe { state.try_borrow_unguarded()? };
146        Ok(unsafe { self.as_ptr().as_ref().unwrap_unchecked() })
147    }
148}
149
150impl<T: Default> BetterRefCell<T> {
151    /// Equivalent to [`RefCell::take`].
152    pub fn take(&self) -> T {
153        self.replace(Default::default())
154    }
155}
156
157impl<T: Clone> Clone for BetterRefCell<T> {
158    /// Equivalent to [`RefCell::clone`].
159    #[inline]
160    #[track_caller]
161    fn clone(&self) -> Self {
162        Self::new(self.borrow().clone())
163    }
164
165    /// Equivalent to [`RefCell::clone_from`].
166    #[inline]
167    #[track_caller]
168    fn clone_from(&mut self, source: &Self) {
169        self.get_mut().clone_from(&source.borrow());
170    }
171}
172
173impl<T: Default> Default for BetterRefCell<T> {
174    /// Equivalent to [`RefCell::default`].
175    #[inline]
176    fn default() -> Self {
177        Self::new(Default::default())
178    }
179}
180
181impl<T: ?Sized + PartialEq> PartialEq for BetterRefCell<T> {
182    /// Equivalent to [`RefCell::eq`].
183    #[inline]
184    fn eq(&self, other: &Self) -> bool {
185        *self.borrow() == *other.borrow()
186    }
187}
188
189impl<T: ?Sized + Eq> Eq for BetterRefCell<T> {}
190
191impl<T: ?Sized + PartialOrd> PartialOrd for BetterRefCell<T> {
192    /// Equivalent to [`RefCell::partial_cmp`].
193    #[inline]
194    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
195        self.borrow().partial_cmp(&*other.borrow())
196    }
197
198    /// Equivalent to [`RefCell::lt`].
199    #[inline]
200    fn lt(&self, other: &Self) -> bool {
201        *self.borrow() < *other.borrow()
202    }
203
204    /// Equivalent to [`RefCell::le`].
205    #[inline]
206    fn le(&self, other: &Self) -> bool {
207        *self.borrow() <= *other.borrow()
208    }
209
210    /// Equivalent to [`RefCell::gt`].
211    #[inline]
212    fn gt(&self, other: &Self) -> bool {
213        *self.borrow() > *other.borrow()
214    }
215
216    /// Equivalent to [`RefCell::ge`].
217    #[inline]
218    fn ge(&self, other: &Self) -> bool {
219        *self.borrow() >= *other.borrow()
220    }
221}
222
223impl<T: ?Sized + Ord> Ord for BetterRefCell<T> {
224    /// Equivalent to [`RefCell::cmp`].
225    #[inline]
226    fn cmp(&self, other: &Self) -> cmp::Ordering {
227        self.borrow().cmp(&*other.borrow())
228    }
229}
230
231impl<T> From<T> for BetterRefCell<T> {
232    /// Equivalent to [`RefCell::from`].
233    fn from(t: T) -> Self {
234        Self::new(t)
235    }
236}
237
238impl<T: ?Sized> BetterRefCell<T> {
239    /// Temporarily re-grant exclusive access to the internal value, allowing `borrow_mut`/`try_borrow_mut` again in `f`.
240    ///
241    /// # Panics
242    ///
243    /// This function panics only if any of the following conditions are violated:
244    /// 1. `borrowed` must point to the internal value stored in this `BetterRefCell`.
245    /// 2. All borrow guards obtained from this `BetterRefCell` within `f` must be dropped before `f` returns.
246    #[track_caller]
247    pub fn unborrow<R>(&self, borrowed: &mut T, f: impl FnOnce() -> R) -> R {
248        if !ptr::eq(self.value.get(), borrowed) {
249            panic_different_address()
250        }
251        let ptr = self.ptr.replace(Some(borrowed.into()));
252        let state = self.state.take();
253        let result = f();
254        let state = self.state.replace(state);
255        if state.try_borrow_mut().is_ok() {
256            self.ptr.set(ptr);
257            result
258        } else {
259            self.state.set(state);
260            panic_borrow_guard_leaked()
261        }
262    }
263
264    /// Temporarily re-grant exclusive access to the internal value, allowing `borrow_mut`/`try_borrow_mut` again in `f`.
265    ///
266    /// # Safety
267    ///
268    /// This function is safe to call if all of the following conditions are met:
269    /// 1. `borrowed` must point to the internal value stored in this `BetterRefCell`.
270    /// 2. All borrow guards obtained from this `BetterRefCell` within `f` must be dropped before `f` returns.
271    /// 3. Any borrow guards obtained from this `BetterRefCell` outside of `f` must not be used within `f`.
272    #[track_caller]
273    pub unsafe fn unborrow_unchecked<R>(&self, borrowed: &mut T, f: impl FnOnce() -> R) -> R {
274        let ptr = self.ptr.replace(Some(borrowed.into()));
275        let state = self.state.take();
276        let result = f();
277        self.state.set(state);
278        self.ptr.set(ptr);
279        result
280    }
281
282    /// Temporarily re-grant shared access to the internal value, allowing `borrow`/`try_borrow` again in `f`.
283    ///
284    /// # Panics
285    ///
286    /// This function panics only if any of the following conditions are violated:
287    /// 1. `borrowed` must point to the internal value stored in this `BetterRefCell`.
288    /// 2. If `borrowed` is (transitively) originated from `RefMut`, then all `Ref`s obtained from this `BetterRefCell` within `f` must be dropped before `f` returns.
289    #[track_caller]
290    pub fn unborrow_ref<R>(&self, borrowed: &T, f: impl FnOnce() -> R) -> R {
291        if !ptr::eq(self.value.get(), borrowed) {
292            panic_different_address()
293        }
294        if unsafe { &mut *self.state.as_ptr() }.try_borrow().is_ok() {
295            return f();
296        }
297        let ptr = self.ptr.replace(Some(borrowed.into()));
298        let state = self.state.take();
299        let guard = unsafe { (*self.state.as_ptr()).try_borrow().unwrap_unchecked() };
300        let result = f();
301        drop(guard);
302        let state = self.state.replace(state);
303        if state.try_borrow_mut().is_ok() {
304            self.ptr.set(ptr);
305            result
306        } else {
307            self.state.set(state);
308            panic_borrow_guard_leaked()
309        }
310    }
311
312    /// Temporarily re-grant shared access to the internal value, allowing `borrow`/`try_borrow` again in `f`.
313    ///
314    /// # Safety
315    ///
316    /// This function is safe to call if all of the following conditions are met:
317    /// 1. `borrowed` must point to the internal value stored in this `BetterRefCell`.
318    /// 2. If `borrowed` is (transitively) originated from `RefMut`, then all `Ref`s obtained from this `BetterRefCell` within `f` must be dropped before `f` returns.
319    /// 3. Any borrow guards obtained from this `BetterRefCell` outside of `f` must not be used within `f`.
320    #[track_caller]
321    pub unsafe fn unborrow_ref_unchecked<R>(&self, borrowed: &T, f: impl FnOnce() -> R) -> R {
322        if unsafe { &mut *self.state.as_ptr() }.try_borrow().is_ok() {
323            return f();
324        }
325        let ptr = self.ptr.replace(Some(borrowed.into()));
326        let state = self.state.take();
327        let result = f();
328        self.state.set(state);
329        self.ptr.set(ptr);
330        result
331    }
332}
333
334#[inline(never)]
335#[track_caller]
336#[cold]
337fn panic_different_address() -> ! {
338    panic!("reference is pointing to a different address")
339}
340
341#[inline(never)]
342#[track_caller]
343#[cold]
344fn panic_borrow_guard_leaked() -> ! {
345    panic!("borrow guard leaked in closure")
346}
347
348#[cfg(test)]
349mod tests {
350    use super::*;
351
352    #[test]
353    fn unborrow() {
354        let cell = BetterRefCell::new(0);
355        let mut guard = cell.borrow_mut();
356        cell.unborrow(&mut guard, || {
357            let mut guard = cell.borrow_mut();
358            assert_eq!(*guard, 0);
359            *guard += 1;
360        });
361        assert_eq!(*guard, 1);
362    }
363
364    #[test]
365    #[should_panic = "reference is pointing to a different address"]
366    fn unborrow_by_invalid_reference() {
367        let cell = BetterRefCell::new(0);
368        cell.unborrow(&mut 0, || {});
369    }
370
371    #[test]
372    #[should_panic = "borrow guard leaked"]
373    fn unborrow_and_leak_ref_guard() {
374        let cell = BetterRefCell::new(0);
375        cell.unborrow(&mut cell.borrow_mut(), || cell.borrow());
376    }
377
378    #[test]
379    #[should_panic = "borrow guard leaked"]
380    fn unborrow_and_leak_mut_guard() {
381        let cell = BetterRefCell::new(0);
382        cell.unborrow(&mut cell.borrow_mut(), || cell.borrow_mut());
383    }
384
385    #[test]
386    fn unborrow_ref_by_ref() {
387        let cell = BetterRefCell::new(0);
388        let guard = cell.borrow();
389        cell.unborrow_ref(&guard, || {
390            let guard = cell.borrow();
391            assert_eq!(*guard, 0);
392        });
393        assert_eq!(*guard, 0);
394    }
395
396    #[test]
397    fn unborrow_ref_by_mut() {
398        let cell = BetterRefCell::new(0);
399        let guard = cell.borrow_mut();
400        cell.unborrow_ref(&guard, || {
401            let guard = cell.borrow();
402            assert_eq!(*guard, 0);
403        });
404        assert_eq!(*guard, 0);
405    }
406
407    #[test]
408    #[should_panic = "reference is pointing to a different address"]
409    fn unborrow_ref_by_invalid_reference() {
410        let cell = BetterRefCell::new(0);
411        cell.unborrow_ref(&0, || {});
412    }
413
414    #[test]
415    fn unborrow_ref_by_ref_and_leak_guard() {
416        let cell = BetterRefCell::new(0);
417        cell.unborrow_ref(&cell.borrow(), || cell.borrow());
418    }
419
420    #[test]
421    #[should_panic = "borrow guard leaked"]
422    fn unborrow_ref_by_mut_and_leak_guard() {
423        let cell = BetterRefCell::new(0);
424        cell.unborrow_ref(&cell.borrow_mut(), || cell.borrow());
425    }
426
427    #[test]
428    fn unborrow_ref_and_drop_outer_guard() {
429        let cell = BetterRefCell::new(0);
430        let outer_guard = cell.borrow();
431        cell.unborrow_ref(&cell.borrow(), || {
432            drop(outer_guard);
433            assert!(cell.try_borrow_mut().is_err());
434        });
435        assert!(cell.try_borrow_mut().is_ok());
436    }
437}