rusteal-core 0.7.0

use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock, RwLock};

use crate::{lock_or_recover, read_or_recover, write_or_recover};

pub struct ClassRegistration {
    pub type_id: u64,
    pub shape: u64,
    pub depends_on: &'static [u64],
    pub create: fn(last_try: bool, shape: u64) -> bool,
    pub register: fn(),
    pub finalize: fn(),
    pub after_defaults: fn(),
}
inventory::collect!(ClassRegistration);

pub struct StructRegistration {
    pub create: fn(),
    pub register: fn(),
    pub finalize: fn(),
}
inventory::collect!(StructRegistration);

pub struct EnumRegistration {
    pub create: fn(),
}
inventory::collect!(EnumRegistration);

pub struct ClassFunctionRegistration {
    pub type_id: u64,
    pub shape: u64,
    pub register_functions: fn(),
    pub class_defaults: Option<fn()>,
}
inventory::collect!(ClassFunctionRegistration);

pub trait ClassDefaultsOutcome {
    fn report(self, class: &str);
}

impl ClassDefaultsOutcome for () {
    fn report(self, _class: &str) {}
}

impl ClassDefaultsOutcome for crate::error::RustealResult<()> {
    fn report(self, class: &str) {
        if let Err(e) = self {
            e.__log_in(
                crate::logging::LOG_WARNING,
                &format!("{class}'s class defaults"),
            );
        }
    }
}

pub fn register_all_from_inventory() {
    for ereg in inventory::iter::<EnumRegistration> {
        (ereg.create)();
    }

    let mut pending: Vec<&ClassRegistration> =
        inventory::iter::<ClassRegistration>.into_iter().collect();

    let class_count = pending.len() as u32;
    let mut ordered: Vec<&ClassRegistration> = Vec::with_capacity(pending.len());

    loop {
        let before = pending.len();

        pending.retain(|reg| {
            if (reg.create)(false, class_shape(reg)) {
                ordered.push(reg);

                false
            } else {
                true
            }
        });

        if pending.is_empty() || pending.len() == before {
            break;
        }
    }

    for reg in pending {
        (reg.create)(true, class_shape(reg));
    }

    for sreg in inventory::iter::<StructRegistration> {
        (sreg.create)();
    }

    for sreg in inventory::iter::<StructRegistration> {
        (sreg.register)();
    }

    for sreg in inventory::iter::<StructRegistration> {
        (sreg.finalize)();
    }

    for reg in &ordered {
        (reg.register)();
    }

    let ordered = dependency_order(&ordered);

    let mut func_reg_count = 0u32;

    for freg in inventory::iter::<ClassFunctionRegistration> {
        (freg.register_functions)();
        func_reg_count += 1;
    }

    for reg in &ordered {
        (reg.finalize)();

        for freg in inventory::iter::<ClassFunctionRegistration> {
            if freg.type_id == reg.type_id
                && let Some(class_defaults) = freg.class_defaults
            {
                class_defaults();
            }
        }

        (reg.after_defaults)();
    }

    let total_funcs = read_or_recover(func_registry()).len();

    let msg = format!(
        "[Rusteal] register_all_from_inventory: {class_count} classes, {func_reg_count} impl blocks, {total_funcs} function callbacks",
    );

    let bytes = msg.as_bytes();

    unsafe {
        crate::ffi_dispatch::logging_log(0, bytes.as_ptr(), bytes.len() as u32);
    }
}

fn dependency_order<'a>(classes: &[&'a ClassRegistration]) -> Vec<&'a ClassRegistration> {
    let ids: Vec<u64> = classes.iter().map(|class| class.type_id).collect();
    let mut state = vec![0u8; classes.len()];
    let mut out = Vec::with_capacity(classes.len());

    fn visit<'a>(
        index: usize,
        classes: &[&'a ClassRegistration],
        ids: &[u64],
        state: &mut [u8],
        out: &mut Vec<&'a ClassRegistration>,
    ) {
        if state[index] != 0 {
            return;
        }

        state[index] = 1;

        for dependency in classes[index].depends_on {
            if let Some(found) = ids.iter().position(|id| id == dependency) {
                visit(found, classes, ids, state, out);
            }
        }

        state[index] = 2;
        out.push(classes[index]);
    }

    for index in 0..classes.len() {
        visit(index, classes, &ids, &mut state, &mut out);
    }

    out
}

