arora-engine 5.3.0

use std::{collections::HashMap, fmt::Debug, ops::DerefMut, pin::Pin, rc::Rc};

use arora_buffers::serde_uuid::deserialize;
use arora_types::{
    call::{Call, CallResult},
    value::Value,
};
use uuid::Uuid;

use crate::{
    call::{serialize_to_arg, CallBridge, CallError, Callable, CallableId, CallableRegistry},
    compiled::CompiledModule,
    executor::{self, Executor},
    module::{DispatchError, Module},
    schema::module::low::ModuleDefinition,
};

use derive_more::{Display, From};

pub struct EngineBuilder {
    executors: HashMap<&'static str, Box<dyn Executor>>,
}

impl Default for EngineBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl EngineBuilder {
    pub fn new() -> Self {
        Self {
            executors: HashMap::new(),
        }
    }

    pub fn add_executor<E: Executor + 'static>(mut self, executor: E) -> Self {
        self.executors.insert(executor.name(), Box::new(executor));
        self
    }

    pub fn build(self) -> Pin<Box<Engine>> {
        Engine::new(self.executors)
    }
}

#[derive(Debug, Display, From, Clone)]
pub enum LoadModuleError {
    ExecutorNotFound,
    MalformedExecutable,
    Internal(String),
}

impl std::error::Error for LoadModuleError {}

impl From<executor::LoadModuleError> for LoadModuleError {
    fn from(e: executor::LoadModuleError) -> Self {
        match e {
            executor::LoadModuleError::MalformedExecutable => LoadModuleError::MalformedExecutable,
            executor::LoadModuleError::Internal(message) => LoadModuleError::Internal(message),
        }
    }
}

#[derive(Debug, Display, From)]
pub enum UnloadModuleError {
    ModuleNotFound,
    Internal(String),
    /// The executor that loaded the module cannot unload it
    /// ([`Executor::can_unload`]): the module stays loaded.
    #[display("the {executor} executor cannot unload a module it loaded")]
    #[from(ignore)]
    Unsupported {
        executor: &'static str,
    },
}

impl std::error::Error for UnloadModuleError {}

impl From<executor::UnloadModuleError> for UnloadModuleError {
    fn from(e: executor::UnloadModuleError) -> Self {
        match e {
            executor::UnloadModuleError::ModuleNotFound => UnloadModuleError::ModuleNotFound,
            executor::UnloadModuleError::Internal(message) => UnloadModuleError::Internal(message),
        }
    }
}

/// [`Engine::reload_module`] failed: nothing was replaced.
#[derive(Debug, Display, Clone)]
pub enum ReloadModuleError {
    /// No module is loaded under the id.
    ModuleNotFound,
    /// The executor that loaded the module cannot unload it
    /// ([`Executor::can_unload`]): the module stays loaded.
    #[display("the {executor} executor cannot unload a module it loaded")]
    Unsupported { executor: &'static str },
    /// The new version did not load; the old one is still loaded.
    Load(LoadModuleError),
}

impl std::error::Error for ReloadModuleError {}

impl From<LoadModuleError> for ReloadModuleError {
    fn from(e: LoadModuleError) -> Self {
        ReloadModuleError::Load(e)
    }
}

impl From<executor::LoadModuleError> for ReloadModuleError {
    fn from(e: executor::LoadModuleError) -> Self {
        ReloadModuleError::Load(e.into())
    }
}

/// [`Engine`] is the main encapsulation of the Arora runtime.
/// It consists of a set of [`Executor`]s and [`Module`]s.
pub struct Engine {
    executors: HashMap<&'static str, Box<dyn Executor>>,
    modules: HashMap<Uuid, Loaded>,
    callables: CallableRegistry,
}

/// A module the engine dispatches to, and the executor that loaded it —
/// `None` for one registered already instantiated.
struct Loaded {
    executor: Option<&'static str>,
    module: Box<dyn Module>,
}

impl Engine {
    /// Create a new [`Engine`] with the given [`Executor`]s.
    fn new(executors: HashMap<&'static str, Box<dyn Executor>>) -> Pin<Box<Engine>> {
        let mut ret = Box::pin(Engine {
            executors,
            modules: HashMap::new(),
            callables: CallableRegistry::new(),
        });

        {
            let engine = &mut *ret.as_mut() as *mut Engine;
            for executor in ret.executors.values_mut() {
                executor.set_engine(engine);
            }
        }

        ret
    }

    /// Load a [`Module`] from the given [`ModuleDefinition`].
    pub fn load_module(
        &mut self,
        module_definition: ModuleDefinition,
    ) -> Result<(), LoadModuleError> {
        let module_id = module_definition.header.id;
        let executor_name = module_definition.header.executor.name.as_str();

        if self.modules.contains_key(&module_id) {
            return Ok(());
        }

        // We haven't loaded the module yet.

        // Find the executor for the module.
        let executor = self
            .executors
            .get_mut(executor_name)
            .ok_or(LoadModuleError::ExecutorNotFound)?;

        let module = executor.load_module(module_definition)?;
        self.modules.insert(
            module_id,
            Loaded {
                executor: Some(executor.name()),
                module,
            },
        );

        Ok(())
    }

