webcomposer/

lua.rs

use crate::core::element::{Element, Render};
pub use mlua::{IntoLua, Lua, LuaSerdeExt, Result as LuaResult, Value as LuaValue};
use mlua::{ObjectLike, prelude::*};
use serde::{Serialize, de::DeserializeOwned};
use std::path::PathBuf;
use uuid::Uuid;

#[cfg(feature = "web")]
use std::sync::{LazyLock, RwLock};
#[cfg(feature = "web")]
use tera::Tera;

#[cfg(feature = "web")]
static TERA: LazyLock<RwLock<Tera>> = LazyLock::new(|| RwLock::new(Tera::default()));

#[macro_export]
macro_rules! wrap_lua_runtime_err {
    ($x:expr) => {
        match $x {
            Ok(x) => x,
            Err(e) => {
                return Err(LuaError::RuntimeError(e.to_string()));
            }
        }
    };

    ($x:expr, Option) => {
        match $x {
            Some(x) => x,
            None => {
                return Err(LuaError::RuntimeError("Unknown error".to_string()));
            }
        }
    };
}

/// Lua VM runtime.
pub struct LuaContext();

impl LuaContext {
    /// Implement file system stuff for the Lua vm since the built-in fs module sucks.
    fn lua_fs(lua: &Lua) -> LuaResult<LuaTable> {
        let v = lua.create_table()?;
        v.set(
            "read",
            lua.create_function(|l, p: String| {
                Ok(wrap_lua_runtime_err!(std::fs::read_to_string(p)).into_lua(l)?)
            })?,
        )?;
        v.set(
            "write",
            lua.create_function(|_, (p, c): (String, String)| {
                wrap_lua_runtime_err!(std::fs::write(p, c));
                Ok(())
            })?,
        )?;
        v.set(
            "copy",
            lua.create_function(|_, (f, t): (String, String)| {
                wrap_lua_runtime_err!(std::fs::copy(f, t));
                Ok(())
            })?,
        )?;
        v.set(
            "create_dir",
            lua.create_function(|_, p: String| {
                wrap_lua_runtime_err!(std::fs::create_dir(p));
                Ok(())
            })?,
        )?;
        v.set(
            "create_dir_all",
            lua.create_function(|_, p: String| {
                wrap_lua_runtime_err!(std::fs::create_dir_all(p));
                Ok(())
            })?,
        )?;
        v.set(
            "remove_file",
            lua.create_function(|_, p: String| {
                wrap_lua_runtime_err!(std::fs::remove_file(p));
                Ok(())
            })?,
        )?;
        v.set(
            "remove_dir",
            lua.create_function(|_, p: String| {
                wrap_lua_runtime_err!(std::fs::remove_dir(p));
                Ok(())
            })?,
        )?;
        v.set(
            "rename",
            lua.create_function(|_, (f, t): (String, String)| {
                wrap_lua_runtime_err!(std::fs::rename(f, t));
                Ok(())
            })?,
        )?;
        v.set(
            "exists",
            lua.create_function(|_, path: String| {
                Ok(wrap_lua_runtime_err!(std::fs::exists(path)))
            })?,
        )?;
        v.set(
            "read_dir",
            lua.create_function(|_, path: String| {
                Ok(wrap_lua_runtime_err!(std::fs::read_dir(path))
                    .map(|x| x.unwrap().path().to_str().unwrap().to_string())
                    .collect::<Vec<String>>())
            })?,
        )?;
        v.set(
            "stat",
            lua.create_function(|l, path: String| {
                let v = l.create_table()?;
                let stat = wrap_lua_runtime_err!(std::fs::metadata(path));

                v.set("is_file", stat.is_file())?;
                v.set("is_dir", stat.is_dir())?;
                v.set("is_symlink", stat.is_symlink())?;
                v.set("len", stat.len())?;

                Ok(v)
            })?,
        )?;
        Ok(v)
    }

    /// WebComposer Lisp Markup
    fn lua_wclm(lua: &Lua) -> LuaResult<LuaTable> {
        let v = lua.create_table()?;
        v.set(
            "parse",
            lua.create_function(|l, v: LuaString| {
                Ok(l.to_value(&crate::parse(&wrap_lua_runtime_err!(v.to_str())))?)
            })?,
        )?;
        v.set(
            "render",
            lua.create_function(|l, v: LuaTable| {
                Ok(l.from_value::<Element>(v.to_value())?.render_safe())
            })?,
        )?;
        Ok(v)
    }

