retex-lua 1.3.0-rc.3

//! expl3's integer arrays (l3intarray) run natively.
//!
//! Under LuaTeX, expl3 keeps its intarrays in Lua tables and accesses
//! them through `\luafunction`s defined in `expl3.lua`; l3regex keeps its
//! matcher state there, so mhchem's `\ce`, tabularray and other
//! regex-heavy documents make millions of these calls. The table lookup
//! they all start with:
//!
//! ```lua
//! local __intarray_table do
//!   local tables = get_luadata and get_luadata'__intarray' or {[0] = {}}
//!   function __intarray_table()
//!     local t = scan_token()
//!     if t ~= intarray_marker then
//!       put_next(t)
//!       tex.error'LaTeX Error: intarray expected'
//!       return tables[0]
//!     end
//!     local i = scan_int()
//!     local current_table = tables[i]
//!     if current_table then return current_table end
//!     current_table = {}
//!     tables[i] = current_table
//!     return current_table
//!   end
//!   ...
//! end
//! ```
//!
//! and the commands, for example:
//!
//! ```lua
//! luacmd('__intarray_gset:w', function()
//!   local i = scan_int()
//!   local t = __intarray_table()
//!   t[i] = scan_int()
//! end, 'protected', 'global')
//!
//! luacmd('__intarray_item:w', function()
//!   local i = scan_int()
//!   local t = __intarray_table()
//!   sprint(-2, t[i])
//! end, 'global')
//! ```
//!
//! Most of their time went to the Lua calls of the token library. The
//! native versions do the same steps with the engine's scanners and leave
//! the arrays where expl3 keeps them, in the `tables` upvalue of
//! `__intarray_table` (so the format keeps dumping them through
//! `register_luadata`): `scan_int` is [`Engine::scan_int`], `scan_token`
//! [`Engine::get_x_token`], `put_next` [`Engine::ins_list`], and
//! `sprint(-2, n)` of an integer the text of `n` printed with the
//! string catcode regime when the call ends. The rare paths (no intarray
//! after the index, an index out of range, entries that are not integers)
//! call the same Lua functions as the original code.
//!
//! A function is recognised when it comes from `expl3.lua` and its
//! bytecode (stripped of line numbers and names) is one of the known
//! versions, when `__intarray_table` is the known one too, and when the
//! library functions it captured (`scan_int`, `tex.sprint`, ...) are the
//! built-in ones; it is classified again when its `\luafunction` slot
//! gets another function. The captured values are read then: they are
//! locals of `expl3.lua` that it never assigns again (the tables they
//! hold, like `tables`, are read when the functions run). Anything else
//! runs in Lua. As for the native
//! macros, `RETEX_NO_NATIVE=1` disables them and `RETEX_NATIVE_DEBUG=1`
//! prints the fingerprints of the functions of `expl3.lua` as they are
//! first called.

use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

use mlua::{Function, Lua, Table, Value};
use retex_engine::Engine;
use retex_engine::input::{CatMode, LuaItem, LuaLineKind};
use retex_engine::token::Token;

use crate::LuaTexHost;
use crate::tokenlib::LToken;

/// What a recognised function does.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Op {
    /// `\__intarray_gset_count:Nw`: pads the array with zeros.
    GsetCount,
    /// `\intarray_count:N`: the length.
    Count,
    /// `\__intarray_gset:wF`: sets an existing entry, or the false branch.
    GsetF,
    /// `\__intarray_gset:w`: sets an entry.
    Gset,
    /// `\intarray_gzero:N`: sets all entries to zero.
    Gzero,
    /// `\__intarray_item:wF`: an existing entry, or the false branch.
    ItemF,
    /// `\__intarray_item:w`: an entry.
    Item,
    /// `\__intarray_to_clist:Nn`: all entries, with a separator.
    ToClist,
    /// `\__intarray_range_to_clist:w`: entries, separated by commas.
    RangeToClist,
    /// `\__intarray_gset_range:w`: sets entries from a comma list.
    GsetRange,
}

/// Fingerprints ([`fingerprint`]) of the functions of `expl3.lua` (expl3
/// 2026-01-19, TeX Live 2025).
const KNOWN: &[(u64, Op)] = &[
    (0x44f9_10d6_1e86_2917, Op::GsetCount),
    (0x4db4_34a5_7226_cc1c, Op::Count),
    (0x3d6d_f53c_5793_9e57, Op::GsetF),
    (0xc6b3_092c_6873_3a65, Op::Gset),
    (0x8e17_137f_277f_d428, Op::Gzero),
    (0x3a64_d664_79e5_3cad, Op::ItemF),
    (0xfe57_e74e_51d1_7ad2, Op::Item),
    (0xbdcd_adf9_896c_e4d6, Op::ToClist),
    (0xf990_d4fc_7891_767c, Op::RangeToClist),
    (0xcd0d_5316_f574_c051, Op::GsetRange),
];