fn class_shape(class: &ClassRegistration) -> u64 {
    let mut blocks: Vec<u64> = inventory::iter::<ClassFunctionRegistration>
        .into_iter()
        .filter(|block| block.type_id == class.type_id)
        .map(|block| block.shape)
        .collect();

    blocks.sort_unstable();
    combine_shapes(class.shape, &blocks)
}

fn combine_shapes(class: u64, blocks: &[u64]) -> u64 {
    let mut hash: u64 = 0xcbf29ce484222325;

    for part in std::iter::once(class).chain(blocks.iter().copied()) {
        for byte in part.to_le_bytes() {
            hash ^= byte as u64;
            hash = hash.wrapping_mul(0x100000001b3);
        }
    }

    hash
}

use rusteal_ffi::UObjectHandle;

pub struct RustTypeInfo {
    pub name: &'static str,
    pub construct_fn: fn() -> *mut u8,
    pub drop_fn: unsafe fn(*mut u8),
}

use crate::ffi_dispatch::NativePtr;

type ReifyFunctionCallback = Arc<dyn Fn(UObjectHandle, *mut u8, NativePtr) + Send + Sync>;

struct FunctionEntry {
    callback: ReifyFunctionCallback,
    type_id: u64,
}

static TYPE_REGISTRY: OnceLock<Mutex<HashMap<u64, RustTypeInfo>>> = OnceLock::new();
static FUNC_REGISTRY: OnceLock<RwLock<Vec<FunctionEntry>>> = OnceLock::new();
static INSTANCE_DATA: OnceLock<RwLock<HashMap<u64, Vec<InstanceEntry>>>> = OnceLock::new();

struct InstanceEntry {
    data: *mut u8,
    type_id: u64,
}

unsafe impl Send for InstanceEntry {}
unsafe impl Sync for InstanceEntry {}

fn type_registry() -> &'static Mutex<HashMap<u64, RustTypeInfo>> {
    TYPE_REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
}

fn func_registry() -> &'static RwLock<Vec<FunctionEntry>> {
    FUNC_REGISTRY.get_or_init(|| RwLock::new(Vec::new()))
}

fn instance_data() -> &'static RwLock<HashMap<u64, Vec<InstanceEntry>>> {
    INSTANCE_DATA.get_or_init(|| RwLock::new(HashMap::new()))
}

pub fn register_type(type_id: u64, info: RustTypeInfo) {
    lock_or_recover(type_registry()).insert(type_id, info);
}

pub fn register_function<F>(type_id: u64, f: F) -> u64
where
    F: Fn(UObjectHandle, *mut u8, NativePtr) + Send + Sync + 'static,
{
    let mut vec = write_or_recover(func_registry());
    let id = vec.len() as u64;

    vec.push(FunctionEntry {
        callback: Arc::new(f),
        type_id,
    });

    id
}

pub fn construct_instance(obj: UObjectHandle, type_id: u64) {
    let types = lock_or_recover(type_registry());

    let Some(info) = types.get(&type_id) else {
        if crate::api::is_api_initialized() {
            let msg = format!("[Rusteal] construct_instance: unknown type_id {type_id}");
            let bytes = msg.as_bytes();

            unsafe {
                crate::ffi_dispatch::logging_log(1, bytes.as_ptr(), bytes.len() as u32);
            }
        }

        return;
    };

    let data = (info.construct_fn)();
    drop(types);

    let key = obj.to_addr();

    let old = {
        let mut map = write_or_recover(instance_data());
        let entries = map.entry(key).or_default();
        match entries.iter_mut().find(|e| e.type_id == type_id) {
            Some(entry) => Some(std::mem::replace(&mut entry.data, data)),
            None => {
                entries.push(InstanceEntry { data, type_id });

                None
            }
        }
    };

    if let Some(old) = old {
        drop_data(type_id, old);
    }
}

fn drop_data(type_id: u64, data: *mut u8) {
    let types = lock_or_recover(type_registry());

    if let Some(info) = types.get(&type_id) {
        unsafe {
            (info.drop_fn)(data);
        }
    }
}

