1#![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 #[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 #[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 #[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 #[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}