nightshade 0.8.0

A cross-platform data-oriented game engine.
Documentation
use crate::ecs::world::{Entity, World};
use crate::plugin::{EngineCommand, EngineEvent};
use anyhow::Result;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::mpsc::{Receiver, Sender, channel};
use wasmtime::{Caller, Engine, Linker, Module, Store};
use winit::event::{ElementState, MouseButton};
use winit::keyboard::KeyCode;

use super::async_results::AsyncResult;
use super::commands::{CommandContext, process_async_results, process_engine_commands};
use super::entity_mapping::EntityMapping;
use super::events::{
    PluginInstance, PluginState, send_bytes_to_plugin, send_engine_event_to_plugin,
};
use super::memory::read_plugin_memory;

type CustomLinkerSetup = Box<dyn FnMut(&mut Linker<PluginState>, &Engine) -> Result<()>>;

pub struct PluginRuntimeConfig {
    pub plugins_base_path: PathBuf,
    pub cleanup_interval_frames: u64,
}

impl Default for PluginRuntimeConfig {
    fn default() -> Self {
        Self {
            plugins_base_path: PathBuf::from("plugins"),
            cleanup_interval_frames: 60,
        }
    }
}

pub struct PluginRuntime {
    engine: Engine,
    plugins: Vec<PluginInstance>,
    plugin_id_to_index: HashMap<u64, usize>,
    entity_mapping: EntityMapping,
    texture_id_to_name: HashMap<u64, String>,
    plugins_base_path: PathBuf,
    canonical_base_path: PathBuf,
    async_sender: Sender<AsyncResult>,
    async_receiver: Receiver<AsyncResult>,
    pending_input_events: Vec<EngineEvent>,
    last_mouse_position: (f32, f32),
    next_texture_id: u64,
    next_plugin_id: u64,
    frames_since_cleanup: u64,
    cleanup_interval: u64,
    custom_linker_setup: Option<CustomLinkerSetup>,
}

impl PluginRuntime {
    pub fn new(config: PluginRuntimeConfig) -> Result<Self> {
        let engine = Engine::default();
        let (async_sender, async_receiver) = channel();
        let canonical_base_path = std::fs::canonicalize(&config.plugins_base_path)
            .unwrap_or_else(|_| config.plugins_base_path.clone());
        Ok(Self {
            engine,
            plugins: Vec::new(),
            plugin_id_to_index: HashMap::new(),
            entity_mapping: EntityMapping::new(),
            texture_id_to_name: HashMap::new(),
            plugins_base_path: config.plugins_base_path,
            canonical_base_path,
            async_sender,
            async_receiver,
            pending_input_events: Vec::new(),
            last_mouse_position: (0.0, 0.0),
            next_texture_id: 1,
            next_plugin_id: 1,
            frames_since_cleanup: 0,
            cleanup_interval: config.cleanup_interval_frames,
            custom_linker_setup: None,
        })
    }

    pub fn with_custom_linker<F>(&mut self, setup: F)
    where
        F: FnMut(&mut Linker<PluginState>, &Engine) -> Result<()> + 'static,
    {
        self.custom_linker_setup = Some(Box::new(setup));
    }

    pub fn load_plugins_from_directory(&mut self, dir: &Path) -> Result<()> {
        if !dir.exists() {
            tracing::info!("Plugin directory does not exist: {:?}", dir);
            return Ok(());
        }

        for entry in std::fs::read_dir(dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.extension().is_some_and(|ext| ext == "wasm") {
                match self.load_plugin(&path) {
                    Ok(_) => tracing::info!("Loaded plugin: {:?}", path),
                    Err(error) => tracing::error!("Failed to load plugin {:?}: {}", path, error),
                }
            }
        }

        Ok(())
    }

    fn load_plugin(&mut self, path: &Path) -> Result<()> {
        let module = Module::from_file(&self.engine, path)?;
        let mut linker: Linker<PluginState> = Linker::new(&self.engine);

        wasmtime_wasi::p1::add_to_linker_sync(&mut linker, |state| &mut state.wasi)?;

        linker.func_wrap(
            "env",
            "host_send_command",
            |mut caller: Caller<'_, PluginState>, ptr: u32, len: u32| {
                if let Some(bytes) = read_plugin_memory(&mut caller, ptr, len) {
                    if let Some(cmd) = EngineCommand::from_bytes(&bytes) {
                        caller.data_mut().pending_commands.push(cmd);
                    } else {
                        tracing::warn!("Failed to deserialize plugin command");
                    }
                } else {
                    tracing::error!("Plugin command buffer out of bounds");
                }
            },
        )?;