/// Fingerprints of `__intarray_table`.
const KNOWN_TABLE: &[u64] = &[0xde6e_911a_e45c_eb5f];

/// A hash of a function's bytecode without its debug information (line
/// numbers, names), which moving the code in its file leaves alone.
fn fingerprint(f: &Function) -> u64 {
    let mut h = retex_engine::fingerprint::Fp::default();
    h.bytes(&f.dump(true));
    h.get()
}

/// A recognised function, with the values it captured.
struct Native {
    op: Op,
    /// expl3's arrays, by number (`tables` of `__intarray_table`).
    tables: Table,
    /// `\__intarray:w`, which starts an intarray variable.
    marker: Token,
    /// The `tex` table of `__intarray_table`.
    table_tex: Table,
    /// The function's `tex` table (`GsetF`, `ItemF`) or `token` table
    /// (`ToClist`).
    lib: Option<Table>,
    /// `\use_none:n` and `\use:n` (`GsetF`, `ItemF`), or `,` (`GsetRange`).
    toks: [Token; 2],
    /// `{\exp_after:wN, \scan_stop:}` (`ItemF`), read when called as the
    /// Lua code reads it.
    expand_after: Option<Table>,
}

/// The library functions the recognised functions may capture, as
/// installed (a document could replace them before expl3 captures them).
struct Builtins {
    scan_int: Function,
    scan_token: Function,
    put_next: Function,
    scan_string: Function,
    sprint: Function,
    concat: Function,
    getupvalue: Function,
}

/// The Lua code of the rare paths, run in Lua as in the original.
struct Helpers {
    not_intarray: Function,
    bad_index_count: Function,
    bad_index_param: Function,
}

/// They receive the `tex` table the original captured.
const HELPERS: &str = "return function(tex) tex.error'LaTeX Error: intarray expected' end,
  function(tex, i) tex.count.l__intarray_bad_index_int = i end,
  function(tex, i) tex.l__intarray_bad_index_int = i end";

/// A `\luafunction` slot's function, and what it is if recognised.
type Slot = (Function, Option<Rc<Native>>);

/// The native intarray functions of a Lua state.
pub(crate) struct Intarrays {
    builtins: Option<Builtins>,
    helpers: Option<Helpers>,
    /// By `\luafunction` slot: the function last seen there (kept alive,
    /// so that no other function gets its address) and what it is.
    slots: RefCell<HashMap<i32, Slot, retex_engine::hash::FastHash>>,
}

impl Intarrays {
    /// Records the library functions of a fresh Lua state.
    pub(crate) fn new(lua: &Lua) -> Self {
        let builtins = (|| -> mlua::Result<Builtins> {
            let g = lua.globals();
            let token: Table = g.get("token")?;
            let tex: Table = g.get("tex")?;
            let table: Table = g.get("table")?;
            let debug: Table = g.get("debug")?;
            Ok(Builtins {
                scan_int: token.get("scan_int")?,
                scan_token: token.get("scan_token")?,
                put_next: token.get("put_next")?,
                scan_string: token.get("scan_string")?,
                sprint: tex.get("sprint")?,
                concat: table.get("concat")?,
                getupvalue: debug.get("getupvalue")?,
            })
        })()
        .ok();
        let helpers = (|| -> mlua::Result<Helpers> {
            let (a, b, c) = lua.load(HELPERS).set_name("=expl3.lua").call(())?;
            Ok(Helpers {
                not_intarray: a,
                bad_index_count: b,
                bad_index_param: c,
            })
        })()
        .ok()
        .filter(|_| std::env::var_os("RETEX_NO_NATIVE").is_none());
        Intarrays {
            builtins,
            helpers,
            slots: RefCell::default(),
        }
    }

