rlua-searcher/src/searcher.rs

use rlua::prelude::LuaError;
use rlua::{Context, MetaMethod, RegistryKey, Table, UserData, UserDataMethods, Value};
use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use std::path::{Path, PathBuf};

use crate::types::Result;
use crate::utils;

/// Stores Lua modules indexed by module name, and provides an `rlua::MetaMethod`
/// to enable `require`ing the stored modules by name in an `rlua::Context`.
struct Searcher {
    /// A `HashMap` of Lua modules in `String` representation, indexed by module name.
    modules: HashMap<String, String>,

    /// An `rlua::RegistryKey` whose value is the Lua environment within which the
    /// user made the request to instantiate a `Searcher` for `modules`.
    globals: RegistryKey,
}

impl Searcher {
    fn new(modules: HashMap<String, String>, globals: RegistryKey) -> Self {
        Self { modules, globals }
    }
}

impl UserData for Searcher {
    fn add_methods<'lua, M: UserDataMethods<'lua, Self>>(methods: &mut M) {
        methods.add_meta_method(MetaMethod::Call, |lua_ctx, this, name: String| {
            match this.modules.get(&name) {
                Some(content) => Ok(Value::Function(
                    lua_ctx
                        .load(content)
                        .set_name(&name)?
                        .set_environment(lua_ctx.registry_value::<Table>(&this.globals)?)?
                        .into_function()?,
                )),
                None => Ok(Value::Nil),
            }
        });
    }
}

/// Like `Searcher`, but with `modules` values encoded as `&'static str` to facilitate
/// compile-time includes of Fennel source code.
struct StaticSearcher {
    modules: HashMap<&'static str, &'static str>,
    globals: RegistryKey,
}

impl StaticSearcher {
    fn new(modules: HashMap<&'static str, &'static str>, globals: RegistryKey) -> Self {
        Self { modules, globals }
    }
}

impl UserData for StaticSearcher {
    fn add_methods<'lua, M: UserDataMethods<'lua, Self>>(methods: &mut M) {
        methods.add_meta_method(MetaMethod::Call, |lua_ctx, this, name: String| {
            match this.modules.get(name.as_str()) {
                Some(content) => Ok(Value::Function(
                    lua_ctx
                        .load(content)
                        .set_name(&name)?
                        .set_environment(lua_ctx.registry_value::<Table>(&this.globals)?)?
                        .into_function()?,
                )),
                None => Ok(Value::Nil),
            }
        });
    }
}

/// Like `Searcher`, but with `modules` values given as paths to files the content of
/// which can be read as Lua source code.
///
/// Facilitates Lua module reloading, and module reloading of any other programming
/// language whose source code can be compiled to Lua.
struct PolySearcher<P>
where
    P: 'static + AsRef<Path> + Send,
{
    modules: HashMap<String, P>,
    globals: RegistryKey,

    /// Function to read file content as Lua source code.
    convert: Box<dyn Fn(PathBuf) -> rlua::Result<String> + Send>,
}

impl<P> PolySearcher<P>
where
    P: 'static + AsRef<Path> + Send,
{
    fn new(
        modules: HashMap<String, P>,
        globals: RegistryKey,
        convert: Box<dyn Fn(PathBuf) -> rlua::Result<String> + Send>,
    ) -> Self {
        Self {
            modules,
            globals,
            convert,
        }
    }
}

impl<P> UserData for PolySearcher<P>
where
    P: 'static + AsRef<Path> + Send,
{
    fn add_methods<'lua, M: UserDataMethods<'lua, Self>>(methods: &mut M) {
        methods.add_meta_method(MetaMethod::Call, |lua_ctx, this, name: String| {
            match this.modules.get(&name) {
                Some(ref path) => {
                    let path = path.as_ref();

                    // Ensure `path` is relative to `$CARGO_MANIFEST_DIR`.
                    let path = if path.is_relative() {
                        utils::runtime_root().join(path)
                    } else {
                        path.to_path_buf()
                    };

                    let content = (this.convert)(path)?;

                    Ok(Value::Function(
                        lua_ctx
                            .load(&content)
                            .set_name(&name)?
                            .set_environment(lua_ctx.registry_value::<Table>(&this.globals)?)?
                            .into_function()?,
                    ))
                }
                None => Ok(Value::Nil),
            }
        });
    }
}

/// Extend `rlua::Context` to support `require`ing Lua modules by name.
pub trait AddSearcher {
    /// Add a `HashMap` of Lua modules indexed by module name to Lua’s
    /// `package.searchers` table in an `rlua::Context`, with lookup functionality
    /// provided by the `rlua_searcher::Searcher` struct.
    fn add_searcher(&self, modules: HashMap<String, String>) -> Result<()>;

    /// Like `add_searcher`, but with Fennel source code encoded as `&'static str`
    /// to facilitate compile-time includes.
    fn add_static_searcher(&self, modules: HashMap<&'static str, &'static str>) -> Result<()>;

    /// Like `add_searcher`, but with `modules` values given as paths to files containing
    /// Lua source code to facilitate module reloading.
    fn add_path_searcher<P>(&self, modules: HashMap<String, P>) -> Result<()>
    where
        P: 'static + AsRef<Path> + Send;

    /// Like `add_path_searcher`, but with user-provided closure for converting source
    /// code to Lua.
    fn add_poly_searcher<P>(
        &self,
        modules: HashMap<String, P>,
        convert: Box<dyn Fn(PathBuf) -> rlua::Result<String> + Send>,
    ) -> Result<()>
    where
        P: 'static + AsRef<Path> + Send;
}

impl<'a> AddSearcher for Context<'a> {
    fn add_searcher(&self, modules: HashMap<String, String>) -> Result<()> {
        let globals = self.globals();
        let searchers: Table = globals.get::<_, Table>("package")?.get("searchers")?;
        let registry_key = self.create_registry_value(globals)?;
        let searcher = Searcher::new(modules, registry_key);
        searchers
            .set(searchers.len()? + 1, searcher)
            .map_err(|e| e.into())
    }

    fn add_static_searcher(&self, modules: HashMap<&'static str, &'static str>) -> Result<()> {
        let globals = self.globals();
        let searchers: Table = globals.get::<_, Table>("package")?.get("searchers")?;
        let registry_key = self.create_registry_value(globals)?;
        let searcher = StaticSearcher::new(modules, registry_key);
        searchers
            .set(searchers.len()? + 1, searcher)
            .map_err(|e| e.into())
    }

    fn add_path_searcher<P>(&self, modules: HashMap<String, P>) -> Result<()>
    where
        P: 'static + AsRef<Path> + Send,
    {
        let convert = Box::new(|path| {
            let mut content = String::new();
            let mut file = File::open(path)
                .map_err(|e| LuaError::RuntimeError(format!("io error: {:#?}", e)))?;
            file.read_to_string(&mut content)
                .map_err(|e| LuaError::RuntimeError(format!("io error: {:#?}", e)))?;
            Ok(content)
        });
        self.add_poly_searcher(modules, convert)
    }

    fn add_poly_searcher<P>(
        &self,
        modules: HashMap<String, P>,
        convert: Box<dyn Fn(PathBuf) -> rlua::Result<String> + Send>,
    ) -> Result<()>
    where
        P: 'static + AsRef<Path> + Send,
    {
        let globals = self.globals();
        let searchers: Table = globals.get::<_, Table>("package")?.get("searchers")?;
        let registry_key = self.create_registry_value(globals)?;
        let searcher = PolySearcher::new(modules, registry_key, convert);
        searchers
            .set(searchers.len()? + 1, searcher)
            .map_err(|e| e.into())
    }
}