himada-dispatch 0.1.1

use std::any::Any;
use crate::{KernelInfo, Dispatch, FnBench};
use himada_core::HardwareDNA;

thread_local! {
    static PLUGIN_STORE: std::cell::RefCell<Vec<Box<dyn Any + Send>>> =
        const { std::cell::RefCell::new(Vec::new()) };
}

#[cfg(test)]
fn clear_plugin_store() {
    PLUGIN_STORE.with(|s| *s.borrow_mut() = Vec::new());
}

pub struct PluginRegistry<F> {
    kernels: Vec<KernelInfo<F>>,
}

impl<F> PluginRegistry<F> {
    pub fn new() -> Self {
        Self { kernels: Vec::new() }
    }

    pub fn register(
        &mut self,
        name: &'static str,
        func: F,
        is_supported: fn(&HardwareDNA) -> bool,
        thermal_priority: u8,
    ) {
        self.kernels.push(KernelInfo { name, func, is_supported, thermal_priority });
    }

    pub fn kernels(&self) -> &[KernelInfo<F>] {
        &self.kernels
    }
}

pub trait KernelPlugin<F>: Send + Sync {
    fn register(&self, registry: &mut PluginRegistry<F>);
}

pub fn load_plugins<F: 'static + Send + Clone>() -> Option<Vec<KernelInfo<F>>> {
    PLUGIN_STORE.with(|store| {
        for item in store.borrow().iter() {
            if let Some(registry) = item.downcast_ref::<PluginRegistry<F>>() {
                let kernels = registry.kernels().to_vec();
                if !kernels.is_empty() {
                    return Some(kernels);
                }
            }
        }
        None
    })
}

pub fn register_plugin_registry<F: 'static + Send>(registry: PluginRegistry<F>) {
    PLUGIN_STORE.with(|store| {
        store.borrow_mut().push(Box::new(registry));
    });
}

impl<F: FnBench> Dispatch<F> {
    pub fn attach_plugins(&mut self) {
        if let Some(extra) = load_plugins::<F>() {
            self.supported.extend(extra);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::DotKernel;

    struct TestPlugin;

    impl KernelPlugin<DotKernel> for TestPlugin {
        fn register(&self, registry: &mut PluginRegistry<DotKernel>) {
            registry.register("test_dot", |a, b| {
                a.iter().zip(b).map(|(x, y)| x * y).sum()
            }, |_: &HardwareDNA| true, 1);
        }
    }

    #[test]
    fn test_plugin_registration() {
        let mut reg = PluginRegistry::<DotKernel>::new();
        let plugin = TestPlugin;
        plugin.register(&mut reg);
        assert_eq!(reg.kernels().len(), 1);
        assert_eq!(reg.kernels()[0].name, "test_dot");
    }

    #[test]
    fn test_plugin_registry_store() {
        clear_plugin_store();
        let mut reg = PluginRegistry::<DotKernel>::new();
        reg.register("stored_dot", |a, b| {
            a.iter().zip(b).map(|(x, y)| x * y).sum()
        }, |_: &HardwareDNA| true, 1);
        register_plugin_registry(reg);
        let loaded = load_plugins::<DotKernel>();
        assert!(loaded.is_some());
        assert_eq!(loaded.unwrap().len(), 1);
    }

    #[test]
    fn test_dispatch_attach_plugins() {
        clear_plugin_store();
        let mut reg = PluginRegistry::<DotKernel>::new();
        reg.register("plugin_dot", |a, b| {
            a.iter().zip(b).map(|(x, y)| x * y).sum()
        }, |_: &HardwareDNA| true, 1);
        register_plugin_registry(reg);

        let mut dispatch = Dispatch::new("test", vec![
            KernelInfo {
                name: "scalar",
                func: (|a: &[f64], b: &[f64]| a.iter().zip(b).map(|(x, y)| x * y).sum()) as fn(&[f64], &[f64]) -> f64,
                is_supported: (|_: &HardwareDNA| true) as fn(&HardwareDNA) -> bool,
                thermal_priority: 1,
            },
        ]);
        dispatch.attach_plugins();
        assert_eq!(dispatch.supported.len(), 2);
        assert_eq!(dispatch.supported[1].name, "plugin_dot");
    }
}