1use std::collections::HashMap;
2use std::sync::{Arc, Mutex, OnceLock, RwLock};
3
4use crate::{lock_or_recover, read_or_recover, write_or_recover};
5
6pub struct ClassRegistration {
7 pub type_id: u64,
8 pub shape: u64,
9 pub depends_on: &'static [u64],
10 pub create: fn(last_try: bool, shape: u64) -> bool,
11 pub register: fn(),
12 pub finalize: fn(),
13 pub after_defaults: fn(),
14}
15inventory::collect!(ClassRegistration);
16
17pub struct StructRegistration {
18 pub create: fn(),
19 pub register: fn(),
20 pub finalize: fn(),
21}
22inventory::collect!(StructRegistration);
23
24pub struct EnumRegistration {
25 pub create: fn(),
26}
27inventory::collect!(EnumRegistration);
28
29pub struct InterfaceRegistration {
30 pub create: fn(),
31 pub register: fn(),
32 pub finalize: fn(),
33}
34inventory::collect!(InterfaceRegistration);
35
36pub struct ClassFunctionRegistration {
37 pub type_id: u64,
38 pub shape: u64,
39 pub register_functions: fn(),
40 pub class_defaults: Option<fn()>,
41}
42inventory::collect!(ClassFunctionRegistration);
43
44pub trait ClassDefaultsOutcome {
45 fn report(self, class: &str);
46}
47
48impl ClassDefaultsOutcome for () {
49 fn report(self, _class: &str) {}
50}
51
52impl ClassDefaultsOutcome for crate::error::RustealResult<()> {
53 fn report(self, class: &str) {
54 if let Err(e) = self {
55 e.__log_in(
56 crate::logging::LOG_WARNING,
57 &format!("{class}'s class defaults"),
58 );
59 }
60 }
61}
62
63pub fn register_all_from_inventory() {
64 for ereg in inventory::iter::<EnumRegistration> {
65 (ereg.create)();
66 }
67
68 for ireg in inventory::iter::<InterfaceRegistration> {
69 (ireg.create)();
70 }
71
72 let mut pending: Vec<&ClassRegistration> =
73 inventory::iter::<ClassRegistration>.into_iter().collect();
74
75 let class_count = pending.len() as u32;
76 let mut ordered: Vec<&ClassRegistration> = Vec::with_capacity(pending.len());
77
78 loop {
79 let before = pending.len();
80
81 pending.retain(|reg| {
82 if (reg.create)(false, class_shape(reg)) {
83 ordered.push(reg);
84
85 false
86 } else {
87 true
88 }
89 });
90
91 if pending.is_empty() || pending.len() == before {
92 break;
93 }
94 }
95
96 for reg in pending {
97 (reg.create)(true, class_shape(reg));
98 }
99
100 for sreg in inventory::iter::<StructRegistration> {
101 (sreg.create)();
102 }
103
104 for sreg in inventory::iter::<StructRegistration> {
105 (sreg.register)();
106 }
107
108 for sreg in inventory::iter::<StructRegistration> {
109 (sreg.finalize)();
110 }
111
112 for ireg in inventory::iter::<InterfaceRegistration> {
113 (ireg.register)();
114 (ireg.finalize)();
115 }
116
117 for reg in &ordered {
118 (reg.register)();
119 }
120
121 let ordered = dependency_order(&ordered);
122
123 let mut func_reg_count = 0u32;
124
125 for freg in inventory::iter::<ClassFunctionRegistration> {
126 (freg.register_functions)();
127 func_reg_count += 1;
128 }
129
130 for reg in &ordered {
131 (reg.finalize)();
132
133 for freg in inventory::iter::<ClassFunctionRegistration> {
134 if freg.type_id == reg.type_id
135 && let Some(class_defaults) = freg.class_defaults
136 {
137 class_defaults();
138 }
139 }
140
141 (reg.after_defaults)();
142 }
143
144 let total_funcs = read_or_recover(func_registry()).len();
145
146 let msg = format!(
147 "[Rusteal] register_all_from_inventory: {class_count} classes, {func_reg_count} impl blocks, {total_funcs} function callbacks",
148 );
149
150 let bytes = msg.as_bytes();
151
152 unsafe {
153 crate::ffi_dispatch::logging_log(0, bytes.as_ptr(), bytes.len() as u32);
154 }
155}
156
157fn dependency_order<'a>(classes: &[&'a ClassRegistration]) -> Vec<&'a ClassRegistration> {
158 let ids: Vec<u64> = classes.iter().map(|class| class.type_id).collect();
159 let mut state = vec![0u8; classes.len()];
160 let mut out = Vec::with_capacity(classes.len());
161
162 fn visit<'a>(
163 index: usize,
164 classes: &[&'a ClassRegistration],
165 ids: &[u64],
166 state: &mut [u8],
167 out: &mut Vec<&'a ClassRegistration>,
168 ) {
169 if state[index] != 0 {
170 return;
171 }
172
173 state[index] = 1;
174
175 for dependency in classes[index].depends_on {
176 if let Some(found) = ids.iter().position(|id| id == dependency) {
177 visit(found, classes, ids, state, out);
178 }
179 }
180
181 state[index] = 2;
182 out.push(classes[index]);
183 }
184
185 for index in 0..classes.len() {
186 visit(index, classes, &ids, &mut state, &mut out);
187 }
188
189 out
190}
191
192fn class_shape(class: &ClassRegistration) -> u64 {
193 let mut blocks: Vec<u64> = inventory::iter::<ClassFunctionRegistration>
194 .into_iter()
195 .filter(|block| block.type_id == class.type_id)
196 .map(|block| block.shape)
197 .collect();
198
199 blocks.sort_unstable();
200 combine_shapes(class.shape, &blocks)
201}
202
203fn combine_shapes(class: u64, blocks: &[u64]) -> u64 {
204 let mut hash: u64 = 0xcbf29ce484222325;
205
206 for part in std::iter::once(class).chain(blocks.iter().copied()) {
207 for byte in part.to_le_bytes() {
208 hash ^= byte as u64;
209 hash = hash.wrapping_mul(0x100000001b3);
210 }
211 }
212
213 hash
214}
215
216use rusteal_ffi::UObjectHandle;
217
218pub struct RustTypeInfo {
219 pub name: &'static str,
220 pub construct_fn: fn() -> *mut u8,
221 pub drop_fn: unsafe fn(*mut u8),
222}
223
224use crate::ffi_dispatch::NativePtr;
225
226type ReifyFunctionCallback = Arc<dyn Fn(UObjectHandle, *mut u8, NativePtr) + Send + Sync>;
227
228struct FunctionEntry {
229 callback: ReifyFunctionCallback,
230 type_id: u64,
231}
232
233static TYPE_REGISTRY: OnceLock<Mutex<HashMap<u64, RustTypeInfo>>> = OnceLock::new();
234static FUNC_REGISTRY: OnceLock<RwLock<Vec<FunctionEntry>>> = OnceLock::new();
235static INSTANCE_DATA: OnceLock<RwLock<HashMap<u64, Vec<InstanceEntry>>>> = OnceLock::new();
236
237struct InstanceEntry {
238 data: *mut u8,
239 type_id: u64,
240}
241
242unsafe impl Send for InstanceEntry {}
243unsafe impl Sync for InstanceEntry {}
244
245fn type_registry() -> &'static Mutex<HashMap<u64, RustTypeInfo>> {
246 TYPE_REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
247}
248
249fn func_registry() -> &'static RwLock<Vec<FunctionEntry>> {
250 FUNC_REGISTRY.get_or_init(|| RwLock::new(Vec::new()))
251}
252
253fn instance_data() -> &'static RwLock<HashMap<u64, Vec<InstanceEntry>>> {
254 INSTANCE_DATA.get_or_init(|| RwLock::new(HashMap::new()))
255}
256
257pub fn register_type(type_id: u64, info: RustTypeInfo) {
258 lock_or_recover(type_registry()).insert(type_id, info);
259}
260
261pub fn register_function<F>(type_id: u64, f: F) -> u64
262where
263 F: Fn(UObjectHandle, *mut u8, NativePtr) + Send + Sync + 'static,
264{
265 let mut vec = write_or_recover(func_registry());
266 let id = vec.len() as u64;
267
268 vec.push(FunctionEntry {
269 callback: Arc::new(f),
270 type_id,
271 });
272
273 id
274}
275
276pub fn construct_instance(obj: UObjectHandle, type_id: u64) {
277 let types = lock_or_recover(type_registry());
278
279 let Some(info) = types.get(&type_id) else {
280 if crate::api::is_api_initialized() {
281 let msg = format!("[Rusteal] construct_instance: unknown type_id {type_id}");
282 let bytes = msg.as_bytes();
283
284 unsafe {
285 crate::ffi_dispatch::logging_log(1, bytes.as_ptr(), bytes.len() as u32);
286 }
287 }
288
289 return;
290 };
291
292 let data = (info.construct_fn)();
293 drop(types);
294
295 let key = obj.to_addr();
296
297 let old = {
298 let mut map = write_or_recover(instance_data());
299 let entries = map.entry(key).or_default();
300 match entries.iter_mut().find(|e| e.type_id == type_id) {
301 Some(entry) => Some(std::mem::replace(&mut entry.data, data)),
302 None => {
303 entries.push(InstanceEntry { data, type_id });
304
305 None
306 }
307 }
308 };
309
310 if let Some(old) = old {
311 drop_data(type_id, old);
312 }
313}
314
315fn drop_data(type_id: u64, data: *mut u8) {
316 let types = lock_or_recover(type_registry());
317
318 if let Some(info) = types.get(&type_id) {
319 unsafe {
320 (info.drop_fn)(data);
321 }
322 }
323}
324
325pub fn drop_instance(obj: UObjectHandle, _type_id: u64) {
326 let key = obj.to_addr();
327 let entries = write_or_recover(instance_data()).remove(&key);
328
329 for entry in entries.into_iter().flatten() {
330 drop_data(entry.type_id, entry.data);
331 }
332}
333
334pub fn invoke_function(callback_id: u64, obj: UObjectHandle, params: NativePtr) {
335 let func = {
336 let vec = read_or_recover(func_registry());
337 vec.get(callback_id as usize)
338 .map(|entry| (entry.callback.clone(), entry.type_id))
339 };
340
341 if let Some((func, type_id)) = func {
342 let rust_data = get_instance_data(obj, type_id);
343 func(obj, rust_data, params);
344 } else if crate::api::is_api_initialized() {
345 let vec_len = read_or_recover(func_registry()).len();
346
347 let msg = format!(
348 "[Rusteal] invoke_function: callback_id {callback_id} not found (registry size = {vec_len})",
349 );
350
351 let bytes = msg.as_bytes();
352
353 unsafe {
354 crate::ffi_dispatch::logging_log(1, bytes.as_ptr(), bytes.len() as u32);
355 }
356 }
357}
358
359pub fn clear_all() {
360 if let Some(instances) = INSTANCE_DATA.get() {
361 let mut map = write_or_recover(instances);
362 let types = lock_or_recover(type_registry());
363
364 for entry in map.drain().flat_map(|(_, entries)| entries) {
365 if let Some(info) = types.get(&entry.type_id) {
366 unsafe {
367 (info.drop_fn)(entry.data);
368 }
369 }
370 }
371
372 drop(types);
373 }
374
375 if let Some(funcs) = FUNC_REGISTRY.get() {
376 write_or_recover(funcs).clear();
377 }
378
379 if let Some(types) = TYPE_REGISTRY.get() {
380 lock_or_recover(types).clear();
381 }
382}
383
384pub fn get_instance_data(obj: UObjectHandle, type_id: u64) -> *mut u8 {
385 let key = obj.to_addr();
386
387 read_or_recover(instance_data())
388 .get(&key)
389 .and_then(|entries| entries.iter().find(|e| e.type_id == type_id))
390 .map(|e| e.data)
391 .unwrap_or(std::ptr::null_mut())
392}
393
394#[cfg(test)]
395mod tests {
396 use super::combine_shapes;
397
398 #[test]
399 fn every_part_of_a_class_makes_its_shape() {
400 let shape = combine_shapes(1, &[2, 3]);
401 assert_eq!(shape, combine_shapes(1, &[2, 3]));
402 assert_ne!(shape, combine_shapes(9, &[2, 3]));
403 assert_ne!(shape, combine_shapes(1, &[2, 4]));
404 assert_ne!(shape, combine_shapes(1, &[2]));
405 }
406}
407
408#[cfg(test)]
409mod order_tests {
410 use super::*;
411
412 fn noop() {}
413
414 fn create(_last_try: bool, _shape: u64) -> bool {
415 true
416 }
417
418 const fn class(type_id: u64, depends_on: &'static [u64]) -> ClassRegistration {
419 ClassRegistration {
420 type_id,
421 shape: 0,
422 depends_on,
423 create,
424 register: noop,
425 finalize: noop,
426 after_defaults: noop,
427 }
428 }
429
430 #[test]
431 fn component_and_parent_classes_come_first() {
432 static GRABBABLE: ClassRegistration = class(1, &[100, 2]);
433 static GRAB: ClassRegistration = class(2, &[101]);
434 static PISTOL: ClassRegistration = class(3, &[100, 2, 4]);
435 static PROJECTILE: ClassRegistration = class(4, &[100]);
436
437 let order: Vec<u64> = dependency_order(&[&GRABBABLE, &PISTOL, &GRAB, &PROJECTILE])
438 .iter()
439 .map(|class| class.type_id)
440 .collect();
441
442 let position = |id| order.iter().position(|o| *o == id).unwrap();
443 assert!(position(2) < position(1));
444 assert!(position(2) < position(3));
445 assert!(position(4) < position(3));
446 assert_eq!(order.len(), 4);
447 }
448
449 #[test]
450 fn cycles_still_order_every_class() {
451 static A: ClassRegistration = class(1, &[2]);
452 static B: ClassRegistration = class(2, &[1]);
453
454 assert_eq!(dependency_order(&[&A, &B]).len(), 2);
455 }
456}