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 create: fn(last_try: bool, shape: u64) -> bool,
10 pub register: fn(),
11 pub finalize: fn(),
12 pub after_defaults: fn(),
13}
14inventory::collect!(ClassRegistration);
15
16pub struct StructRegistration {
17 pub create: fn(),
18 pub register: fn(),
19 pub finalize: fn(),
20}
21inventory::collect!(StructRegistration);
22
23pub struct ClassFunctionRegistration {
24 pub type_id: u64,
25 pub shape: u64,
26 pub register_functions: fn(),
27 pub class_defaults: Option<fn()>,
28}
29inventory::collect!(ClassFunctionRegistration);
30
31pub trait ClassDefaultsOutcome {
32 fn report(self, class: &str);
33}
34
35impl ClassDefaultsOutcome for () {
36 fn report(self, _class: &str) {}
37}
38
39impl ClassDefaultsOutcome for crate::error::RustealResult<()> {
40 fn report(self, class: &str) {
41 if let Err(e) = self {
42 let msg = format!("[Rusteal] {class}: class defaults failed: {e}");
43 let bytes = msg.as_bytes();
44
45 unsafe {
46 crate::ffi_dispatch::logging_log(1, bytes.as_ptr(), bytes.len() as u32);
47 }
48 }
49 }
50}
51
52pub fn register_all_from_inventory() {
53 let mut pending: Vec<&ClassRegistration> =
54 inventory::iter::<ClassRegistration>.into_iter().collect();
55
56 let class_count = pending.len() as u32;
57 let mut ordered: Vec<&ClassRegistration> = Vec::with_capacity(pending.len());
58
59 loop {
60 let before = pending.len();
61
62 pending.retain(|reg| {
63 if (reg.create)(false, class_shape(reg)) {
64 ordered.push(reg);
65
66 false
67 } else {
68 true
69 }
70 });
71
72 if pending.is_empty() || pending.len() == before {
73 break;
74 }
75 }
76
77 for reg in pending {
78 (reg.create)(true, class_shape(reg));
79 }
80
81 for sreg in inventory::iter::<StructRegistration> {
82 (sreg.create)();
83 }
84
85 for sreg in inventory::iter::<StructRegistration> {
86 (sreg.register)();
87 }
88
89 for sreg in inventory::iter::<StructRegistration> {
90 (sreg.finalize)();
91 }
92
93 for reg in &ordered {
94 (reg.register)();
95 }
96
97 let mut func_reg_count = 0u32;
98
99 for freg in inventory::iter::<ClassFunctionRegistration> {
100 (freg.register_functions)();
101 func_reg_count += 1;
102 }
103
104 for reg in &ordered {
105 (reg.finalize)();
106
107 for freg in inventory::iter::<ClassFunctionRegistration> {
108 if freg.type_id == reg.type_id
109 && let Some(class_defaults) = freg.class_defaults
110 {
111 class_defaults();
112 }
113 }
114
115 (reg.after_defaults)();
116 }
117
118 let total_funcs = read_or_recover(func_registry()).len();
119
120 let msg = format!(
121 "[Rusteal] register_all_from_inventory: {class_count} classes, {func_reg_count} impl blocks, {total_funcs} function callbacks",
122 );
123
124 let bytes = msg.as_bytes();
125
126 unsafe {
127 crate::ffi_dispatch::logging_log(0, bytes.as_ptr(), bytes.len() as u32);
128 }
129}
130
131fn class_shape(class: &ClassRegistration) -> u64 {
132 let mut blocks: Vec<u64> = inventory::iter::<ClassFunctionRegistration>
133 .into_iter()
134 .filter(|block| block.type_id == class.type_id)
135 .map(|block| block.shape)
136 .collect();
137
138 blocks.sort_unstable();
139 combine_shapes(class.shape, &blocks)
140}
141
142fn combine_shapes(class: u64, blocks: &[u64]) -> u64 {
143 let mut hash: u64 = 0xcbf29ce484222325;
144
145 for part in std::iter::once(class).chain(blocks.iter().copied()) {
146 for byte in part.to_le_bytes() {
147 hash ^= byte as u64;
148 hash = hash.wrapping_mul(0x100000001b3);
149 }
150 }
151
152 hash
153}
154
155use rusteal_ffi::UObjectHandle;
156
157pub struct RustTypeInfo {
158 pub name: &'static str,
159 pub construct_fn: fn() -> *mut u8,
160 pub drop_fn: unsafe fn(*mut u8),
161}
162
163use crate::ffi_dispatch::NativePtr;
164
165type ReifyFunctionCallback = Arc<dyn Fn(UObjectHandle, *mut u8, NativePtr) + Send + Sync>;
166
167struct FunctionEntry {
168 callback: ReifyFunctionCallback,
169 type_id: u64,
170}
171
172static TYPE_REGISTRY: OnceLock<Mutex<HashMap<u64, RustTypeInfo>>> = OnceLock::new();
173static FUNC_REGISTRY: OnceLock<RwLock<Vec<FunctionEntry>>> = OnceLock::new();
174static INSTANCE_DATA: OnceLock<RwLock<HashMap<u64, Vec<InstanceEntry>>>> = OnceLock::new();
175
176struct InstanceEntry {
177 data: *mut u8,
178 type_id: u64,
179}
180
181unsafe impl Send for InstanceEntry {}
182unsafe impl Sync for InstanceEntry {}
183
184fn type_registry() -> &'static Mutex<HashMap<u64, RustTypeInfo>> {
185 TYPE_REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
186}
187
188fn func_registry() -> &'static RwLock<Vec<FunctionEntry>> {
189 FUNC_REGISTRY.get_or_init(|| RwLock::new(Vec::new()))
190}
191
192fn instance_data() -> &'static RwLock<HashMap<u64, Vec<InstanceEntry>>> {
193 INSTANCE_DATA.get_or_init(|| RwLock::new(HashMap::new()))
194}
195
196pub fn register_type(type_id: u64, info: RustTypeInfo) {
197 lock_or_recover(type_registry()).insert(type_id, info);
198}
199
200pub fn register_function<F>(type_id: u64, f: F) -> u64
201where
202 F: Fn(UObjectHandle, *mut u8, NativePtr) + Send + Sync + 'static,
203{
204 let mut vec = write_or_recover(func_registry());
205 let id = vec.len() as u64;
206
207 vec.push(FunctionEntry {
208 callback: Arc::new(f),
209 type_id,
210 });
211
212 id
213}
214
215pub fn construct_instance(obj: UObjectHandle, type_id: u64) {
216 let types = lock_or_recover(type_registry());
217
218 let Some(info) = types.get(&type_id) else {
219 if crate::api::is_api_initialized() {
220 let msg = format!("[Rusteal] construct_instance: unknown type_id {type_id}");
221 let bytes = msg.as_bytes();
222
223 unsafe {
224 crate::ffi_dispatch::logging_log(1, bytes.as_ptr(), bytes.len() as u32);
225 }
226 }
227
228 return;
229 };
230
231 let data = (info.construct_fn)();
232 drop(types);
233
234 let key = obj.to_addr();
235
236 let old = {
237 let mut map = write_or_recover(instance_data());
238 let entries = map.entry(key).or_default();
239 match entries.iter_mut().find(|e| e.type_id == type_id) {
240 Some(entry) => Some(std::mem::replace(&mut entry.data, data)),
241 None => {
242 entries.push(InstanceEntry { data, type_id });
243
244 None
245 }
246 }
247 };
248
249 if let Some(old) = old {
250 drop_data(type_id, old);
251 }
252}
253
254fn drop_data(type_id: u64, data: *mut u8) {
255 let types = lock_or_recover(type_registry());
256
257 if let Some(info) = types.get(&type_id) {
258 unsafe {
259 (info.drop_fn)(data);
260 }
261 }
262}
263
264pub fn drop_instance(obj: UObjectHandle, _type_id: u64) {
265 let key = obj.to_addr();
266 let entries = write_or_recover(instance_data()).remove(&key);
267
268 for entry in entries.into_iter().flatten() {
269 drop_data(entry.type_id, entry.data);
270 }
271}
272
273pub fn invoke_function(callback_id: u64, obj: UObjectHandle, params: NativePtr) {
274 let func = {
275 let vec = read_or_recover(func_registry());
276 vec.get(callback_id as usize)
277 .map(|entry| (entry.callback.clone(), entry.type_id))
278 };
279
280 if let Some((func, type_id)) = func {
281 let rust_data = get_instance_data(obj, type_id);
282 func(obj, rust_data, params);
283 } else if crate::api::is_api_initialized() {
284 let vec_len = read_or_recover(func_registry()).len();
285
286 let msg = format!(
287 "[Rusteal] invoke_function: callback_id {callback_id} not found (registry size = {vec_len})",
288 );
289
290 let bytes = msg.as_bytes();
291
292 unsafe {
293 crate::ffi_dispatch::logging_log(1, bytes.as_ptr(), bytes.len() as u32);
294 }
295 }
296}
297
298pub fn clear_all() {
299 if let Some(instances) = INSTANCE_DATA.get() {
300 let mut map = write_or_recover(instances);
301 let types = lock_or_recover(type_registry());
302
303 for entry in map.drain().flat_map(|(_, entries)| entries) {
304 if let Some(info) = types.get(&entry.type_id) {
305 unsafe {
306 (info.drop_fn)(entry.data);
307 }
308 }
309 }
310
311 drop(types);
312 }
313
314 if let Some(funcs) = FUNC_REGISTRY.get() {
315 write_or_recover(funcs).clear();
316 }
317
318 if let Some(types) = TYPE_REGISTRY.get() {
319 lock_or_recover(types).clear();
320 }
321}
322
323pub fn get_instance_data(obj: UObjectHandle, type_id: u64) -> *mut u8 {
324 let key = obj.to_addr();
325
326 read_or_recover(instance_data())
327 .get(&key)
328 .and_then(|entries| entries.iter().find(|e| e.type_id == type_id))
329 .map(|e| e.data)
330 .unwrap_or(std::ptr::null_mut())
331}
332
333#[cfg(test)]
334mod tests {
335 use super::combine_shapes;
336
337 #[test]
338 fn every_part_of_a_class_makes_its_shape() {
339 let shape = combine_shapes(1, &[2, 3]);
340 assert_eq!(shape, combine_shapes(1, &[2, 3]));
341 assert_ne!(shape, combine_shapes(9, &[2, 3]));
342 assert_ne!(shape, combine_shapes(1, &[2, 4]));
343 assert_ne!(shape, combine_shapes(1, &[2]));
344 }
345}