    /// Runs `f`, the function of slot `n`, natively if it is one of
    /// expl3's intarray functions; `None` otherwise.
    pub(crate) fn call(
        &self,
        host: &LuaTexHost,
        engine: &mut Engine,
        n: i32,
        f: &Function,
    ) -> Option<Result<(), String>> {
        let cached = self
            .slots
            .borrow()
            .get(&n)
            .filter(|(g, _)| g.to_pointer() == f.to_pointer())
            .map(|(_, native)| native.clone());
        let native = match cached {
            Some(native) => native?,
            None => {
                let native = self.classify(f).map(Rc::new);
                self.slots
                    .borrow_mut()
                    .insert(n, (f.clone(), native.clone()));
                native?
            }
        };
        let b = self.builtins.as_ref()?;
        let h = self.helpers.as_ref()?;
        // `token.scan_string` is looked up when the function runs.
        if let (Op::ToClist, Some(token)) = (native.op, &native.lib)
            && !same_fn(
                token.raw_get::<Value>("scan_string").ok().as_ref(),
                &b.scan_string,
            )
        {
            return None;
        }
        Some(run(host, engine, &native, b, h).map_err(|e| crate::lua_error_message(&e)))
    }

    /// Writes entries of the arrays, for native code doing what the TeX
    /// code would do with `\__intarray_gset:w` (see
    /// [`retex_engine::LuaHost::intarray_set`]): through the `tables` of
    /// `\__intarray_gset:w`'s function `f` (slot `n`), if recognised.
    pub(crate) fn set(
        &self,
        lua: &Lua,
        n: i32,
        f: Option<&Function>,
        writes: &[(i64, i64, &[i64])],
    ) -> bool {
        let Some(f) = f else {
            return false;
        };
        let cached = self
            .slots
            .borrow()
            .get(&n)
            .filter(|(g, _)| g.to_pointer() == f.to_pointer())
            .map(|(_, native)| native.clone());
        let native = match cached {
            Some(native) => native,
            None => {
                let native = self.classify(f).map(Rc::new);
                self.slots
                    .borrow_mut()
                    .insert(n, (f.clone(), native.clone()));
                native
            }
        };
        let Some(native) = native.filter(|n| n.op == Op::Gset) else {
            return false;
        };
        let tables = native.tables.clone();
        let mut targets = Vec::with_capacity(writes.len());
        for &(array, _, _) in writes {
            // As `__intarray_table` makes a missing array.
            let t = match tables.raw_get::<Value>(array) {
                Ok(Value::Table(t)) => t,
                Ok(Value::Nil) => {
                    let Ok(t) = lua.create_table().and_then(|t| {
                        tables.raw_set(array, &t)?;
                        Ok(t)
                    }) else {
                        return false;
                    };
                    t
                }
                _ => return false,
            };
            targets.push(t);
        }
        for (t, &(_, first, values)) in targets.iter().zip(writes) {
            for (k, &v) in values.iter().enumerate() {
                if t.raw_set(first + k as i64, v).is_err() {
                    return false;
                }
            }
        }
        true
    }

    /// Recognises one of the intarray functions.
    fn classify(&self, f: &Function) -> Option<Native> {
        let b = self.builtins.as_ref()?;
        let source = f.info().source?;
        if !source.ends_with("expl3.lua") {
            return None;
        }
        let fp = fingerprint(f);
        let op = KNOWN.iter().find(|(k, _)| *k == fp).map(|(_, op)| *op);
        if std::env::var_os("RETEX_NATIVE_DEBUG").is_some() {
            eprintln!(
                "[native] expl3.lua:{}: 0x{fp:016x} {}",
                f.info().line_defined.unwrap_or(0),
                op.map_or("no match".to_owned(), |op| format!("{op:?}"))
            );
        }
        let op = op?;
        let up = upvalues(b, f)?;
        let get = |name: &str| up.iter().find(|(n, _)| n == name).map(|(_, v)| v);
        let is = |name: &str, builtin: &Function| same_fn(get(name), builtin);
        let tok = |name: &str| match get(name) {
            Some(Value::UserData(u)) => u.borrow::<LToken>().ok().map(|t| Token(t.0)),
            _ => None,
        };
        let Some(Value::Function(table_fn)) = get("__intarray_table") else {
            return None;
        };
        if !KNOWN_TABLE.contains(&fingerprint(table_fn)) {
            return None;
        }
        let tup = upvalues(b, table_fn)?;
        let tget = |name: &str| tup.iter().find(|(n, _)| n == name).map(|(_, v)| v);
        let tis = |name: &str, builtin: &Function| same_fn(tget(name), builtin);
        if !(tis("scan_token", &b.scan_token)
            && tis("put_next", &b.put_next)
            && tis("scan_int", &b.scan_int))
        {
            return None;
        }
        let Some(Value::Table(table_tex)) = tget("tex") else {
            return None;
        };
        let Some(Value::Table(tables)) = tget("tables") else {
            return None;
        };
        let marker = match tget("intarray_marker") {
            Some(Value::UserData(u)) => Token(u.borrow::<LToken>().ok()?.0),
            _ => return None,
        };
        // The captured values each function uses besides the table.
        let scans = is("scan_int", &b.scan_int);
        let prints = is("sprint", &b.sprint);
        let mut toks = [Token(0); 2];
        let mut expand_after = None;
        let mut lib = None;
        let ok = match op {
            Op::GsetCount | Op::Gset => scans,
            Op::Count | Op::Item => prints && (op == Op::Count || scans),
            Op::Gzero => true,
            Op::ToClist => {
                lib = get("token").and_then(|v| v.as_table()).cloned();
                prints && is("concat", &b.concat) && lib.is_some()
            }
            Op::RangeToClist => scans && prints && is("concat", &b.concat),
            Op::GsetF | Op::ItemF => {
                let (Some(none), Some(one)) = (tok("use_none"), tok("use_i")) else {
                    return None;
                };
                toks = [none, one];
                if op == Op::ItemF {
                    let Some(Value::Table(t)) = get("expand_after_scan_stop") else {
                        return None;
                    };
                    expand_after = Some(t.clone());
                }
                scans
                    && is("put_next", &b.put_next)
                    && (op == Op::GsetF || (prints && is("scan_token", &b.scan_token)))
                    && {
                        lib = get("tex").and_then(|v| v.as_table()).cloned();
                        lib.is_some()
                    }
            }
            Op::GsetRange => {
                toks[0] = tok("comma")?;
                scans && is("scan_token", &b.scan_token) && is("put_next", &b.put_next)
            }
        };
        ok.then(|| Native {
            op,
            tables: tables.clone(),
            marker,
            table_tex: table_tex.clone(),
            lib,
            toks,
            expand_after,
        })
    }
}

