rafx_base/
trust_cell.rs

1// This is originally from the shred project. (I wanted to avoid pulling in all dependencies)
2//
3// Shouldn't have any funtional changes (i.e. added some allow dead_code, ran cargo fmt)
4//
5// Original source available on github: https://github.com/slide-rs/shred
6//
7// shred is distributed under the terms of both the MIT license and the Apache License (Version 2.0).
8//
9// See LICENSE-APACHE and LICENSE-MIT.
10
11//! Helper module for some internals, most users don't need to interact with it.
12
13use 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/// Marker struct for an invalid borrow error
37#[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/// An immutable reference to data in a `TrustCell`.
58///
59/// Access the value via `std::ops::Deref` (e.g. `*val`)
60#[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    // ADDED: I had some code like this:
68    //
69    //   let w : &'a systems::TrustCell<resource::ResourceMap> = &*self.resource_map;
70    //   let w_borrow : systems::trust_cell::Ref<'a, resource::ResourceMap> = w.borrow();
71    //   let w_ref : &'a resource::ResourceMap = &*w_borrow;
72    //
73    // It did not compile because it complained w_borrow did not live long enough.
74    // Adding this getter allowed me to use the ref as-is.
75    //
76    //TODO: Make a minimum example and ask why the above code doesn't work
77    // pub fn value(&self) -> &'a T {
78    //     self.value
79    // }
80
81    /// Makes a new `Ref` for a component of the borrowed data which preserves
82    /// the existing borrow.
83    ///
84    /// The `TrustCell` is already immutably borrowed, so this cannot fail.
85    ///
86    /// This is an associated function that needs to be used as `Ref::map(...)`.
87    /// A method would interfere with methods of the same name on the contents
88    /// of a `Ref` used through `Deref`. Further this preserves the borrow of
89    /// the value and hence does the proper cleanup when it's dropped.
90    ///
91    /// # Examples
92    ///
93    /// This can be used to avoid pointer indirection when a boxed item is
94    /// stored in the `TrustCell`.
95    ///
96    /// ```
97    /// use rafx_base::trust_cell::{Ref, TrustCell};
98    ///
99    /// let cb = TrustCell::new(Box::new(5));
100    ///
101    /// // Borrowing the cell causes the `Ref` to store a reference to the `Box`, which is a
102    /// // pointer to the value on the heap, not the actual value.
103    /// let boxed_ref: Ref<'_, Box<usize>> = cb.borrow();
104    /// assert_eq!(**boxed_ref, 5); // Notice the double deref to get the actual value.
105    ///
106    /// // By using `map` we can let `Ref` store a reference directly to the value on the heap.
107    /// let pure_ref: Ref<'_, usize> = Ref::map(boxed_ref, Box::as_ref);
108    ///
109    /// assert_eq!(*pure_ref, 5);
110    /// ```
111    ///
112    /// We can also use `map` to get a reference to a sub-part of the borrowed
113    /// value.
114    ///
115    /// ```rust
116    /// # use rafx_base::trust_cell::{TrustCell, Ref};
117    ///
118    /// let c = TrustCell::new((5, 'b'));
119    /// let b1: Ref<'_, (u32, char)> = c.borrow();
120    /// let b2: Ref<'_, u32> = Ref::map(b1, |t| &t.0);
121    /// assert_eq!(*b2, 5);
122    /// ```
123    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        // Extract the values from the `Ref` through a pointer so that we do not run
132        // `Drop`. Because the returned `Ref` has the same lifetime `'a` as the
133        // given `Ref`, the lifetime we created through turning the pointer into
134        // a ref is valid.
135        let flag = unsafe { &*(self.flag as *const _) };
136        let value = unsafe { &*(self.value as *const _) };
137
138        // We have to forget self so that we do not run `Drop`. Further it's safe
139        // because we are creating a new `Ref`, with the same flag, which will
140        // run the cleanup when it's dropped.
141        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/// A mutable reference to data in a `TrustCell`.
175///
176/// Access the value via `std::ops::DerefMut` (e.g. `*val`)
177#[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    /// Makes a new `RefMut` for a component of the borrowed data which
185    /// preserves the existing borrow.
186    ///
187    /// The `TrustCell` is already mutably borrowed, so this cannot fail.
188    ///
189    /// This is an associated function that needs to be used as
190    /// `RefMut::map(...)`. A method would interfere with methods of the
191    /// same name on the contents of a `RefMut` used through `DerefMut`.
192    /// Further this preserves the borrow of the value and hence does the
193    /// proper cleanup when it's dropped.
194    ///
195    /// # Examples
196    ///
197    /// This can also be used to avoid pointer indirection when a boxed item is
198    /// stored in the `TrustCell`.
199    ///
200    /// ```
201    /// use rafx_base::trust_cell::{RefMut, TrustCell};
202    ///
203    /// let cb = TrustCell::new(Box::new(5));
204    ///
205    /// // Borrowing the cell causes the `RefMut` to store a reference to the `Box`, which is a
206    /// // pointer to the value on the heap, and not a reference directly to the value.
207    /// let boxed_ref: RefMut<'_, Box<usize>> = cb.borrow_mut();
208    /// assert_eq!(**boxed_ref, 5); // Notice the double deref to get the actual value.
209    ///
210    /// // By using `map` we can let `RefMut` store a reference directly to the value on the heap.
211    /// let pure_ref: RefMut<'_, usize> = RefMut::map(boxed_ref, Box::as_mut);
212    ///
213    /// assert_eq!(*pure_ref, 5);
214    /// ```
215    ///
216    /// We can also use `map` to get a reference to a sub-part of the borrowed
217    /// value.
218    ///
219    /// ```rust
220    /// # use rafx_base::trust_cell::{TrustCell, RefMut};
221    ///
222    /// let c = TrustCell::new((5, 'b'));
223    ///
224    /// let b1: RefMut<'_, (u32, char)> = c.borrow_mut();
225    /// let b2: RefMut<'_, u32> = RefMut::map(b1, |t| &mut t.0);
226    /// assert_eq!(*b2, 5);
227    /// ```
228    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        // Extract the values from the `RefMut` through a pointer so that we do not run
237        // `Drop`. Because the returned `RefMut` has the same lifetime `'a` as
238        // the given `RefMut`, the lifetime we created through turning the
239        // pointer into a ref is valid.
240        let flag = unsafe { &*(self.flag as *const _) };
241        let value = unsafe { &mut *(self.value as *mut _) };
242
243        // We have to forget self so that we do not run `Drop`. Further it's safe
244        // because we are creating a new `RefMut`, with the same flag, which
245        // will run the cleanup when it's dropped.
246        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/// A custom cell container that is a `RefCell` with thread-safety.
276#[derive(Debug)]
277pub struct TrustCell<T> {
278    flag: AtomicUsize,
279    inner: UnsafeCell<T>,
280}
281
282impl<T> TrustCell<T> {
283    /// Create a new cell, similar to `RefCell::new`
284    pub fn new(val: T) -> Self {
285        TrustCell {
286            flag: AtomicUsize::new(0),
287            inner: UnsafeCell::new(val),
288        }
289    }
290
291    /// Consumes this cell and returns ownership of `T`.
292    pub fn into_inner(self) -> T {
293        self.inner.into_inner()
294    }
295
296    /// Get an immutable reference to the inner data.
297    ///
298    /// Absence of write accesses is checked at run-time.
299    ///
300    /// # Panics
301    ///
302    /// This function will panic if there is a mutable reference to the data
303    /// already in use.
304    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    /// Get an immutable reference to the inner data.
315    ///
316    /// Absence of write accesses is checked at run-time. If access is not
317    /// possible, an error is returned.
318    #[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    /// Get a mutable reference to the inner data.
329    ///
330    /// Exclusive access is checked at run-time.
331    ///
332    /// # Panics
333    ///
334    /// This function will panic if there are any references to the data already
335    /// in use.
336    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    /// Get a mutable reference to the inner data.
347    ///
348    /// Exclusive access is checked at run-time. If access is not possible, an
349    /// error is returned.#
350    #[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    /// Gets exclusive access to the inner value, bypassing the Cell.
361    ///
362    /// Exclusive access is checked at compile time.
363    #[allow(dead_code)]
364    pub fn get_mut(&mut self) -> &mut T {
365        // safe because we have exclusive access via &mut self
366        unsafe { &mut *self.inner.get() }
367    }
368
369    /// Make sure we are allowed to aquire a read lock, and increment the read
370    /// count by 1
371    fn check_flag_read(&self) -> Result<(), InvalidBorrow> {
372        // Check that no write reference is out, then try to increment the read count
373        // and return once successful.
374        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                    // try again..
391                }
392            }
393        }
394    }
395
396    /// Make sure we are allowed to aquire a write lock, and then set the write
397    /// lock flag.
398    fn check_flag_write(&self) -> Result<(), InvalidBorrow> {
399        // Check we have 0 references out, and then set the ref count to usize::MAX to
400        // indicate a write lock.
401        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}