rusteal_core/
pinned.rs

1use std::collections::HashMap;
2use std::marker::PhantomData;
3use std::ops::Deref;
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::sync::{Arc, Mutex, OnceLock};
6
7use crate::lock_or_recover;
8
9use rusteal_ffi::UObjectHandle;
10
11use crate::error::{RustealError, RustealResult};
12use crate::ffi_dispatch;
13use crate::object_ref::{Checked, UObjectRef};
14use crate::traits::{HasParent, UeClass, UeHandle, ValidHandle};
15
16fn alive_registry() -> &'static Mutex<HashMap<u64, Arc<AtomicBool>>> {
17    static REGISTRY: OnceLock<Mutex<HashMap<u64, Arc<AtomicBool>>>> = OnceLock::new();
18    REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
19}
20
21pub fn notify_pinned_destroyed(handle: UObjectHandle) {
22    if let Ok(registry) = alive_registry().lock()
23        && let Some(flag) = registry.get(&handle.to_addr())
24    {
25        flag.store(false, Ordering::Relaxed);
26    }
27}
28
29pub fn clear_all() {
30    if let Ok(mut registry) = alive_registry().lock() {
31        registry.clear();
32    }
33}
34
35#[repr(C)]
36pub struct Pinned<T: UeClass> {
37    handle: UObjectHandle,
38    alive: Arc<AtomicBool>,
39    _marker: PhantomData<*const T>,
40}
41
42unsafe impl<T: UeClass> Send for Pinned<T> {}
43
44impl<T: UeClass> Pinned<T> {
45    #[track_caller]
46    pub fn new(obj: UObjectRef<T>) -> RustealResult<Self> {
47        if !obj.is_valid() {
48            return Err(RustealError::object_destroyed());
49        }
50
51        let alive = Arc::new(AtomicBool::new(true));
52
53        lock_or_recover(alive_registry()).insert(obj.raw().to_addr(), alive.clone());
54
55        unsafe {
56            ffi_dispatch::lifecycle_add_gc_root(obj.raw());
57            ffi_dispatch::lifecycle_register_pinned(obj.raw());
58        }
59
60        Ok(Pinned {
61            handle: obj.raw(),
62            alive,
63            _marker: PhantomData,
64        })
65    }
66
67    #[inline]
68    pub fn is_alive(&self) -> bool {
69        self.alive.load(Ordering::Relaxed)
70    }
71
72    #[inline]
73    pub fn handle(&self) -> UObjectHandle {
74        self.handle
75    }
76
77    #[inline]
78    pub fn as_ref(&self) -> UObjectRef<T> {
79        unsafe { UObjectRef::from_raw(self.handle) }
80    }
81
82    #[inline]
83    pub fn as_checked(&self) -> Checked<T> {
84        debug_assert!(self.is_alive(), "Pinned object has been destroyed");
85
86        Checked::new_unchecked(self.handle)
87    }
88}
89
90impl<T: UeClass> Drop for Pinned<T> {
91    fn drop(&mut self) {
92        lock_or_recover(alive_registry()).remove(&self.handle.to_addr());
93
94        unsafe {
95            ffi_dispatch::lifecycle_unregister_pinned(self.handle);
96            ffi_dispatch::lifecycle_remove_gc_root(self.handle);
97        }
98    }
99}
100
101impl<T: UeClass> ValidHandle for Pinned<T> {
102    #[inline]
103    fn handle(&self) -> UObjectHandle {
104        debug_assert!(self.is_alive(), "Pinned object has been destroyed");
105
106        self.handle
107    }
108}
109
110impl<T: UeClass> UeHandle for Pinned<T> {
111    #[inline]
112    fn checked_handle(&self) -> RustealResult<UObjectHandle> {
113        if self.alive.load(Ordering::Relaxed) {
114            Ok(self.handle)
115        } else {
116            Err(RustealError::object_destroyed())
117        }
118    }
119
120    #[inline]
121    fn raw_handle(&self) -> UObjectHandle {
122        self.handle
123    }
124}
125
126impl<T: UeClass> std::fmt::Debug for Pinned<T> {
127    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128        f.debug_struct("Pinned")
129            .field("handle", &self.handle)
130            .field("alive", &self.is_alive())
131            .finish()
132    }
133}
134
135impl<T: HasParent> Deref for Pinned<T> {
136    type Target = Pinned<T::Parent>;
137    #[inline]
138    fn deref(&self) -> &Pinned<T::Parent> {
139        unsafe { &*(self as *const _ as *const Pinned<T::Parent>) }
140    }
141}
rusteal_core/pinned.rs

rusteal_core/
pinned.rs
