retex-lua 0.9.0

//! Embedded Lua 5.3 and the LuaTeX Lua API.
//!
//! [`LuaTexHost`] implements the engine's [`LuaHost`] trait. During a
//! `\directlua` call the engine is lent to Lua as scoped userdata stored in
//! the registry under [`ENGINE_KEY`]. Library functions borrow it from
//! there, so no Lua object can outlive the borrow. Nested calls
//! (Lua → TeX → Lua) install a fresh scoped reference and restore the
//! previous one afterwards.
//!
//! The libraries follow the LuaTeX reference manual (chapter 10, "The TeX
//! related libraries") and are checked against the behaviour of the
//! reference `luatex` binary.

use std::rc::Rc;

use mlua::{AnyUserData, Lua, Value};
use retex_engine::input::LuaItem;
use retex_engine::{Engine, LuaHost};

mod callback;
mod compat;
mod extras;
mod fontlib;
mod fontloaderlib;
mod fonttable;
mod kpselib;
mod langlib;
mod lfslib;
mod lualib;
pub mod luaotfload;
pub mod nodelib;
mod otfaux;
mod otfnodes;
mod pdflib;
mod siolib;
mod statuslib;
mod texiolib;
mod texlib;
mod tokenlib;

/// Registry key of the scoped engine reference.
pub const ENGINE_KEY: &str = "retex.engine";

/// Text printed by `tex.print` and friends during the current calls, one
/// buffer per nesting level.
#[derive(Default)]
pub(crate) struct PrintStack(pub Vec<Vec<LuaItem>>);

/// Runs `f` with the engine currently lent to Lua.
pub(crate) fn with_engine<R>(lua: &Lua, f: impl FnOnce(&mut Engine) -> R) -> mlua::Result<R> {
    let ud: AnyUserData = lua
        .named_registry_value(ENGINE_KEY)
        .map_err(|e| mlua::Error::runtime(format!("no engine: {e}")))?;
    ud.borrow_mut_scoped::<Engine, R>(f)
        .map_err(|e| mlua::Error::runtime(format!("engine borrow: {e}")))
}

/// Registry key of the factory used by [`create_function_plain_errors`].
const PLAIN_ERRORS_KEY: &str = "retex.plainerrors";

/// Like `Lua::create_function`, but errors reach Lua code as plain
/// strings, as `luaL_error` raises them in LuaTeX's C libraries
/// (`pcall(pdf.obj)` returns `false, "pdf.obj() needs at least one
/// argument"`), instead of mlua's error userdata with a traceback.
pub(crate) fn create_function_plain_errors<A, R, F>(lua: &Lua, f: F) -> mlua::Result<mlua::Function>
where
    A: mlua::FromLuaMulti,
    R: mlua::IntoLuaMulti,
    F: Fn(&Lua, A) -> mlua::Result<R> + 'static,
{
    let inner = lua.create_function(move |lua, args: A| {
        let mut out = match f(lua, args) {
            Ok(r) => r.into_lua_multi(lua)?,
            Err(e) => {
                let msg = lua.create_string(lua_error_message(&e))?;
                return Ok(mlua::MultiValue::from_vec(vec![
                    Value::Boolean(false),
                    Value::String(msg),
                ]));
            }
        };
        out.push_front(Value::Boolean(true));
        Ok(out)
    })?;
    let factory: mlua::Function = match lua.named_registry_value::<Value>(PLAIN_ERRORS_KEY)? {
        Value::Function(f) => f,
        _ => {
            let f = lua
                .load(
                    "local f = ...
                     local function finish(ok, ...)
                       if ok then return ... end
                       error((...), 0)
                     end
                     return function(...) return finish(f(...)) end",
                )
                .set_name("=[retex]")
                .into_function()?;
            lua.set_named_registry_value(PLAIN_ERRORS_KEY, &f)?;
            f
        }
    };
    factory.call(inner)
}

/// Appends an item to the current print buffer.
pub(crate) fn push_print(lua: &Lua, item: LuaItem) {
    if let Some(mut ps) = lua.app_data_mut::<PrintStack>()
        && let Some(top) = ps.0.last_mut()
    {
        top.push(item);
    }
}