    /// Load a [`Module`] from `module`, compiled ahead: the executor its header
    /// names instantiates the compiled code instead of compiling the
    /// executable, so loading one [`CompiledModule`] into many engines compiles
    /// it once. Otherwise as [`load_module`](Self::load_module): an id already
    /// loaded is answered with Ok and left as it is.
    pub fn load_compiled_module(&mut self, module: &CompiledModule) -> Result<(), LoadModuleError> {
        let module_id = module.header().id;
        if self.modules.contains_key(&module_id) {
            return Ok(());
        }
        let executor = self
            .executors
            .get_mut(module.header().executor.name.as_str())
            .ok_or(LoadModuleError::ExecutorNotFound)?;
        let instance = executor.load_compiled_module(module)?;
        self.modules.insert(
            module_id,
            Loaded {
                executor: Some(executor.name()),
                module: instance,
            },
        );
        Ok(())
    }

    /// Replace the module loaded under `module`'s id with a fresh instance of
    /// `module`: from the next dispatch, its functions run the new code, in a
    /// new instance with its own memory and globals.
    ///
    /// The new instance is created before the old one is touched, so a module
    /// that fails to load leaves the old one loaded and serving, and the error
    /// is the load's. The old instance is dropped once replaced, which releases
    /// its memory. An id with no module loaded fails with
    /// [`ReloadModuleError::ModuleNotFound`]: loading is
    /// [`load_compiled_module`](Self::load_compiled_module)'s job. A module
    /// whose executor cannot unload it ([`Executor::can_unload`]) is refused
    /// with [`ReloadModuleError::Unsupported`], before anything loads.
    ///
    /// Not reachable from a dispatch: a module cannot reload itself (or any
    /// other) while a call into it is on the stack, since the instance it runs
    /// in would be dropped under it. Call it between dispatches.
    pub fn reload_module(&mut self, module: &CompiledModule) -> Result<(), ReloadModuleError> {
        let module_id = module.header().id;
        let loaded = self
            .modules
            .get(&module_id)
            .ok_or(ReloadModuleError::ModuleNotFound)?;
        if let Some(executor) = self.refusing_unload(loaded) {
            return Err(ReloadModuleError::Unsupported { executor });
        }
        let executor = self
            .executors
            .get_mut(module.header().executor.name.as_str())
            .ok_or(LoadModuleError::ExecutorNotFound)?;
        let instance = executor.load_compiled_module(module)?;
        self.modules.insert(
            module_id,
            Loaded {
                executor: Some(executor.name()),
                module: instance,
            },
        );
        Ok(())
    }

    /// Whether a module — loaded or registered — dispatches under `id`.
    pub fn contains_module(&self, id: &Uuid) -> bool {
        self.modules.contains_key(id)
    }

    /// Remove the module loaded under `id` from dispatch and drop its
    /// instance: a call to it then fails with
    /// [`DispatchError::ModuleNotFound`]. Like
    /// [`reload_module`](Self::reload_module), not reachable from a dispatch.
    ///
    /// Any module leaves this way, a host module registered with
    /// [`register_module`](Self::register_module) included: which modules a
    /// host lets go is the host's decision. A module whose executor cannot
    /// unload it ([`Executor::can_unload`]) is refused with
    /// [`UnloadModuleError::Unsupported`] and stays loaded.
    pub fn unload_module(&mut self, id: &Uuid) -> Result<(), UnloadModuleError> {
        let loaded = self
            .modules
            .get(id)
            .ok_or(UnloadModuleError::ModuleNotFound)?;
        if let Some(executor) = self.refusing_unload(loaded) {
            return Err(UnloadModuleError::Unsupported { executor });
        }
        self.modules.remove(id);
        Ok(())
    }

    /// The executor that loaded `loaded`, when it cannot unload it.
    fn refusing_unload(&self, loaded: &Loaded) -> Option<&'static str> {
        let executor = loaded.executor?;
        let can_unload = self
            .executors
            .get(executor)
            .is_none_or(|executor| executor.can_unload());
        (!can_unload).then_some(executor)
    }

    /// Register an already-instantiated [`Module`] under `id`, making its
    /// functions reachable through [`CallBridge::arora_call`] exactly like a
    /// loaded module's. Where [`load_module`](Self::load_module) instantiates a
    /// module from a definition through an [`Executor`], this hands one over
    /// directly — how host-side modules (e.g. a
    /// [`HostModule`](crate::module::HostModule) built from closures)
    /// enter the same dispatch as guest code. Registering an id that is
    /// already present replaces the module.
    pub fn register_module(&mut self, id: Uuid, module: Box<dyn Module>) {
        self.modules.insert(
            id,
            Loaded {
                executor: None,
                module,
            },
        );
    }

