1
use std::cell::UnsafeCell;
2
use std::mem::ManuallyDrop;
3
use std::ops::Deref;
4
use std::ops::DerefMut;
5

            
6
use crate::storage::GlobalTermPoolGuard;
7
use crate::storage::THREAD_TERM_POOL;
8

            
9
/// A mutex that prevents garbage collection by holding a shared read lock on
10
/// the [super::GlobalTermPool] for the duration of the guard's lifetime.
11
/// Returns a [GcMutexGuard] on access.
12
///
13
/// # Safety
14
///
15
/// The `GcMutex` returns guards that are tied to the thread-local storage of
16
/// [crate::storage::THREAD_TERM_POOL]. This means that the guard must be
17
/// dropped before this thread-local storage is dropped. Otherwise
18
/// use-after-free will occur, which is undefined behaviour.
19
pub struct GcMutex<T> {
20
    inner: UnsafeCell<T>,
21
}
22

            
23
// SAFETY: Sharing GcMutex<T> across threads gives shared &T access, so T must be Send.
24
unsafe impl<T: Send> Send for GcMutex<T> {}
25
// SAFETY: Sharing &GcMutex<T> across threads gives shared &T access, so T must be Sync.
26
unsafe impl<T: Send + Sync> Sync for GcMutex<T> {}
27

            
28
impl<T> GcMutex<T> {
29
1166058
    pub fn new(value: T) -> GcMutex<T> {
30
1166058
        GcMutex {
31
1166058
            inner: UnsafeCell::new(value),
32
1166058
        }
33
1166058
    }
34

            
35
    /// Provides shared access to the underlying value, returning a [GcMutexReadGuard].
36
    ///
37
    /// The returned guard holds a read lock on the global term pool, preventing
38
    /// garbage collection for its lifetime. It only provides immutable access.
39
80132321
    pub fn lock(&self) -> GcMutexReadGuard<'_, T> {
40
        GcMutexReadGuard {
41
80132321
            mutex: self,
42
80132321
            guard: ManuallyDrop::new(THREAD_TERM_POOL.with(|tp| unsafe {
43
80132321
                std::mem::transmute::<_, GlobalTermPoolGuard<'_>>(
44
80132321
                    tp.term_pool().read_recursive().expect("Lock poisoned!"),
45
                )
46
80132321
            })),
47
        }
48
80132321
    }
49

            
50
    /// Provides exclusive mutable access to the underlying value, returning a [GcMutexGuard].
51
    ///
52
    /// Takes `&mut self` so only one mutable guard can exist at a time; the borrow
53
    /// checker enforces that no other guard (read or write) coexists.
54
115647291
    pub fn lock_mut(&mut self) -> GcMutexGuard<'_, T> {
55
        GcMutexGuard {
56
115647291
            mutex: self,
57
115647291
            guard: ManuallyDrop::new(THREAD_TERM_POOL.with(|tp| unsafe {
58
115647291
                std::mem::transmute::<_, GlobalTermPoolGuard<'_>>(
59
115647291
                    tp.term_pool().read_recursive().expect("Lock poisoned!"),
60
                )
61
115647291
            })),
62
        }
63
115647291
    }
64
}
65

            
66
/// A read-only guard produced by [GcMutex::lock].  Holds a shared read lock on
67
/// the global term pool for its lifetime, preventing garbage collection.
68
pub struct GcMutexReadGuard<'a, T> {
69
    mutex: &'a GcMutex<T>,
70

            
71
    /// Only used to avoid garbage collection, will be released on drop.
72
    guard: ManuallyDrop<GlobalTermPoolGuard<'a>>,
73
}
74

            
75
impl<T> Deref for GcMutexReadGuard<'_, T> {
76
    type Target = T;
77

            
78
720256455
    fn deref(&self) -> &Self::Target {
79
720256455
        unsafe { &*self.mutex.inner.get() }
80
720256455
    }
81
}
82

            
83
impl<T> Drop for GcMutexReadGuard<'_, T> {
84
80132321
    fn drop(&mut self) {
85
80132321
        if self.guard.read_depth() == 1 {
86
33468118
            THREAD_TERM_POOL.with(|tp| unsafe { tp.trigger_delayed_garbage_collection(&mut self.guard) })
87
46664203
        } else {
88
46664203
            unsafe { ManuallyDrop::drop(&mut self.guard) };
89
46664203
        }
90
80132321
    }
91
}
92

            
93
/// A read-write guard produced by [GcMutex::lock_mut].  Provides both
94
/// [Deref] and [DerefMut].  Because [GcMutex::lock_mut] takes `&mut self`,
95
/// the borrow checker guarantees this is the only live guard for its lifetime.
96
pub struct GcMutexGuard<'a, T> {
97
    mutex: &'a GcMutex<T>,
98

            
99
    /// Only used to avoid garbage collection, will be released on drop.
100
    guard: ManuallyDrop<GlobalTermPoolGuard<'a>>,
101
}
102

            
103
impl<T> Deref for GcMutexGuard<'_, T> {
104
    type Target = T;
105

            
106
149367080
    fn deref(&self) -> &Self::Target {
107
149367080
        unsafe { &*self.mutex.inner.get() }
108
149367080
    }
109
}
110

            
111
impl<T> DerefMut for GcMutexGuard<'_, T> {
112
181611810
    fn deref_mut(&mut self) -> &mut Self::Target {
113
181611810
        unsafe { &mut *self.mutex.inner.get() }
114
181611810
    }
115
}
116

            
117
impl<T> Drop for GcMutexGuard<'_, T> {
118
115647291
    fn drop(&mut self) {
119
115647291
        if self.guard.read_depth() == 1 {
120
            // If this is the last guard, we can trigger garbage collection when it was delayed earlier.
121
46142169
            THREAD_TERM_POOL.with(|tp| unsafe { tp.trigger_delayed_garbage_collection(&mut self.guard) })
122
69505122
        } else {
123
69505122
            // Just drop the guard
124
69505122
            unsafe { ManuallyDrop::drop(&mut self.guard) };
125
69505122
        }
126
115647291
    }
127
}