/// The Lua interpreter with LuaTeX's libraries.
pub struct LuaTexHost {
    lua: Lua,
}

/// A sampling profiler for performance work (`RETEX_LUA_PROFILE=1`):
/// every 1000 Lua VM instructions, the running function is charged one
/// sample; the busiest functions are printed to stderr at exit.
const PROFILER_START: &str = r#"
local samples, getinfo = {}, debug.getinfo
debug.sethook(function()
  local i = getinfo(2, "S")
  if i then
    local k = i.short_src .. ":" .. i.linedefined
    samples[k] = (samples[k] or 0) + 1
  end
end, "", 1000)
__retex_profile = samples
"#;

const PROFILER_REPORT: &str = r#"
local s = __retex_profile
if not s then return end
debug.sethook()
local t, total = {}, 0
for k, n in pairs(s) do t[#t + 1] = {k, n}; total = total + n end
table.sort(t, function(a, b) return a[2] > b[2] end)
io.stderr:write(string.format("[lua profile] %d samples (1 per 1000 instructions)\n", total))
for i = 1, math.min(25, #t) do
  io.stderr:write(string.format("[lua profile] %6d  %5.1f%%  %s\n", t[i][2], 100 * t[i][2] / total, t[i][1]))
end
"#;

impl Drop for LuaTexHost {
    fn drop(&mut self) {
        if std::env::var_os("RETEX_LUA_PROFILE").is_some() {
            let _ = self.lua.load(PROFILER_REPORT).exec();
        }
    }
}

impl LuaTexHost {
    /// Creates the interpreter and installs the libraries.
    pub fn new() -> mlua::Result<Rc<Self>> {
        // SAFETY: LuaTeX documents expect the full standard library,
        // including `debug`. This is the crate's single audited unsafe site:
        // mlua requires `unsafe` to load the debug library, which can break
        // Lua's own invariants but not Rust's memory safety outside Lua.
        #[allow(unsafe_code)]
        let lua = unsafe { Lua::unsafe_new_with(mlua::StdLib::ALL, mlua::LuaOptions::new()) };
        lua.set_app_data(PrintStack::default());
        retex_engine::timing::mark("create: Lua state");
        compat::install(&lua)?;
        texlib::install(&lua)?;
        texiolib::install(&lua)?;
        lualib::install(&lua)?;
        callback::install(&lua)?;
        kpselib::install(&lua)?;
        statuslib::install(&lua)?;
        lfslib::install(&lua)?;
        nodelib::install(&lua)?;
        tokenlib::install(&lua)?;
        extras::install(&lua)?;
        luaotfload::install(&lua)?;
        langlib::install(&lua)?;
        fontlib::install(&lua)?;
        fontloaderlib::install(&lua)?;
        pdflib::install(&lua)?;
        retex_engine::timing::mark("create: Lua libraries");
        retex_lpeg::install(&lua)?;
        retex_engine::timing::mark("create: LPeg");
        siolib::install(&lua)?;
        // The command line, as LuaTeX's `arg`: the program name as invoked
        // at index 0, then the arguments verbatim (hyperxmp reads it).
        let arg = lua.create_table()?;
        for (i, a) in std::env::args_os().enumerate() {
            arg.raw_set(i as i64, lua.create_string(a.as_encoded_bytes())?)?;
        }
        lua.globals().set("arg", arg)?;
        // The job's output directory, for the commands it runs.
        let output_directory = lua.create_function(|lua, ()| {
            Ok(crate::with_engine(lua, |e| {
                e.files
                    .output_dir
                    .as_ref()
                    .map(|d| d.to_string_lossy().into_owned())
            })
            .ok()
            .flatten())
        })?;
        lua.load(include_str!("shell.lua"))
            .set_name("=shell")
            .call::<()>(output_directory)?;
        if std::env::var_os("RETEX_LUA_PROFILE").is_some() {
            lua.load(PROFILER_START).set_name("=profiler").exec()?;
        }
        Ok(Rc::new(LuaTexHost { lua }))
    }

    /// The underlying Lua state.
    pub fn lua(&self) -> &Lua {
        &self.lua
    }

    /// Installs a host into an engine.
    pub fn attach(engine: &mut Engine) -> mlua::Result<Rc<Self>> {
        let host = LuaTexHost::new()?;
        engine.lua = Some(host.clone() as Rc<dyn LuaHost>);
        Ok(host)
    }
}

impl LuaTexHost {
    /// Lends the engine to Lua while `f` runs (see the module docs).
    fn lend<R>(
        &self,
        engine: &mut Engine,
        f: impl FnOnce(&Lua) -> mlua::Result<R>,
    ) -> mlua::Result<R> {
        let lua = &self.lua;
        lua.scope(|scope| {
            let ud = scope.create_any_userdata_ref_mut(&mut *engine)?;
            let prev: Value = lua.named_registry_value(ENGINE_KEY)?;
            lua.set_named_registry_value(ENGINE_KEY, ud)?;
            let r = f(lua);
            lua.set_named_registry_value(ENGINE_KEY, prev)?;
            r
        })
    }
}

impl LuaHost for LuaTexHost {
    fn run(&self, engine: &mut Engine, name: &str, code: &[u8]) -> Result<(), String> {
        let lua = &self.lua;
        if let Some(mut ps) = lua.app_data_mut::<PrintStack>() {
            ps.0.push(Vec::new());
        }
        let chunk_name = format!("=[{name}]");
        let result = lua.scope(|scope| {
            let ud = scope.create_any_userdata_ref_mut(&mut *engine)?;
            let prev: Value = lua.named_registry_value(ENGINE_KEY)?;
            lua.set_named_registry_value(ENGINE_KEY, ud)?;
            let r = lua.load(code).set_name(chunk_name.as_str()).exec();
            lua.set_named_registry_value(ENGINE_KEY, prev)?;
            r
        });
        let items = lua
            .app_data_mut::<PrintStack>()
            .and_then(|mut ps| ps.0.pop())
            .unwrap_or_default();
        engine.push_lua_output(items);
        result.map_err(|e| lua_error_message(&e))
    }

    fn call_simple(&self, engine: &mut Engine, name: &str) -> Result<(), String> {
        let Some(f) = callback::get(&self.lua, name) else {
            return Ok(());
        };
        self.lend(engine, |_lua| f.call::<()>(()))
            .map_err(|e| lua_error_message(&e))
    }

    fn has_callback(&self, name: &str) -> bool {
        callback::get(&self.lua, name).is_some()
    }

    fn define_font(&self, engine: &mut Engine, name: &[u8], size: i32) -> Option<u32> {
        let f = callback::get(&self.lua, "define_font")?;
        let next_id = engine.fonts.fonts.len() as i64;
        let name = self.lua.create_string(name).ok()?;
        let r = self.lend(engine, |_lua| {
            f.call::<Value>((name, i64::from(size), next_id))
        });
        match r {
            Ok(Value::Integer(n)) if n > 0 => Some(n as u32),
            Ok(Value::Number(n)) if n > 0.0 => Some(n as u32),
            Ok(Value::Table(t)) => {
                let info = fonttable::font_from_table(&t).ok()?;
                Some(engine.register_font(info))
            }
            Ok(_) => None,
            Err(e) => {
                engine.lua_error(&lua_error_message(&e));
                None
            }
        }
    }

    fn list_filter(
        &self,
        engine: &mut Engine,
        callback_name: &str,
        head: retex_engine::nodes::NodeRef,
        group: &str,
    ) -> retex_engine::ListFilter {
        use retex_engine::ListFilter;
        let Some(f) = callback::get(&self.lua, callback_name) else {
            return ListFilter::Keep(head);
        };
        let r = self.lend(engine, |lua| {
            let h = if head.is_null() {
                Value::Nil
            } else {
                Value::UserData(lua.create_userdata(nodelib::LNode(head.0))?)
            };
            f.call::<Value>((h, group))
        });
        match r {
            Ok(Value::Boolean(false)) => ListFilter::Flush,
            Ok(Value::UserData(u)) => match u.borrow::<nodelib::LNode>() {
                Ok(n) => ListFilter::Keep(retex_engine::nodes::NodeRef(n.0)),
                Err(_) => ListFilter::Keep(head),
            },
            Ok(Value::Integer(n)) => ListFilter::Keep(retex_engine::nodes::NodeRef(n as u32)),
            Ok(_) => ListFilter::Keep(head),
            Err(e) => {
                engine.lua_error(&lua_error_message(&e));
                ListFilter::Keep(head)
            }
        }
    }

    fn glyph_info(&self, engine: &mut Engine, p: retex_engine::nodes::NodeRef) -> Option<Vec<u8>> {
        let f = callback::get(&self.lua, "glyph_info")?;
        let r = self.lend(engine, |lua| {
            let n = lua.create_userdata(nodelib::LNode(p.0))?;
            f.call::<Value>(n)
        });
        match r {
            Ok(Value::String(s)) => Some(s.as_bytes().to_vec()),
            Ok(_) => None,
            Err(e) => {
                engine.lua_error(&lua_error_message(&e));
                None
            }
        }
    }

    fn dump(&self, _engine: &mut Engine) -> Vec<u8> {
        // Store the functions in `lua.bytecode` as (index, bytecode) pairs.
        let mut out = retex_engine::format::Writer::default();
        let mut entries: Vec<(i64, Vec<u8>)> = Vec::new();
        if let Ok(lt) = self.lua.globals().get::<mlua::Table>("lua")
            && let Ok(bc) = lt.get::<mlua::Table>("bytecode")
        {
            for (k, v) in bc.pairs::<i64, mlua::Function>().flatten() {
                entries.push((k, v.dump(false)));
            }
        }
        entries.sort_by_key(|(k, _)| *k);
        out.u(entries.len() as u64);
        for (k, b) in entries {
            out.i(k);
            out.bytes(&b);
        }
        out.buf
    }

    fn undump(&self, _engine: &mut Engine, data: &[u8]) -> Result<(), String> {
        if data.is_empty() {
            return Ok(());
        }
        let mut r = retex_engine::format::Reader::new(data);
        let n = r.u().map_err(|e| e.to_string())?;
        let lt: mlua::Table = self.lua.globals().get("lua").map_err(|e| e.to_string())?;
        let bc: mlua::Table = lt.get("bytecode").map_err(|e| e.to_string())?;
        for _ in 0..n {
            let k = r.i().map_err(|e| e.to_string())?;
            let b = r.bytes().map_err(|e| e.to_string())?;
            let f = self
                .lua
                .load(b)
                .set_mode(mlua::ChunkMode::Binary)
                .into_function()
                .map_err(|e| e.to_string())?;
            bc.set(k, f).map_err(|e| e.to_string())?;
        }
        Ok(())
    }
}

/// Formats a Lua error the way LuaTeX reports it: the message without
/// mlua's decorations.
pub fn lua_error_message(e: &mlua::Error) -> String {
    match e {
        mlua::Error::RuntimeError(s) | mlua::Error::SyntaxError { message: s, .. } => s.clone(),
        mlua::Error::CallbackError { cause, .. } => lua_error_message(cause),
        other => other.to_string(),
    }
}

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

    #[test]
    fn arg_holds_the_command_line() {
        let host = LuaTexHost::new().expect("host");
        let (n, first): (i64, String) = host.lua().load("return #arg, arg[0]").eval().expect("arg");
        assert_eq!(n as usize, std::env::args().count() - 1);
        assert_eq!(first, std::env::args().next().unwrap());
    }

    #[test]
    fn lpeg_is_installed() {
        let host = LuaTexHost::new().expect("host");
        let r: (i64, bool, String, bool) = host
            .lua()
            .load(
                r#"return lpeg.match(lpeg.R"az"^1, "abc1"), package.loaded.lpeg == lpeg,
                   lpeg.version(), package.loaded.re == nil"#,
            )
            .eval()
            .expect("lpeg");
        assert_eq!(r, (4, true, "1.0.1".to_string(), true));
    }
}