    /// Dispatch a method call to a module. `arg` must be a raw Arora Buffer.
    pub fn dispatch(
        &mut self,
        module_id: &Uuid,
        function_id: &Uuid,
        arg: &[u8],
    ) -> Result<Box<[u8]>, DispatchError> {
        let module = self
            .modules
            .get_mut(module_id)
            .ok_or(DispatchError::ModuleNotFound { id: *module_id })?;

        module.module.dispatch(function_id, arg)
    }
}

pub type EngineRef = *mut Engine;

impl CallBridge for Engine {
    fn arora_call(&mut self, call: Call) -> Result<CallResult, CallError> {
        let module = call.module_id.ok_or_else(|| CallError::Generic {
            message: "call is missing its module id".to_string(),
        })?;
        let call_id = call.id;
        let result_data = self
            .dispatch(&module, &call_id, serialize_to_arg(call).as_ref())
            .map_err(Into::<CallError>::into)?;
        if let Value::Structure(structure) = deserialize(result_data.as_ref()) {
            if call_id != structure.id {
                Err(CallError::Internal {
                    message: format!(
                        "result id {} differs from function id {}",
                        structure.id, call_id
                    ),
                })?
            }
            let mut ret = None;
            let mut mutated = Vec::with_capacity(structure.fields.len() - 1);
            for field in structure.fields {
                // First field must be the return value.
                if ret.is_none() {
                    assert_eq!(field.id, call_id);
                    ret = Some(*field.value);
                } else {
                    mutated.push(field);
                }
            }
            let ret = ret.ok_or(CallError::Internal {
                message: "call result did not contain a return value".to_string(),
            })?;
            Ok(CallResult { ret, mutated })
        } else {
            Err(CallError::Internal {
                message: "returned data was not a structure".to_string(),
            })
        }
    }

    fn arora_register_callable(&mut self, callable: Rc<dyn Callable>) -> CallableId {
        self.callables.register_callable(callable).unwrap()
    }

    fn arora_unregister_callable(&mut self, callable_id: &CallableId) {
        self.callables.unregister_callable(callable_id).unwrap()
    }

    fn arora_call_indirect(&mut self, callable_id: &CallableId) -> Result<Value, CallError> {
        self.callables.find_callable(callable_id)?.call(self)
    }
}

pub type PinnedEngine = Pin<Box<Engine>>;

impl CallBridge for PinnedEngine {
    fn arora_call(&mut self, call: Call) -> Result<CallResult, CallError> {
        self.deref_mut().arora_call(call)
    }

    fn arora_register_callable(&mut self, callable: Rc<dyn Callable>) -> CallableId {
        self.deref_mut().arora_register_callable(callable)
    }

    fn arora_unregister_callable(&mut self, callable_id: &CallableId) {
        self.deref_mut().arora_unregister_callable(callable_id)
    }

    fn arora_call_indirect(&mut self, callable_id: &CallableId) -> Result<Value, CallError> {
        self.deref_mut().arora_call_indirect(callable_id)
    }
}

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

    #[test]
    fn function_modules_dispatch_like_loaded_modules() {
        let mut engine = EngineBuilder::new().build();
        let module_id = Uuid::from_u128(1);
        let function = Uuid::from_u128(2);
        let module = ModuleBuilder::new(module_id)
            .function(function, |call: Call| {
                let Value::I32(n) = call.args[0].value.as_ref() else {
                    return Err(CallError::Guest {
                        message: "expected an i32".to_string(),
                    });
                };
                Ok(CallResult {
                    ret: Value::I32(n * 2),
                    mutated: Vec::new(),
                })
            })
            .build();
        engine.register_module(module.id(), Box::new(module));

        let call = Call {
            module_id: Some(module_id),
            id: function,
            args: vec![arora_types::value::StructureField {
                id: Uuid::from_u128(3),
                value: Box::new(Value::I32(21)),
            }],
        };
        let result = engine.arora_call(call).unwrap();
        assert_eq!(result.ret, Value::I32(42));

        // An unattached function id fails like a missing guest function.
        let miss = Call {
            module_id: Some(module_id),
            id: Uuid::from_u128(4),
            args: Vec::new(),
        };
        assert!(matches!(
            engine.arora_call(miss),
            Err(CallError::FunctionNotFound { .. })
        ));

        // An unloaded module is gone from dispatch.
        assert!(engine.contains_module(&module_id));
        engine.unload_module(&module_id).expect("unload");
        assert!(!engine.contains_module(&module_id));
        assert!(matches!(
            engine.arora_call(Call {
                module_id: Some(module_id),
                id: function,
                args: Vec::new(),
            }),
            Err(CallError::ModuleNotFound { .. })
        ));
        assert!(matches!(
            engine.unload_module(&module_id),
            Err(UnloadModuleError::ModuleNotFound)
        ));

        // An unregistered module id fails as before.
        let other = Uuid::from_u128(9);
        let elsewhere = Call {
            module_id: Some(other),
            id: function,
            args: Vec::new(),
        };
        assert!(matches!(
            engine.arora_call(elsewhere),
            Err(CallError::ModuleNotFound { .. })
        ));
    }
}