rusteal_core/
reify_registry.rs

1// Reify registry: manages Rust-side type info, function callbacks, and instance data
2// for runtime-created UE classes.
3//
4// Three registries:
5// 1. Type registry: maps type_id -> RustTypeInfo (constructor, destructor, name)
6// 2. Function registry: maps callback_id -> Rust function closure
7// 3. Instance data: maps UObject pointer -> the allocated Rust data of each
8//    Rust class it is (a Rust class whose parent is a Rust class has both)
9
10use std::collections::HashMap;
11use std::sync::{Arc, Mutex, OnceLock, RwLock};
12
13use crate::{lock_or_recover, read_or_recover, write_or_recover};
14
15// ---------------------------------------------------------------------------
16// Inventory-based auto-registration
17// ---------------------------------------------------------------------------
18
19/// Submitted by `#[uclass]` — holds register + finalize fn pointers.
20pub struct ClassRegistration {
21    /// The class's Rust type ID.
22    pub type_id: u64,
23    /// Create the UClass (type info, parent, class handle). `false` while its
24    /// parent, a Rust class not created yet, is missing; on the last try the
25    /// missing parent is logged and it returns `true`.
26    pub create: fn(last_try: bool) -> bool,
27    /// Add its properties and components, which may name other Rust classes
28    /// (`SubclassOf<MyCharacter>`): every class exists by then.
29    pub register: fn(),
30    pub finalize: fn(),
31    /// Once its class defaults are written: a config class loads its ini
32    /// values over them.
33    pub after_defaults: fn(),
34}
35inventory::collect!(ClassRegistration);
36
37/// Submitted by `#[ustruct]`.
38pub struct StructRegistration {
39    /// Create the UScriptStruct (its handle).
40    pub create: fn(),
41    /// Add its properties, which may name Rust classes and structs: every
42    /// class and struct exists by then.
43    pub register: fn(),
44    /// Link it, after the Rust structs it holds.
45    pub finalize: fn(),
46}
47inventory::collect!(StructRegistration);
48
49/// Submitted by `#[uclass_impl]` — holds register_functions fn pointer, and the
50/// `#[class_defaults]` method's caller when the block has one.
51pub struct ClassFunctionRegistration {
52    /// The Rust type ID of the class the block implements.
53    pub type_id: u64,
54    pub register_functions: fn(),
55    pub class_defaults: Option<fn()>,
56}
57inventory::collect!(ClassFunctionRegistration);
58
59/// What a `#[class_defaults]` method may return: nothing, or a result whose
60/// error is logged.
61pub trait ClassDefaultsOutcome {
62    fn report(self, class: &str);
63}
64
65impl ClassDefaultsOutcome for () {
66    fn report(self, _class: &str) {}
67}
68
69impl ClassDefaultsOutcome for crate::error::RustealResult<()> {
70    fn report(self, class: &str) {
71        if let Err(e) = self {
72            let msg = format!("[Rusteal] {class}: class defaults failed: {e}");
73            let bytes = msg.as_bytes();
74            unsafe {
75                crate::ffi_dispatch::logging_log(1, bytes.as_ptr(), bytes.len() as u32);
76            }
77        }
78    }
79}
80
81/// Registration in phases, parents before children: create every class →
82/// create, fill and link every struct → add every class's properties and
83/// components → register every function →
84/// finalize each class and write its class defaults. A Rust class whose parent
85/// is a Rust class is created after it and finalized after the parent's class
86/// defaults, so its class default object starts from them.
87pub fn register_all_from_inventory() {
88    let mut pending: Vec<&ClassRegistration> = inventory::iter::<ClassRegistration>.into_iter().collect();
89    let class_count = pending.len() as u32;
90    let mut ordered: Vec<&ClassRegistration> = Vec::with_capacity(pending.len());
91    loop {
92        let before = pending.len();
93        pending.retain(|reg| {
94            if (reg.create)(false) {
95                ordered.push(reg);
96                false
97            } else {
98                true
99            }
100        });
101        if pending.is_empty() || pending.len() == before {
102            break;
103        }
104    }
105    // Parents never created: log them.
106    for reg in pending {
107        (reg.create)(true);
108    }
109
110    // Structs, which classes' properties and functions may hold, and whose
111    // properties may name classes.
112    for sreg in inventory::iter::<StructRegistration> {
113        (sreg.create)();
114    }
115    for sreg in inventory::iter::<StructRegistration> {
116        (sreg.register)();
117    }
118    for sreg in inventory::iter::<StructRegistration> {
119        (sreg.finalize)();
120    }
121
122    for reg in &ordered {
123        (reg.register)();
124    }
125    let mut func_reg_count = 0u32;
126    for freg in inventory::iter::<ClassFunctionRegistration> {
127        (freg.register_functions)();
128        func_reg_count += 1;
129    }
130    for reg in &ordered {
131        (reg.finalize)();
132        for freg in inventory::iter::<ClassFunctionRegistration> {
133            if freg.type_id == reg.type_id
134                && let Some(class_defaults) = freg.class_defaults
135            {
136                class_defaults();
137            }
138        }
139        (reg.after_defaults)();
140    }
141
142    // Log registration summary (helps diagnose hot-reload issues).
143    let total_funcs = read_or_recover(func_registry()).len();
144    let msg = format!(
145        "[Rusteal] register_all_from_inventory: {} classes, {} impl blocks, {} function callbacks",
146        class_count, func_reg_count, total_funcs,
147    );
148    let bytes = msg.as_bytes();
149    unsafe {
150        crate::ffi_dispatch::logging_log(0, bytes.as_ptr(), bytes.len() as u32);
151    }
152}
153
154use rusteal_ffi::UObjectHandle;
155
156/// Information about a Rust type registered for reification.
157pub struct RustTypeInfo {
158    /// Human-readable type name (for debugging).
159    pub name: &'static str,
160    /// Allocate and return a default-initialized instance. The returned pointer
161    /// must be freeable by `drop_fn`.
162    pub construct_fn: fn() -> *mut u8,
163    /// Drop and deallocate an instance previously created by `construct_fn`.
164    pub drop_fn: unsafe fn(*mut u8),
165}
166
167use crate::ffi_dispatch::NativePtr;
168
169// Type for reify function callbacks: (obj, rust_data, params)
170// Uses Arc so we can clone the reference out of the registry and release
171// the lock before invoking the callback (prevents deadlock if the callback
172// makes FFI calls that re-enter Rust).
173// `params` is an opaque FFI buffer pointer (`NativePtr` = `*mut u8`).
174type ReifyFunctionCallback = Arc<dyn Fn(UObjectHandle, *mut u8, NativePtr) + Send + Sync>;
175
176/// A registered function and the Rust type whose data it is called with.
177struct FunctionEntry {
178    callback: ReifyFunctionCallback,
179    type_id: u64,
180}
181
182// ---------------------------------------------------------------------------
183// Statics
184// ---------------------------------------------------------------------------
185
186static TYPE_REGISTRY: OnceLock<Mutex<HashMap<u64, RustTypeInfo>>> = OnceLock::new();
187static FUNC_REGISTRY: OnceLock<RwLock<Vec<FunctionEntry>>> = OnceLock::new();
188static INSTANCE_DATA: OnceLock<RwLock<HashMap<u64, Vec<InstanceEntry>>>> = OnceLock::new();
189
190struct InstanceEntry {
191    data: *mut u8,
192    type_id: u64,
193}
194
195// SAFETY: The raw pointer in InstanceEntry is only accessed on the game thread.
196unsafe impl Send for InstanceEntry {}
197unsafe impl Sync for InstanceEntry {}
198
199fn type_registry() -> &'static Mutex<HashMap<u64, RustTypeInfo>> {
200    TYPE_REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
201}
202
203fn func_registry() -> &'static RwLock<Vec<FunctionEntry>> {
204    FUNC_REGISTRY.get_or_init(|| RwLock::new(Vec::new()))
205}
206
207fn instance_data() -> &'static RwLock<HashMap<u64, Vec<InstanceEntry>>> {
208    INSTANCE_DATA.get_or_init(|| RwLock::new(HashMap::new()))
209}
210
211// ---------------------------------------------------------------------------
212// Type registry
213// ---------------------------------------------------------------------------
214
215/// Register a Rust type for reification. Must be called before `create_class`.
216pub fn register_type(type_id: u64, info: RustTypeInfo) {
217    lock_or_recover(type_registry())
218        .insert(type_id, info);
219}
220
221// ---------------------------------------------------------------------------
222// Function registry
223// ---------------------------------------------------------------------------
224
225/// Register a Rust function callback of the Rust type `type_id`, which it is
226/// called with the data of, and return its unique callback ID.
227pub fn register_function<F>(type_id: u64, f: F) -> u64
228where
229    F: Fn(UObjectHandle, *mut u8, NativePtr) + Send + Sync + 'static,
230{
231    let mut vec = write_or_recover(func_registry());
232    let id = vec.len() as u64;
233    vec.push(FunctionEntry { callback: Arc::new(f), type_id });
234    id
235}
236
237// ---------------------------------------------------------------------------
238// Instance lifecycle
239// ---------------------------------------------------------------------------
240
241/// Construct the Rust instance of the type `type_id` for a newly created
242/// UObject, once per Rust class the object is. Called from the C++ class
243/// constructor via `construct_rust_instance` callback.
244pub fn construct_instance(obj: UObjectHandle, type_id: u64) {
245    let types = lock_or_recover(type_registry());
246    let Some(info) = types.get(&type_id) else {
247        // Log warning — type not registered (might be a CDO before registration completes)
248        if crate::api::is_api_initialized() {
249            let msg = format!("[Rusteal] construct_instance: unknown type_id {type_id}");
250            let bytes = msg.as_bytes();
251            unsafe {
252                crate::ffi_dispatch::logging_log(1, bytes.as_ptr(), bytes.len() as u32);
253            }
254        }
255        return;
256    };
257    let data = (info.construct_fn)();
258    drop(types); // Release lock before acquiring instance_data lock
259
260    let key = obj.to_addr();
261    let old = {
262        let mut map = write_or_recover(instance_data());
263        let entries = map.entry(key).or_default();
264        match entries.iter_mut().find(|e| e.type_id == type_id) {
265            Some(entry) => Some(std::mem::replace(&mut entry.data, data)),
266            None => {
267                entries.push(InstanceEntry { data, type_id });
268                None
269            }
270        }
271    };
272    // A new object at the address of one whose deletion went unnoticed.
273    if let Some(old) = old {
274        drop_data(type_id, old);
275    }
276}
277
278fn drop_data(type_id: u64, data: *mut u8) {
279    let types = lock_or_recover(type_registry());
280    if let Some(info) = types.get(&type_id) {
281        unsafe {
282            (info.drop_fn)(data);
283        }
284    }
285}
286
287/// Drop and remove the Rust instances of a destroyed UObject, every Rust
288/// class's it is. Called from the C++ delete listener via
289/// `drop_rust_instance` callback.
290pub fn drop_instance(obj: UObjectHandle, _type_id: u64) {
291    let key = obj.to_addr();
292    let entries = write_or_recover(instance_data()).remove(&key);
293    for entry in entries.into_iter().flatten() {
294        drop_data(entry.type_id, entry.data);
295    }
296}
297
298/// Invoke a registered Rust function callback.
299/// Called from the C++ thunk via `invoke_rust_function` callback.
300pub fn invoke_function(callback_id: u64, obj: UObjectHandle, params: NativePtr) {
301    // Clone the callback Arc out of the registry and release the read lock
302    // BEFORE invoking the callback. This prevents deadlocks if the callback
303    // makes FFI calls that re-enter Rust.
304    let func = {
305        let vec = read_or_recover(func_registry());
306        vec.get(callback_id as usize)
307            .map(|entry| (entry.callback.clone(), entry.type_id))
308    };
309
310    if let Some((func, type_id)) = func {
311        // The data of the Rust class the function belongs to.
312        let rust_data = get_instance_data(obj, type_id);
313        func(obj, rust_data, params);
314    } else if crate::api::is_api_initialized() {
315        let vec_len = read_or_recover(func_registry()).len();
316        let msg = format!(
317            "[Rusteal] invoke_function: callback_id {} not found (registry size = {})",
318            callback_id, vec_len,
319        );
320        let bytes = msg.as_bytes();
321        unsafe {
322            crate::ffi_dispatch::logging_log(1, bytes.as_ptr(), bytes.len() as u32);
323        }
324    }
325}
326
327/// Clear all registries and drop all instance data.
328/// Called during shutdown before DLL unload (enables hot reload).
329pub fn clear_all() {
330    // 1. Drop all instance data, using the type's drop_fn.
331    if let Some(instances) = INSTANCE_DATA.get() {
332        let mut map = write_or_recover(instances);
333        let types = lock_or_recover(type_registry());
334        for entry in map.drain().flat_map(|(_, entries)| entries) {
335            if let Some(info) = types.get(&entry.type_id) {
336                unsafe {
337                    (info.drop_fn)(entry.data);
338                }
339            }
340        }
341        drop(types);
342    }
343    // 2. Clear function registry.
344    if let Some(funcs) = FUNC_REGISTRY.get() {
345        write_or_recover(funcs).clear();
346    }
347    // 3. Clear type registry.
348    if let Some(types) = TYPE_REGISTRY.get() {
349        lock_or_recover(types).clear();
350    }
351}
352
353/// Get the Rust instance data pointer of the type `type_id` for a UObject.
354/// Returns null if no instance data is registered.
355pub fn get_instance_data(obj: UObjectHandle, type_id: u64) -> *mut u8 {
356    let key = obj.to_addr();
357    read_or_recover(instance_data())
358        .get(&key)
359        .and_then(|entries| entries.iter().find(|e| e.type_id == type_id))
360        .map(|e| e.data)
361        .unwrap_or(std::ptr::null_mut())
362}
rusteal_core/reify_registry.rs

rusteal_core/
reify_registry.rs