    /// Tera integration.
    #[cfg(feature = "web")]
    fn lua_tera(lua: &Lua) -> LuaResult<LuaTable> {
        let v = lua.create_table()?;
        v.set(
            "add",
            lua.create_function(|_, (name, value): (String, String)| {
                let mut tera = TERA.write().expect("failed to lock tera state for write");
                wrap_lua_runtime_err!(tera.add_raw_template(&name, &value));
                Ok(())
            })?,
        )?;
        v.set(
            "render",
            lua.create_function(|l, (tmpl, ctx): (String, LuaTable)| {
                let tera = TERA.read().expect("failed to lock tera state for read");
                Ok(wrap_lua_runtime_err!(tera.render(
                    &tmpl,
                    &wrap_lua_runtime_err!(tera::Context::from_value(
                        l.from_value::<serde_json::Value>(ctx.to_value())?,
                    )),
                ))
                .into_lua(l)?)
            })?,
        )?;
        Ok(v)
    }

    /// Language features.
    fn lua_features(lua: &Lua, args: String) -> LuaResult<LuaTable> {
        let v = lua.create_table()?;
        v.set("NULL", lua.null())?;
        v.set("args", {
            let t = lua.create_table()?;

            for arg in args.split(" ") {
                t.push(arg)?;
            }

            t
        })?;
        v.set("rargs", args)?;
        v.set(
            "match",
            lua.create_function(|_, (value, matches): (LuaValue, LuaTable)| {
                // search for match
                if let Ok(x) = matches.get::<LuaValue>(&value)
                    && !x.is_nil()
                {
                    if x.is_function() {
                        return Ok(x.as_function().unwrap().call::<LuaValue>(())?);
                    } else {
                        return Ok(x);
                    }
                }

                // catch-all
                if let Ok(x) = matches.get::<LuaValue>("_")
                    && !x.is_nil()
                {
                    if x.is_function() {
                        return Ok(x.as_function().unwrap().call::<LuaValue>(())?);
                    } else {
                        return Ok(x);
                    }
                }

                // error (no catch-all and no match)
                Err(LuaError::RuntimeError(format!(
                    "No matches available ({})",
                    value.as_string().unwrap().to_str()?
                )))
            })?,
        )?;
        v.set(
            "throw",
            lua.create_function(|_, value: String| {
                Err::<(), LuaError>(LuaError::RuntimeError(value))
            })?,
        )?;
        v.set("macros", {
            let v = lua.create_table()?;
            v.set(
                "define",
                lua.create_function(|l, (name, func): (String, LuaFunction)| {
                    l.globals().set(format!("__WEBC_MACROS_{name}"), func)?;
                    Ok(())
                })?,
            )?;
            v.set(
                "call",
                lua.create_async_function(async |l, (name, args): (String, LuaTable)| {
                    let f: LuaFunction = l.globals().get(format!("__WEBC_MACROS_{name}"))?;
                    Ok(f.call_async::<LuaTable>(args).await?)
                })?,
            )?;
            v
        })?;
        Ok(v)
    }

    fn lua_tools(lua: &Lua) -> LuaResult<LuaTable> {
        let v = lua.create_table()?;
        v.set(
            "sha256",
            lua.create_function(|_, input: String| Ok(crate::globber::hash(input)))?,
        )?;
        v.set(
            "uuid",
            lua.create_function(|_, ()| Ok(Uuid::new_v4().to_string()))?,
        )?;
        v.set(
            "uuid7",
            lua.create_function(|_, ()| Ok(Uuid::now_v7().to_string()))?,
        )?;
        v.set(
            "id",
            lua.create_function(|_, ()| Ok(tritools::id::Id::new().printable()))?,
        )?;
        v.set(
            "salt",
            lua.create_function(|_, ()| Ok(tritools::encoding::salt()))?,
        )?;
        v.set(
            "salt_len",
            lua.create_function(|_, x: usize| Ok(tritools::encoding::salt_len(x)))?,
        )?;
        v.set(
            "unix_epoch_timestamp",
            lua.create_function(|_, ()| Ok(tritools::time::unix_epoch_timestamp()))?,
        )?;
        v.set(
            "epoch_timestamp",
            lua.create_function(|_, e: u32| Ok(tritools::time::epoch_timestamp(e)))?,
        )?;
        v.set(
            "compress",
            lua.create_function(|_, x: Vec<u8>| Ok(tritools::encoding::compress(x)))?,
        )?;
        v.set(
            "decompress",
            lua.create_function(|_, x: Vec<u8>| Ok(tritools::encoding::decompress(x)))?,
        )?;
        Ok(v)
    }