        if let Some(ref mut custom_setup) = self.custom_linker_setup {
            custom_setup(&mut linker, &self.engine)?;
        }

        let wasi = wasmtime_wasi::WasiCtxBuilder::new()
            .inherit_stdio()
            .build_p1();

        let state = PluginState::new(wasi);
        let mut store = Store::new(&self.engine, state);
        let instance = linker.instantiate(&mut store, &module)?;

        let on_init = instance
            .get_typed_func::<(), ()>(&mut store, "on_init")
            .ok();
        let on_frame = instance
            .get_typed_func::<(), ()>(&mut store, "on_frame")
            .ok();

        let plugin_id = self.generate_plugin_id();
        let index = self.plugins.len();

        self.plugins.push(PluginInstance {
            id: plugin_id,
            store,
            instance,
            on_init,
            on_frame,
        });

        self.plugin_id_to_index.insert(plugin_id, index);

        Ok(())
    }

    pub fn call_on_init(&mut self, world: &mut World) {
        let commands = self.collect_init_commands();
        self.process_commands(world, commands);
    }

    fn collect_init_commands(&mut self) -> Vec<(u64, EngineCommand)> {
        let mut commands = Vec::new();

        for plugin in &mut self.plugins {
            if let Some(ref on_init) = plugin.on_init {
                if let Err(error) = on_init.call(&mut plugin.store, ()) {
                    tracing::error!("Plugin {} on_init failed: {}", plugin.id, error);
                }
                for command in plugin.store.data_mut().pending_commands.drain(..) {
                    commands.push((plugin.id, command));
                }
            }
        }

        commands
    }

    pub fn run_frame(&mut self, world: &mut World) {
        process_async_results(
            world,
            &mut self.plugins,
            &self.plugin_id_to_index,
            &mut self.entity_mapping,
            &mut self.texture_id_to_name,
            &self.async_receiver,
        );

        self.cleanup_stale_entities(world);
        self.flush_pending_input_events();

        let delta_time = world.resources.window.timing.delta_time;
        let frame_count = world.resources.window.timing.frame_counter as u64;

        self.dispatch_event_to_all(&EngineEvent::FrameStart {
            delta_time,
            frame_count,
        });

        let mouse_pos = world.resources.input.mouse.position;
        let new_pos = (mouse_pos.x, mouse_pos.y);
        if new_pos != self.last_mouse_position {
            self.last_mouse_position = new_pos;
            self.dispatch_event_to_all(&EngineEvent::MouseMoved {
                x: new_pos.0,
                y: new_pos.1,
            });
        }

        let commands = self.collect_frame_commands();
        self.process_commands(world, commands);
    }

    fn collect_frame_commands(&mut self) -> Vec<(u64, EngineCommand)> {
        let mut commands = Vec::new();

        for plugin in &mut self.plugins {
            if let Some(ref on_frame) = plugin.on_frame {
                if let Err(error) = on_frame.call(&mut plugin.store, ()) {
                    tracing::error!("Plugin {} on_frame failed: {}", plugin.id, error);
                }
                for command in plugin.store.data_mut().pending_commands.drain(..) {
                    commands.push((plugin.id, command));
                }
            }
        }

        commands
    }

    fn process_commands(&mut self, world: &mut World, commands: Vec<(u64, EngineCommand)>) {
        let mut context = CommandContext {
            plugins: &mut self.plugins,
            plugin_id_to_index: &self.plugin_id_to_index,
            entity_mapping: &mut self.entity_mapping,
            texture_id_to_name: &mut self.texture_id_to_name,
            async_sender: &self.async_sender,
            next_texture_id: &mut self.next_texture_id,
            base_path: self.plugins_base_path.clone(),
            canonical_base_path: self.canonical_base_path.clone(),
        };
        process_engine_commands(world, &mut context, commands);
    }

    pub fn queue_keyboard_event(&mut self, key_code: KeyCode, state: ElementState) {
        let event = if state.is_pressed() {
            EngineEvent::KeyPressed {
                key_code: key_code as u32,
            }
        } else {
            EngineEvent::KeyReleased {
                key_code: key_code as u32,
            }
        };
        self.pending_input_events.push(event);
    }

    pub fn queue_mouse_event(&mut self, state: ElementState, button: MouseButton) {
        let button_id = match button {
            MouseButton::Left => 0,
            MouseButton::Right => 1,
            MouseButton::Middle => 2,
            MouseButton::Back => 3,
            MouseButton::Forward => 4,
            MouseButton::Other(id) => id as u32 + 5,
        };

        let event = match state {
            ElementState::Pressed => EngineEvent::MouseButtonPressed { button: button_id },
            ElementState::Released => EngineEvent::MouseButtonReleased { button: button_id },
        };
        self.pending_input_events.push(event);
    }

    fn flush_pending_input_events(&mut self) {
        let events = std::mem::take(&mut self.pending_input_events);
        for event in events {
            self.dispatch_event_to_all(&event);
        }
    }

    fn cleanup_stale_entities(&mut self, world: &World) {
        self.frames_since_cleanup += 1;
        if self.frames_since_cleanup < self.cleanup_interval {
            return;
        }
        self.frames_since_cleanup = 0;

        self.entity_mapping.cleanup_invalid(|entity| {
            world.get_local_transform(entity).is_some()
                || world.get_global_transform(entity).is_some()
        });
    }

    pub fn dispatch_event_to_all(&mut self, event: &EngineEvent) {
        for plugin in &mut self.plugins {
            send_engine_event_to_plugin(plugin, event);
        }
    }

    pub fn dispatch_event_to_plugin(&mut self, plugin_id: u64, event: &EngineEvent) {
        if let Some(&index) = self.plugin_id_to_index.get(&plugin_id)
            && let Some(plugin) = self.plugins.get_mut(index)
        {
            send_engine_event_to_plugin(plugin, event);
        }
    }

    pub fn dispatch_custom_event_to_plugin(
        &mut self,
        plugin_id: u64,
        event_bytes: &[u8],
        alloc_fn: &str,
        receive_fn: &str,
    ) {
        if let Some(&index) = self.plugin_id_to_index.get(&plugin_id)
            && let Some(plugin) = self.plugins.get_mut(index)
        {
            send_bytes_to_plugin(plugin, event_bytes, alloc_fn, receive_fn);
        }
    }

    pub fn dispatch_custom_event_to_all<F>(
        &mut self,
        event_bytes: &[u8],
        alloc_fn: &str,
        receive_fn: &str,
        filter: F,
    ) where
        F: Fn(&PluginInstance) -> bool,
    {
        for plugin in &mut self.plugins {
            if filter(plugin) {
                send_bytes_to_plugin(plugin, event_bytes, alloc_fn, receive_fn);
            }
        }
    }

    pub fn dispatch_custom_event_to_all_unfiltered(
        &mut self,
        event_bytes: &[u8],
        alloc_fn: &str,
        receive_fn: &str,
    ) {
        for plugin in &mut self.plugins {
            send_bytes_to_plugin(plugin, event_bytes, alloc_fn, receive_fn);
        }
    }

    pub fn drain_custom_commands(&mut self) -> Vec<(u64, Vec<u8>)> {
        let mut all_commands = Vec::new();
        for plugin in &mut self.plugins {
            let plugin_id = plugin.id;
            for bytes in plugin.store.data_mut().pending_custom_commands.drain(..) {
                all_commands.push((plugin_id, bytes));
            }
        }
        all_commands
    }

    pub fn register_entity(&mut self, entity: Entity) -> u64 {
        self.entity_mapping.register(entity)
    }

    pub fn get_entity(&self, plugin_entity_id: u64) -> Option<Entity> {
        self.entity_mapping.get(plugin_entity_id)
    }

    pub fn unregister_entity(&mut self, plugin_entity_id: u64) {
        self.entity_mapping.unregister(plugin_entity_id);
    }

    pub fn plugin_has_export(&mut self, plugin_id: u64, name: &str) -> bool {
        if let Some(&index) = self.plugin_id_to_index.get(&plugin_id)
            && let Some(plugin) = self.plugins.get_mut(index)
        {
            return plugin
                .instance
                .get_typed_func::<u32, u32>(&mut plugin.store, name)
                .is_ok();
        }
        false
    }

    pub fn plugin_ids(&self) -> Vec<u64> {
        self.plugins.iter().map(|p| p.id).collect()
    }

    fn generate_plugin_id(&mut self) -> u64 {
        let id = self.next_plugin_id;
        self.next_plugin_id += 1;
        id
    }
}