better_refcell/
lib.rs

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