    fn lua_string(lua: &Lua) -> LuaResult<LuaTable> {
        let v = lua.create_table()?;
        v.set(
            "replace",
            lua.create_function(|_, (haystack, from, to): (String, String, String)| {
                Ok(haystack.replace(&from, &to).to_string())
            })?,
        )?;
        v.set(
            "replacen",
            lua.create_function(
                |_, (haystack, from, to, n): (String, String, String, usize)| {
                    Ok(haystack.replacen(&from, &to, n).to_string())
                },
            )?,
        )?;
        v.set(
            "trim",
            lua.create_function(|_, input: String| Ok(input.trim().to_string()))?,
        )?;
        v.set(
            "trim_end",
            lua.create_function(|_, input: String| Ok(input.trim_end().to_string()))?,
        )?;
        v.set(
            "trim_start",
            lua.create_function(|_, input: String| Ok(input.trim_start().to_string()))?,
        )?;
        v.set(
            "regex_does_match",
            lua.create_function(|_, (pat, haystack): (String, String)| {
                let regex = regex::RegexBuilder::new(&pat)
                    .multi_line(true)
                    .build()
                    .unwrap();

                Ok(regex.captures(&haystack).is_some())
            })?,
        )?;
        v.set(
            "chars",
            lua.create_function(|_, input: String| Ok(input.chars().collect::<Vec<char>>()))?,
        )?;
        v.set(
            "mime",
            lua.create_function(|_, input: String| {
                Ok(mime_guess::from_path(input).first().unwrap().to_string())
            })?,
        )?;
        v.set(
            "mime_ext",
            lua.create_function(|_, input: String| {
                Ok(mime_guess::from_ext(&input).first().unwrap().to_string())
            })?,
        )?;
        Ok(v)
    }

    fn lua_array(lua: &Lua) -> LuaResult<LuaTable> {
        let v = lua.create_table()?;
        v.set(
            "contains",
            lua.create_function(|_, (input, v): (LuaTable, LuaValue)| {
                Ok(input
                    .pairs::<LuaNumber, LuaValue>()
                    .find(|x| v == x.as_ref().unwrap().1)
                    .is_some())
            })?,
        )?;
        v.set(
            "push",
            lua.create_function(|_, (input, v): (LuaTable, LuaValue)| Ok(input.push(v)))?,
        )?;
        v.set(
            "pop",
            lua.create_function(|_, input: LuaTable| Ok(input.pop::<LuaValue>()))?,
        )?;
        v.set(
            "remove",
            lua.create_function(|_, (input, k): (LuaTable, LuaValue)| Ok(input.raw_remove(k)))?,
        )?;
        v.set(
            "concat",
            lua.create_function(|_, (a, b): (LuaTable, LuaTable)| {
                let a: Vec<(LuaValue, LuaValue)> = a.pairs().map(|x| x.unwrap()).collect();
                let b: Vec<(LuaValue, LuaValue)> = b.pairs().map(|x| x.unwrap()).collect();
                let c = &[a, b].concat();
                Ok(c.iter().map(|x| x.1.to_owned()).collect::<Vec<LuaValue>>())
            })?,
        )?;
        v.set(
            "concat_table",
            lua.create_function(|l, (a, b): (LuaTable, LuaTable)| {
                let a: Vec<(LuaValue, LuaValue)> = a.pairs().map(|x| x.unwrap()).collect();
                let b: Vec<(LuaValue, LuaValue)> = b.pairs().map(|x| x.unwrap()).collect();
                let c = &[a, b].concat();
                let d = l.create_table()?;

                for (x, y) in c {
                    d.set(x, y)?;
                }

                Ok(d)
            })?,
        )?;
        Ok(v)
    }

    fn lua_memory(lua: &Lua) -> LuaResult<LuaTable> {
        let v = lua.create_table()?;
        v.set(
            "clone",
            lua.create_function(|_, v: LuaValue| Ok(v.clone()))?,
        )?;
        v.set(
            "str_bytes",
            lua.create_function(|_, v: String| Ok(v.as_bytes().to_vec()))?,
        )?;
        v.set(
            "str_from_bytes",
            lua.create_function(|_, v: Vec<u8>| Ok(String::from_utf8_lossy(&v).to_string()))?,
        )?;
        v.set(
            "num_le_bytes",
            lua.create_function(|_, v: LuaNumber| Ok(v.to_le_bytes().to_vec()))?,
        )?;
        v.set(
            "num_be_bytes",
            lua.create_function(|_, v: LuaNumber| Ok(v.to_be_bytes().to_vec()))?,
        )?;
        v.set("used", lua.create_function(|l, ()| Ok(l.used_memory()))?)?;
        v.set(
            "size_of_val",
            lua.create_function(|_, v: LuaValue| Ok(std::mem::size_of_val(&v)))?,
        )?;
        v.set(
            "compile",
            lua.create_function(|_, v: String| {
                let c = mlua::Compiler::new();
                Ok(wrap_lua_runtime_err!(c.compile(v)))
            })?,
        )?;
        Ok(v)
    }

    /// Create a new [`Lua`].
    pub fn lua() -> Lua {
        Lua::new()
    }