/// Whether `v` is the function `f`.
fn same_fn(v: Option<&Value>, f: &Function) -> bool {
    matches!(v, Some(Value::Function(g)) if g == f)
}

/// The upvalues of `f` by name.
fn upvalues(b: &Builtins, f: &Function) -> Option<Vec<(String, Value)>> {
    let mut out = Vec::new();
    for i in 1.. {
        let (name, v): (Option<String>, Value) = b.getupvalue.call((f, i)).ok()?;
        let Some(name) = name else { break };
        out.push((name, v));
    }
    Some(out)
}

/// `token.scan_int()`.
fn scan_int(e: &mut Engine) -> i64 {
    e.scan_int();
    i64::from(e.cur_val)
}

/// `token.scan_token()`.
fn scan_token(e: &mut Engine) -> Token {
    e.get_x_token();
    e.cur_tok
}

/// `token.put_next(t)`.
fn put_next(e: &mut Engine, toks: &[Token]) {
    e.ins_list(Rc::from(toks));
}

/// Lua's truth.
fn truthy(v: &Value) -> bool {
    !matches!(v, Value::Nil | Value::Boolean(false))
}

/// `sprint(-2, n)` of an integer.
fn print_int(n: i64) -> LuaItem {
    LuaItem::Text {
        text: n.to_string().into_bytes(),
        kind: LuaLineKind::Sprint,
        catmode: CatMode::String,
    }
}

/// `__intarray_table()`.
fn table(host: &LuaTexHost, e: &mut Engine, nat: &Native, h: &Helpers) -> mlua::Result<Table> {
    let t = scan_token(e);
    if t != nat.marker {
        put_next(e, &[t]);
        host.lend(e, |_| h.not_intarray.call::<()>(&nat.table_tex))?;
        return nat.tables.raw_get(0);
    }
    let i = scan_int(e);
    match nat.tables.raw_get::<Value>(i)? {
        Value::Table(t) => Ok(t),
        v if truthy(&v) => Err(mlua::Error::runtime(format!(
            "attempt to index a {} value",
            v.type_name()
        ))),
        _ => {
            let t = host.lua.create_table()?;
            nat.tables.raw_set(i, &t)?;
            Ok(t)
        }
    }
}

/// `sprint(-2, v)`, when the call ends: an integer natively, anything
/// else by `tex.sprint`.
fn sprint(host: &LuaTexHost, e: &mut Engine, b: &Builtins, v: Value) -> mlua::Result<()> {
    match v {
        Value::Integer(n) => {
            e.push_lua_output(vec![print_int(n)]);
            Ok(())
        }
        v => host
            .run_chunk(e, |_| b.sprint.call::<()>((-2, v)))
            .map_err(mlua::Error::runtime),
    }
}

