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