1use std::prelude::v1::*;
14
15#[cfg(feature = "std")]
16use std::error::Error;
17
18use std::{
19 cell::UnsafeCell,
20 fmt::{Display, Error as FormatError, Formatter},
21 ops::{Deref, DerefMut},
22 sync::atomic::{AtomicUsize, Ordering},
23 usize,
24};
25
26macro_rules! borrow_panic {
27 ($s:expr) => {{
28 panic!(
29 "Tried to fetch data of type {:?}, but it was already borrowed{}.",
30 core::any::type_name::<T>(),
31 $s,
32 )
33 }};
34}
35
36#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
38pub struct InvalidBorrow;
39
40impl Display for InvalidBorrow {
41 fn fmt(
42 &self,
43 f: &mut Formatter,
44 ) -> Result<(), FormatError> {
45 write!(f, "Tried to borrow when it was illegal")
46 }
47}
48
49#[cfg(feature = "std")]
50impl Error for InvalidBorrow {
51 fn description(&self) -> &str {
52 "This error is returned when you try to borrow immutably when it's already \
53 borrowed mutably or you try to borrow mutably when it's already borrowed"
54 }
55}
56
57#[derive(Debug)]
61pub struct Ref<'a, T: ?Sized + 'a> {
62 flag: &'a AtomicUsize,
63 pub value: &'a T,
64}
65
66impl<'a, T: ?Sized> Ref<'a, T> {
67 pub fn map<U, F>(
124 self,
125 f: F,
126 ) -> Ref<'a, U>
127 where
128 F: FnOnce(&T) -> &U,
129 U: ?Sized,
130 {
131 let flag = unsafe { &*(self.flag as *const _) };
136 let value = unsafe { &*(self.value as *const _) };
137
138 std::mem::forget(self);
142
143 Ref {
144 flag,
145 value: f(value),
146 }
147 }
148}
149
150impl<'a, T: ?Sized> Deref for Ref<'a, T> {
151 type Target = T;
152
153 fn deref(&self) -> &T {
154 self.value
155 }
156}
157
158impl<'a, T: ?Sized> Drop for Ref<'a, T> {
159 fn drop(&mut self) {
160 self.flag.fetch_sub(1, Ordering::Release);
161 }
162}
163
164impl<'a, T: ?Sized> Clone for Ref<'a, T> {
165 fn clone(&self) -> Self {
166 self.flag.fetch_add(1, Ordering::Release);
167 Ref {
168 flag: self.flag,
169 value: self.value,
170 }
171 }
172}
173
174#[derive(Debug)]
178pub struct RefMut<'a, T: ?Sized + 'a> {
179 flag: &'a AtomicUsize,
180 value: &'a mut T,
181}
182
183impl<'a, T: ?Sized> RefMut<'a, T> {
184 pub fn map<U, F>(
229 self,
230 f: F,
231 ) -> RefMut<'a, U>
232 where
233 F: FnOnce(&mut T) -> &mut U,
234 U: ?Sized,
235 {
236 let flag = unsafe { &*(self.flag as *const _) };
241 let value = unsafe { &mut *(self.value as *mut _) };
242
243 std::mem::forget(self);
247
248 RefMut {
249 flag,
250 value: f(value),
251 }
252 }
253}
254
255impl<'a, T: ?Sized> Deref for RefMut<'a, T> {
256 type Target = T;
257
258 fn deref(&self) -> &T {
259 self.value
260 }
261}
262
263impl<'a, T: ?Sized> DerefMut for RefMut<'a, T> {
264 fn deref_mut(&mut self) -> &mut T {
265 self.value
266 }
267}
268
269impl<'a, T: ?Sized> Drop for RefMut<'a, T> {
270 fn drop(&mut self) {
271 self.flag.store(0, Ordering::Release)
272 }
273}
274
275#[derive(Debug)]
277pub struct TrustCell<T> {
278 flag: AtomicUsize,
279 inner: UnsafeCell<T>,
280}
281
282impl<T> TrustCell<T> {
283 pub fn new(val: T) -> Self {
285 TrustCell {
286 flag: AtomicUsize::new(0),
287 inner: UnsafeCell::new(val),
288 }
289 }
290
291 pub fn into_inner(self) -> T {
293 self.inner.into_inner()
294 }
295
296 pub fn borrow(&self) -> Ref<T> {
305 self.check_flag_read()
306 .unwrap_or_else(|_| borrow_panic!(" mutably"));
307
308 Ref {
309 flag: &self.flag,
310 value: unsafe { &*self.inner.get() },
311 }
312 }
313
314 #[allow(dead_code)]
319 pub fn try_borrow(&self) -> Result<Ref<T>, InvalidBorrow> {
320 self.check_flag_read()?;
321
322 Ok(Ref {
323 flag: &self.flag,
324 value: unsafe { &*self.inner.get() },
325 })
326 }
327
328 pub fn borrow_mut(&self) -> RefMut<T> {
337 self.check_flag_write()
338 .unwrap_or_else(|_| borrow_panic!(""));
339
340 RefMut {
341 flag: &self.flag,
342 value: unsafe { &mut *self.inner.get() },
343 }
344 }
345
346 #[allow(dead_code)]
351 pub fn try_borrow_mut(&self) -> Result<RefMut<T>, InvalidBorrow> {
352 self.check_flag_write()?;
353
354 Ok(RefMut {
355 flag: &self.flag,
356 value: unsafe { &mut *self.inner.get() },
357 })
358 }
359
360 #[allow(dead_code)]
364 pub fn get_mut(&mut self) -> &mut T {
365 unsafe { &mut *self.inner.get() }
367 }
368
369 fn check_flag_read(&self) -> Result<(), InvalidBorrow> {
372 loop {
375 let val = self.flag.load(Ordering::Acquire);
376
377 if val == usize::MAX {
378 return Err(InvalidBorrow);
379 }
380
381 match self
382 .flag
383 .compare_exchange_weak(val, val + 1, Ordering::AcqRel, Ordering::Relaxed)
384 {
385 Ok(_result) => {
386 debug_assert_eq!(val, _result);
387 return Ok(());
388 }
389 _ => {
390 }
392 }
393 }
394 }
395
396 fn check_flag_write(&self) -> Result<(), InvalidBorrow> {
399 match self
402 .flag
403 .compare_exchange(0, usize::MAX, Ordering::AcqRel, Ordering::Relaxed)
404 {
405 Ok(0) => Ok(()),
406 _ => Err(InvalidBorrow),
407 }
408 }
409}
410
411unsafe impl<T> Sync for TrustCell<T> where T: Sync {}
412
413impl<T> Default for TrustCell<T>
414where
415 T: Default,
416{
417 fn default() -> Self {
418 TrustCell::new(Default::default())
419 }
420}
421
422#[cfg(test)]
423mod tests {
424 use super::*;
425
426 #[test]
427 fn allow_multiple_reads() {
428 let cell: TrustCell<_> = TrustCell::new(5);
429
430 let a = cell.borrow();
431 let b = cell.borrow();
432
433 assert_eq!(10, *a + *b);
434 }
435
436 #[test]
437 fn allow_clone_reads() {
438 let cell: TrustCell<_> = TrustCell::new(5);
439
440 let a = cell.borrow();
441 let b = a.clone();
442
443 assert_eq!(10, *a + *b);
444 }
445
446 #[test]
447 fn allow_single_write() {
448 let cell: TrustCell<_> = TrustCell::new(5);
449
450 {
451 let mut a = cell.borrow_mut();
452 *a += 2;
453 *a += 3;
454 }
455
456 assert_eq!(10, *cell.borrow());
457 }
458
459 #[test]
460 #[should_panic(expected = "but it was already borrowed mutably")]
461 fn panic_write_and_read() {
462 let cell: TrustCell<_> = TrustCell::new(5);
463
464 let mut a = cell.borrow_mut();
465 *a = 7;
466
467 assert_eq!(7, *cell.borrow());
468 }
469
470 #[test]
471 #[should_panic(expected = "but it was already borrowed")]
472 fn panic_write_and_write() {
473 let cell: TrustCell<_> = TrustCell::new(5);
474
475 let mut a = cell.borrow_mut();
476 *a = 7;
477
478 assert_eq!(7, *cell.borrow_mut());
479 }
480
481 #[test]
482 #[should_panic(expected = "but it was already borrowed")]
483 fn panic_read_and_write() {
484 let cell: TrustCell<_> = TrustCell::new(5);
485
486 let _a = cell.borrow();
487
488 assert_eq!(7, *cell.borrow_mut());
489 }
490
491 #[test]
492 fn try_write_and_read() {
493 let cell: TrustCell<_> = TrustCell::new(5);
494
495 let mut a = cell.try_borrow_mut().unwrap();
496 *a = 7;
497
498 assert!(cell.try_borrow().is_err());
499 }
500
501 #[test]
502 fn try_write_and_write() {
503 let cell: TrustCell<_> = TrustCell::new(5);
504
505 let mut a = cell.try_borrow_mut().unwrap();
506 *a = 7;
507
508 assert!(cell.try_borrow_mut().is_err());
509 }
510
511 #[test]
512 fn try_read_and_write() {
513 let cell: TrustCell<_> = TrustCell::new(5);
514
515 let _a = cell.try_borrow().unwrap();
516
517 assert!(cell.try_borrow_mut().is_err());
518 }
519
520 #[test]
521 fn cloned_borrow_does_not_allow_write() {
522 let cell: TrustCell<_> = TrustCell::new(5);
523
524 let a = cell.borrow();
525 let _b = a.clone();
526 drop(a);
527
528 assert!(cell.try_borrow_mut().is_err());
529 }
530
531 #[test]
532 fn ref_with_non_sized() {
533 let r: Ref<'_, [i32]> = Ref {
534 flag: &AtomicUsize::new(1),
535 value: &[2, 3, 4, 5][..],
536 };
537
538 assert_eq!(&*r, &[2, 3, 4, 5][..]);
539 }
540
541 #[test]
542 fn ref_with_non_sized_clone() {
543 let r: Ref<'_, [i32]> = Ref {
544 flag: &AtomicUsize::new(1),
545 value: &[2, 3, 4, 5][..],
546 };
547 let rr = r.clone();
548
549 assert_eq!(&*r, &[2, 3, 4, 5][..]);
550 assert_eq!(r.flag.load(Ordering::SeqCst), 2);
551
552 assert_eq!(&*rr, &[2, 3, 4, 5][..]);
553 assert_eq!(rr.flag.load(Ordering::SeqCst), 2);
554 }
555
556 #[test]
557 fn ref_with_trait_obj() {
558 let ra: Ref<'_, dyn std::any::Any> = Ref {
559 flag: &AtomicUsize::new(1),
560 value: &2i32,
561 };
562
563 assert_eq!(ra.downcast_ref::<i32>().unwrap(), &2i32);
564 }
565
566 #[test]
567 fn ref_mut_with_non_sized() {
568 let mut r: RefMut<'_, [i32]> = RefMut {
569 flag: &AtomicUsize::new(1),
570 value: &mut [2, 3, 4, 5][..],
571 };
572
573 assert_eq!(&mut *r, &mut [2, 3, 4, 5][..]);
574 }
575
576 #[test]
577 fn ref_mut_with_trait_obj() {
578 let mut ra: RefMut<'_, dyn std::any::Any> = RefMut {
579 flag: &AtomicUsize::new(1),
580 value: &mut 2i32,
581 };
582
583 assert_eq!(ra.downcast_mut::<i32>().unwrap(), &mut 2i32);
584 }
585
586 #[test]
587 fn ref_map_box() {
588 let cell = TrustCell::new(Box::new(10));
589
590 let r: Ref<'_, Box<usize>> = cell.borrow();
591 assert_eq!(&**r, &10);
592
593 let rr: Ref<'_, usize> = cell.borrow().map(Box::as_ref);
594 assert_eq!(&*rr, &10);
595 }
596
597 #[test]
598 fn ref_map_preserves_flag() {
599 let cell = TrustCell::new(Box::new(10));
600
601 let r: Ref<'_, Box<usize>> = cell.borrow();
602 assert_eq!(cell.flag.load(Ordering::SeqCst), 1);
603 let _nr: Ref<'_, usize> = r.map(Box::as_ref);
604 assert_eq!(cell.flag.load(Ordering::SeqCst), 1);
605 }
606
607 #[test]
608 fn ref_map_retains_borrow() {
609 let cell = TrustCell::new(Box::new(10));
610
611 let _r: Ref<'_, usize> = cell.borrow().map(Box::as_ref);
612 assert_eq!(cell.flag.load(Ordering::SeqCst), 1);
613
614 let _rr: Ref<'_, usize> = cell.borrow().map(Box::as_ref);
615 assert_eq!(cell.flag.load(Ordering::SeqCst), 2);
616 }
617
618 #[test]
619 fn ref_map_drops_borrow() {
620 let cell = TrustCell::new(Box::new(10));
621
622 let r: Ref<'_, usize> = cell.borrow().map(Box::as_ref);
623
624 assert_eq!(cell.flag.load(Ordering::SeqCst), 1);
625 drop(r);
626 assert_eq!(cell.flag.load(Ordering::SeqCst), 0);
627 }
628
629 #[test]
630 fn ref_mut_map_box() {
631 let cell = TrustCell::new(Box::new(10));
632
633 {
634 let mut r: RefMut<'_, Box<usize>> = cell.borrow_mut();
635 assert_eq!(&mut **r, &mut 10);
636 }
637 {
638 let mut rr: RefMut<'_, usize> = cell.borrow_mut().map(Box::as_mut);
639 assert_eq!(&mut *rr, &mut 10);
640 }
641 }
642
643 #[test]
644 fn ref_mut_map_preserves_flag() {
645 let cell = TrustCell::new(Box::new(10));
646
647 let r: RefMut<'_, Box<usize>> = cell.borrow_mut();
648 assert_eq!(cell.flag.load(Ordering::SeqCst), std::usize::MAX);
649 let _nr: RefMut<'_, usize> = r.map(Box::as_mut);
650 assert_eq!(cell.flag.load(Ordering::SeqCst), std::usize::MAX);
651 }
652
653 #[test]
654 #[should_panic(expected = "but it was already borrowed")]
655 fn ref_mut_map_retains_mut_borrow() {
656 let cell = TrustCell::new(Box::new(10));
657
658 let _rr: RefMut<'_, usize> = cell.borrow_mut().map(Box::as_mut);
659
660 let _ = cell.borrow_mut();
661 }
662
663 #[test]
664 fn ref_mut_map_drops_borrow() {
665 let cell = TrustCell::new(Box::new(10));
666
667 let r: RefMut<'_, usize> = cell.borrow_mut().map(Box::as_mut);
668
669 assert_eq!(cell.flag.load(Ordering::SeqCst), std::usize::MAX);
670 drop(r);
671 assert_eq!(cell.flag.load(Ordering::SeqCst), 0);
672 }
673}