/// `concat(t, sep, from, to)`: natively when the entries are integers.
fn concat(
    host: &LuaTexHost,
    e: &mut Engine,
    b: &Builtins,
    t: &Table,
    sep: &[u8],
    from: i64,
    to: i64,
) -> mlua::Result<()> {
    let mut text = Vec::new();
    let mut i = from;
    while i <= to {
        let Value::Integer(n) = t.raw_get::<Value>(i)? else {
            let sep = host.lua.create_string(sep)?;
            return host
                .run_chunk(e, |_| {
                    let s: Value = b.concat.call((t, sep, from, to))?;
                    b.sprint.call::<()>((-2, s))
                })
                .map_err(mlua::Error::runtime);
        };
        if i > from {
            text.extend_from_slice(sep);
        }
        text.extend_from_slice(n.to_string().as_bytes());
        i += 1;
    }
    e.push_lua_output(vec![LuaItem::Text {
        text,
        kind: LuaLineKind::Sprint,
        catmode: CatMode::String,
    }]);
    Ok(())
}

/// Runs a recognised function.
fn run(
    host: &LuaTexHost,
    e: &mut Engine,
    nat: &Native,
    b: &Builtins,
    h: &Helpers,
) -> mlua::Result<()> {
    match nat.op {
        Op::GsetCount => {
            let t = table(host, e, nat, h)?;
            let n = scan_int(e);
            for i in t.raw_len() as i64 + 1..=n {
                t.raw_set(i, 0)?;
            }
        }
        Op::Count => {
            let t = table(host, e, nat, h)?;
            e.push_lua_output(vec![print_int(t.raw_len() as i64)]);
        }
        Op::GsetF => {
            let i = scan_int(e);
            let t = table(host, e, nat, h)?;
            if truthy(&t.raw_get::<Value>(i)?) {
                t.raw_set(i, scan_int(e))?;
                put_next(e, &nat.toks[..1]);
            } else {
                host.lend(e, |_| h.bad_index_count.call::<()>((nat.lib.clone(), i)))?;
                scan_int(e);
                put_next(e, &nat.toks[1..]);
            }
        }
        Op::Gset => {
            let i = scan_int(e);
            let t = table(host, e, nat, h)?;
            t.raw_set(i, scan_int(e))?;
        }
        Op::Gzero => {
            let t = table(host, e, nat, h)?;
            for i in 1..=t.raw_len() as i64 {
                t.raw_set(i, 0)?;
            }
        }
        Op::ItemF => {
            let i = scan_int(e);
            let t = table(host, e, nat, h)?;
            let item = t.raw_get::<Value>(i)?;
            let found = truthy(&item);
            if found {
                put_next(e, &nat.toks[..1]);
            } else {
                host.lend(e, |_| h.bad_index_param.call::<()>((nat.lib.clone(), i)))?;
                put_next(e, &nat.toks[1..]);
            }
            // `put_next` of a table: its tokens up to the first that is
            // not one.
            let mut toks = Vec::with_capacity(2);
            if let Some(ea) = &nat.expand_after {
                for v in ea.clone().sequence_values::<Value>() {
                    match v? {
                        Value::UserData(u) => match u.borrow::<LToken>() {
                            Ok(t) => toks.push(Token(t.0)),
                            Err(_) => return Err(put_next_fatal(e)),
                        },
                        _ => return Err(put_next_fatal(e)),
                    }
                }
            }
            put_next(e, &toks);
            scan_token(e);
            if found {
                sprint(host, e, b, item)?;
            }
        }
        Op::Item => {
            let i = scan_int(e);
            let t = table(host, e, nat, h)?;
            sprint(host, e, b, t.raw_get::<Value>(i)?)?;
        }
        Op::ToClist => {
            let t = table(host, e, nat, h)?;
            let sep = crate::tokenlib::scan_string(e);
            concat(host, e, b, &t, &sep, 1, t.raw_len() as i64)?;
        }
        Op::RangeToClist => {
            let t = table(host, e, nat, h)?;
            let from = scan_int(e);
            let to = scan_int(e);
            concat(host, e, b, &t, b",", from, to)?;
        }
        Op::GsetRange => {
            let mut from = scan_int(e);
            let t = table(host, e, nat, h)?;
            let comma = nat.toks[0];
            loop {
                let tok = scan_token(e);
                if tok == comma {
                    while scan_token(e) == comma {}
                    break;
                }
                put_next(e, &[tok]);
                t.raw_set(from, scan_int(e))?;
                scan_token(e);
                from += 1;
            }
        }
    }
    Ok(())
}

/// `put_next` of a table holding something else than tokens: LuaTeX's
/// fatal error (see the token library).
fn put_next_fatal(e: &mut Engine) -> mlua::Error {
    e.luatex_fatal("token lib", "lua <token> expected in put_next (1)")
}