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