pub fn drop_instance(obj: UObjectHandle, _type_id: u64) {
    let key = obj.to_addr();
    let entries = write_or_recover(instance_data()).remove(&key);

    for entry in entries.into_iter().flatten() {
        drop_data(entry.type_id, entry.data);
    }
}

pub fn invoke_function(callback_id: u64, obj: UObjectHandle, params: NativePtr) {
    let func = {
        let vec = read_or_recover(func_registry());
        vec.get(callback_id as usize)
            .map(|entry| (entry.callback.clone(), entry.type_id))
    };

    if let Some((func, type_id)) = func {
        let rust_data = get_instance_data(obj, type_id);
        func(obj, rust_data, params);
    } else if crate::api::is_api_initialized() {
        let vec_len = read_or_recover(func_registry()).len();

        let msg = format!(
            "[Rusteal] invoke_function: callback_id {callback_id} not found (registry size = {vec_len})",
        );

        let bytes = msg.as_bytes();

        unsafe {
            crate::ffi_dispatch::logging_log(1, bytes.as_ptr(), bytes.len() as u32);
        }
    }
}

pub fn clear_all() {
    if let Some(instances) = INSTANCE_DATA.get() {
        let mut map = write_or_recover(instances);
        let types = lock_or_recover(type_registry());

        for entry in map.drain().flat_map(|(_, entries)| entries) {
            if let Some(info) = types.get(&entry.type_id) {
                unsafe {
                    (info.drop_fn)(entry.data);
                }
            }
        }

        drop(types);
    }

    if let Some(funcs) = FUNC_REGISTRY.get() {
        write_or_recover(funcs).clear();
    }

    if let Some(types) = TYPE_REGISTRY.get() {
        lock_or_recover(types).clear();
    }
}

pub fn get_instance_data(obj: UObjectHandle, type_id: u64) -> *mut u8 {
    let key = obj.to_addr();

    read_or_recover(instance_data())
        .get(&key)
        .and_then(|entries| entries.iter().find(|e| e.type_id == type_id))
        .map(|e| e.data)
        .unwrap_or(std::ptr::null_mut())
}

#[cfg(test)]
mod tests {
    use super::combine_shapes;

    #[test]
    fn every_part_of_a_class_makes_its_shape() {
        let shape = combine_shapes(1, &[2, 3]);
        assert_eq!(shape, combine_shapes(1, &[2, 3]));
        assert_ne!(shape, combine_shapes(9, &[2, 3]));
        assert_ne!(shape, combine_shapes(1, &[2, 4]));
        assert_ne!(shape, combine_shapes(1, &[2]));
    }
}

#[cfg(test)]
mod order_tests {
    use super::*;

    fn noop() {}

    fn create(_last_try: bool, _shape: u64) -> bool {
        true
    }

    const fn class(type_id: u64, depends_on: &'static [u64]) -> ClassRegistration {
        ClassRegistration {
            type_id,
            shape: 0,
            depends_on,
            create,
            register: noop,
            finalize: noop,
            after_defaults: noop,
        }
    }

    #[test]
    fn component_and_parent_classes_come_first() {
        static GRABBABLE: ClassRegistration = class(1, &[100, 2]);
        static GRAB: ClassRegistration = class(2, &[101]);
        static PISTOL: ClassRegistration = class(3, &[100, 2, 4]);
        static PROJECTILE: ClassRegistration = class(4, &[100]);

        let order: Vec<u64> = dependency_order(&[&GRABBABLE, &PISTOL, &GRAB, &PROJECTILE])
            .iter()
            .map(|class| class.type_id)
            .collect();

        let position = |id| order.iter().position(|o| *o == id).unwrap();
        assert!(position(2) < position(1));
        assert!(position(2) < position(3));
        assert!(position(4) < position(3));
        assert_eq!(order.len(), 4);
    }

    #[test]
    fn cycles_still_order_every_class() {
        static A: ClassRegistration = class(1, &[2]);
        static B: ClassRegistration = class(2, &[1]);

        assert_eq!(dependency_order(&[&A, &B]).len(), 2);
    }
}