    /// Execute the given lua script to and capture the returned `T`.
    pub fn lua_exec_with(
        lua: &Lua,
        script: Vec<u8>,
        entry_path: String,
        args: String,
    ) -> LuaResult<LuaValue> {
        let entry_path_path = PathBuf::from(entry_path.clone());

        // register modules
        lua.register_module("@webc/fs", Self::lua_fs(&lua)?)?;
        lua.register_module("@webc/json", crate::serde::lua_json(&lua)?)?;
        lua.register_module("@webc/toml", crate::serde::lua_toml(&lua)?)?;
        lua.register_module("@webc/wclm", Self::lua_wclm(&lua)?)?;
        lua.register_module("@webc/features", Self::lua_features(&lua, args)?)?;
        lua.register_module("@webc/async", crate::tokio::lua_async(&lua)?)?;
        lua.register_module("@webc/tools", Self::lua_tools(&lua)?)?;
        lua.register_module("@webc/string", Self::lua_string(&lua)?)?;
        lua.register_module("@webc/array", Self::lua_array(&lua)?)?;
        lua.register_module("@webc/memory", Self::lua_memory(&lua)?)?;
        lua.register_module("@webc/object", crate::object::lua_object(&lua)?)?;
        lua.register_module("@webc/process", crate::process::lua_process(&lua)?)?;
        #[cfg(feature = "web")]
        lua.register_module("@webc/tera", Self::lua_tera(&lua)?)?;
        #[cfg(feature = "web")]
        lua.register_module("@webc/http", crate::http::lua_http(&lua)?)?;
        #[cfg(feature = "globber")]
        lua.register_module("@webc/globber", crate::globber::lua_globber(&lua)?)?;
        #[cfg(feature = "db")]
        lua.register_module("@webc/sqlite", crate::db::sqlite::lua_sqlite(&lua)?)?;
        #[cfg(feature = "db")]
        lua.register_module("@webc/psql", crate::db::psql::lua_psql(&lua)?)?;
        #[cfg(feature = "db")]
        lua.register_module("@webc/redis", crate::db::redis::lua_redis(&lua)?)?;

        // add collections
        crate::collections::array::LuaArray::create(&lua);
        crate::collections::hashmap::LuaHashMap::create(&lua);

        // execute
        let temp = entry_path_path.canonicalize().unwrap();
        let entry_path_path_str = temp.to_str().unwrap();

        Ok(lua
            .load(script)
            .set_name(format!(
                "@{}",
                entry_path_path_str
                    .strip_suffix(".luau")
                    .unwrap_or(entry_path_path_str)
            ))
            .eval()?)
    }

    /// [`Self::lua_exec_with`], except it captures and deserializes the result.
    pub fn lua_exec_capture_with<T>(
        lua: &Lua,
        script: Vec<u8>,
        entry_path: String,
        args: String,
    ) -> LuaResult<T>
    where
        T: Serialize + DeserializeOwned,
    {
        lua.from_value(Self::lua_exec_with(lua, script, entry_path, args)?)
    }

    /// Execute the given lua script to and capture the returned `T`.
    pub fn lua_exec(script: Vec<u8>, entry_path: String, args: String) -> LuaResult<LuaValue> {
        Self::lua_exec_with(&Self::lua(), script, entry_path, args)
    }

    /// Execute the given lua script to and capture the returned `T`.
    pub fn lua_exec_with_instance(
        lua: &Lua,
        script: Vec<u8>,
        entry_path: String,
        args: String,
    ) -> LuaResult<LuaValue> {
        Self::lua_exec_with(&lua, script, entry_path, args)
    }

    /// Execute the given lua script to and capture the returned `T`.
    pub fn lua_exec_capture<T>(script: Vec<u8>, entry_path: String, args: String) -> LuaResult<T>
    where
        T: Serialize + DeserializeOwned,
    {
        Self::lua_exec_capture_with(&Self::lua(), script, entry_path, args)
    }

    /// Execute the given lua script to edit `T`.
    pub fn lua_exec_on_with<T>(
        lua: Lua,
        ctx: T,
        script: Vec<u8>,
        entry_path: String,
        args: String,
    ) -> LuaResult<T>
    where
        T: Serialize + DeserializeOwned,
    {
        lua.globals().set("ctx", lua.to_value(&ctx)?)?;
        Self::lua_exec_capture_with(&lua, script, entry_path, args)
    }

    /// Execute the given lua script to edit `T`.
    pub fn lua_exec_on<T>(ctx: T, script: Vec<u8>, entry_path: String, args: String) -> LuaResult<T>
    where
        T: Serialize + DeserializeOwned,
    {
        let lua = Self::lua();
        lua.globals().set("ctx", lua.to_value(&ctx)?)?;
        Self::lua_exec_capture_with(&lua, script, entry_path, args)
    }
}