retex-lua 0.9.0

//! LuaTeX's extensions to the standard Lua libraries (LuaTeX manual §4.2):
//! `md5`, extra `string` and `os` functions, and `unicode.utf8`.

use mlua::{Lua, MultiValue, Table, Value, Variadic};

/// Installs the extensions.
pub fn install(lua: &Lua) -> mlua::Result<()> {
    // md5
    let md5 = lua.create_table()?;
    md5.set(
        "sum",
        lua.create_function(|lua, s: mlua::String| {
            lua.create_string(retex_pdf::md5::md5(&s.as_bytes()))
        })?,
    )?;
    md5.set(
        "sumhexa",
        lua.create_function(|_, s: mlua::String| {
            Ok(hex(&retex_pdf::md5::md5(&s.as_bytes()), false))
        })?,
    )?;
    md5.set(
        "sumHEXA",
        lua.create_function(|_, s: mlua::String| {
            Ok(hex(&retex_pdf::md5::md5(&s.as_bytes()), true))
        })?,
    )?;
    lua.globals().set("md5", md5)?;

    // string extras
    let string: Table = lua.globals().get("string")?;
    string.set(
        "utfvalues",
        lua.create_function(|lua, s: mlua::String| {
            let chars = retex_engine::utf8x::chars(&s.as_bytes());
            let state = lua.create_table()?;
            for (i, c) in chars.iter().enumerate() {
                state.set(i + 1, *c)?;
            }
            state.set("i", 0)?;
            let iter = lua.create_function(|_, st: Table| {
                let i: i64 = st.get("i")?;
                st.set("i", i + 1)?;
                st.get::<Value>(i + 1)
            })?;
            Ok((iter, state))
        })?,
    )?;
    string.set(
        "utfcharacters",
        lua.create_function(|lua, s: mlua::String| {
            let bytes = s.as_bytes().to_vec();
            let state = lua.create_table()?;
            let mut i = 0;
            let mut k = 1;
            while i < bytes.len() {
                let (_, n) = retex_engine::utf8x::decode(&bytes, i);
                state.set(k, lua.create_string(&bytes[i..i + n])?)?;
                i += n;
                k += 1;
            }
            state.set("i", 0)?;
            let iter = lua.create_function(|_, st: Table| {
                let i: i64 = st.get("i")?;
                st.set("i", i + 1)?;
                st.get::<Value>(i + 1)
            })?;
            Ok((iter, state))
        })?,
    )?;
    string.set(
        "characters",
        lua.create_function(|lua, s: mlua::String| {
            let bytes = s.as_bytes().to_vec();
            let state = lua.create_table()?;
            for (i, b) in bytes.iter().enumerate() {
                state.set(i + 1, lua.create_string([*b])?)?;
            }
            state.set("i", 0)?;
            let iter = lua.create_function(|_, st: Table| {
                let i: i64 = st.get("i")?;
                st.set("i", i + 1)?;
                st.get::<Value>(i + 1)
            })?;
            Ok((iter, state))
        })?,
    )?;
    string.set(
        "explode",
        lua.create_function(|lua, (s, sep): (mlua::String, Option<mlua::String>)| {
            let s = s.as_bytes().to_vec();
            let out = lua.create_table()?;
            let (seps, plus) = match &sep {
                Some(p) => {
                    let p = p.as_bytes().to_vec();
                    let plus = p.ends_with(b"+");
                    let core = if plus { p[..p.len() - 1].to_vec() } else { p };
                    (core, plus)
                }
                None => (b" ".to_vec(), true),
            };
            let mut parts: Vec<&[u8]> = s.split(|b| seps.contains(b)).collect();
            if plus {
                parts.retain(|p| !p.is_empty());
            }
            for (i, p) in parts.iter().enumerate() {
                out.set(i + 1, lua.create_string(p)?)?;
            }
            Ok(out)
        })?,
    )?;
    string.set(
        "bytes",
        lua.create_function(|lua, s: mlua::String| {
            let bytes = s.as_bytes().to_vec();
            let state = lua.create_table()?;
            for (i, b) in bytes.iter().enumerate() {
                state.set(i + 1, *b)?;
            }
            state.set("i", 0)?;
            let iter = lua.create_function(|_, st: Table| {
                let i: i64 = st.get("i")?;
                st.set("i", i + 1)?;
                st.get::<Value>(i + 1)
            })?;
            Ok((iter, state))
        })?,
    )?;
    string.set(
        "utflength",
        lua.create_function(|_, s: mlua::String| {
            Ok(retex_engine::utf8x::chars(&s.as_bytes()).len())
        })?,
    )?;
    // The pair iterators, `utfvalue` and `utfcharacter` (manual §4.2.5);
    // `bytetable` is undocumented (observed: the bytes as an array).
    lua.load(
        r#"
        local byte, sub, char, upairs = string.byte, string.sub, utf8.char, utf8.codes
        function string.bytepairs(s)
          local i = -1
          return function()
            i = i + 2
            if i <= #s then return byte(s, i), byte(s, i + 1) end
          end
        end
        function string.characterpairs(s)
          local i = -1
          return function()
            i = i + 2
            if i <= #s then return sub(s, i, i), sub(s, i + 1, i + 1) end
          end
        end
        function string.bytetable(s) return { byte(s, 1, -1) } end
        function string.utfvalue(s)
          local t = {}
          local ok = pcall(function()
            for _, c in upairs(s) do t[#t + 1] = c end
          end)
          if not ok then return end
          return table.unpack(t)
        end
        function string.utfcharacter(...) return char(...) end
        "#,
    )
    .set_name("=[string extras]")
    .exec()?;

    // os extras
    let os: Table = lua.globals().get("os")?;
    os.set("name", "linux")?;
    os.set("type", "unix")?;
    let selfdir = std::env::current_exe()
        .ok()
        .and_then(|p| p.parent().map(|d| d.to_string_lossy().into_owned()))
        .unwrap_or_default();
    os.set("selfdir", selfdir)?;
    os.set(
        "gettimeofday",
        lua.create_function(|_, ()| {
            Ok(std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map(|d| d.as_secs_f64())
                .unwrap_or(0.0))
        })?,
    )?;
    os.set(
        "uname",
        lua.create_function(|lua, ()| {
            let t = lua.create_table()?;
            t.set("sysname", "Linux")?;
            t.set("machine", std::env::consts::ARCH)?;
            t.set("nodename", "localhost")?;
            t.set("release", "")?;
            t.set("version", "")?;
            Ok(t)
        })?,
    )?;
    os.set(
        "setenv",
        lua.create_function(|_, (_k, _v): (String, Option<String>)| Ok(true))?,
    )?;
    os.set(
        "tmpdir",
        lua.create_function(|_, ()| {
            let d = std::env::temp_dir().join(format!("luatex.{}", std::process::id()));
            Ok(std::fs::create_dir_all(&d)
                .ok()
                .map(|_| d.to_string_lossy().into_owned()))
        })?,
    )?;
    os.set(
        "exec",
        lua.create_function(|_, _a: MultiValue| Ok((Value::Nil, "restricted")))?,
    )?;
    os.set(
        "spawn",
        lua.create_function(|_, _a: MultiValue| Ok((Value::Nil, "restricted")))?,
    )?;

    // unicode.utf8 (a subset of slnunicode)
    let utf8t = lua.create_table()?;
    utf8t.set(
        "len",
        lua.create_function(|_, s: mlua::String| {
            Ok(retex_engine::utf8x::chars(&s.as_bytes()).len())
        })?,
    )?;
    utf8t.set(
        "char",
        lua.create_function(|lua, cs: Variadic<i64>| {
            let mut out = Vec::new();
            for c in cs.iter() {
                retex_engine::utf8x::push(&mut out, *c as u32);
            }
            lua.create_string(&out)
        })?,
    )?;
    utf8t.set(
        "byte",
        lua.create_function(|_, (s, i, j): (mlua::String, Option<i64>, Option<i64>)| {
            let chars = retex_engine::utf8x::chars(&s.as_bytes());
            let n = chars.len() as i64;
            let norm = |k: i64| if k < 0 { n + k + 1 } else { k };
            let i = norm(i.unwrap_or(1)).max(1);
            let j = norm(j.unwrap_or(i)).min(n);
            Ok(Variadic::from_iter(
                (i..=j).map(|k| i64::from(chars[(k - 1) as usize])),
            ))
        })?,
    )?;
    utf8t.set(
        "sub",
        lua.create_function(|lua, (s, i, j): (mlua::String, i64, Option<i64>)| {
            let chars = retex_engine::utf8x::chars(&s.as_bytes());
            let n = chars.len() as i64;
            let norm = |k: i64| if k < 0 { n + k + 1 } else { k };
            let i = norm(i).max(1);
            let j = norm(j.unwrap_or(-1)).min(n);
            let mut out = Vec::new();
            if i <= j {
                for &c in &chars[(i - 1) as usize..j as usize] {
                    retex_engine::utf8x::push(&mut out, c);
                }
            }
            lua.create_string(&out)
        })?,
    )?;
    for (name, upper) in [("upper", true), ("lower", false)] {
        utf8t.set(
            name,
            lua.create_function(move |lua, s: mlua::String| {
                let st = String::from_utf8_lossy(&s.as_bytes()).into_owned();
                lua.create_string(if upper {
                    st.to_uppercase()
                } else {
                    st.to_lowercase()
                })
            })?,
        )?;
    }
    let unicode = lua.create_table()?;
    unicode.set("utf8", utf8t)?;
    lua.globals().set("unicode", unicode)?;
    Ok(())
}

fn hex(b: &[u8], upper: bool) -> String {
    b.iter()
        .map(|x| {
            if upper {
                format!("{x:02X}")
            } else {
                format!("{x:02x}")
            }
        })
        .collect()
}

#[cfg(test)]
mod tests {
    fn eval(code: &str) -> String {
        let host = crate::LuaTexHost::new().unwrap();
        host.lua().load(code).eval::<String>().unwrap()
    }

    // Expected values observed with texlua (LuaTeX 1.22).
    #[test]
    fn pair_iterators_and_utf_values() {
        let r = eval(
            r##"local out = {}
            for a, b in string.bytepairs("abc") do out[#out + 1] = a .. ":" .. tostring(b) end
            for a, b in string.characterpairs("abc") do out[#out + 1] = "[" .. a .. b .. "]" end
            local t = string.bytetable("abc")
            out[#out + 1] = #t .. "," .. t[3]
            out[#out + 1] = table.concat({string.utfvalue("aé€")}, ",")
            out[#out + 1] = string.utfcharacter(97, 233, 8364)
            out[#out + 1] = select("#", string.utfvalue("\255"))
            return table.concat(out, " ")"##,
        );
        assert_eq!(r, "97:98 99:nil [ab] [c] 3,99 97,233,8364 aé€ 0");
    }
}