retex-lua 1.0.0

//! The `token` library (LuaTeX manual §10.4).

use std::rc::Rc;

use mlua::{Lua, MetaMethod, MultiValue, Table, UserData, UserDataMethods, Value};
use retex_engine::Engine;
use retex_engine::cmd::Cmd;
use retex_engine::eqtb::{Loc, Meaning, Val};
use retex_engine::hash::{CsId, CsKind, NO_CS};
use retex_engine::token::{Macro, Token};

use crate::with_engine;

/// A token userdata.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct LToken(pub u32);

/// A character code from Lua: a number that is no character would make a
/// token in the range of control sequences (LuaTeX makes some token or
/// other); it is character 0 here.
fn valid_char(c: i64) -> u32 {
    u32::try_from(c)
        .ok()
        .filter(|&c| c <= 0x10_FFFF)
        .unwrap_or(0)
}

impl UserData for LToken {
    fn add_methods<M: UserDataMethods<Self>>(m: &mut M) {
        m.add_meta_method(MetaMethod::Index, |lua, this, key: String| {
            token_field(lua, Token(this.0), &key)
        });
        m.add_meta_method(MetaMethod::Eq, |_, this, other: mlua::AnyUserData| {
            Ok(other
                .borrow::<LToken>()
                .map(|o| o.0 == this.0)
                .unwrap_or(false))
        });
        m.add_meta_method(MetaMethod::ToString, |_, this, ()| {
            Ok(format!("<token {}>", this.0))
        });
    }
}

/// The command and modifier a token currently stands for.
fn token_meaning(e: &Engine, t: Token) -> (Cmd, u32, Option<Rc<Macro>>) {
    if t.is_cs() {
        let m = e.eqtb.meaning(t.cs_id());
        (m.cmd, m.chr, m.mac.clone())
    } else {
        (t.cmd(), t.chr(), None)
    }
}

fn token_field(lua: &Lua, t: Token, key: &str) -> mlua::Result<Value> {
    enum R {
        I(i64),
        S(Vec<u8>),
        B(bool),
        Nil,
    }
    let r = with_engine(lua, |e| {
        let (cmd, chr, mac) = token_meaning(e, t);
        match key {
            "command" => R::I(i64::from(cmd as u8)),
            "cmdname" => R::S(cmd.name().as_bytes().to_vec()),
            "mode" => R::I(i64::from(chr)),
            // Observed in LuaTeX: the register number for registers made
            // with \countdef and friends, nil for the internal parameters
            // of the same commands, for the code tables and for modifiers
            // beyond the character range (\relax, \par); otherwise the
            // modifier.
            "index" => match cmd {
                Cmd::AssignInt
                | Cmd::AssignAttr
                | Cmd::AssignDimen
                | Cmd::AssignGlue
                | Cmd::AssignMuGlue
                | Cmd::AssignToks => {
                    if chr >= retex_engine::scan::REG_OFFSET {
                        R::I(i64::from(chr - retex_engine::scan::REG_OFFSET))
                    } else {
                        R::Nil
                    }
                }
                Cmd::DefCharCode | Cmd::DefDelCode | Cmd::DefFamily => R::Nil,
                _ if chr > 0x10_FFFF => R::Nil,
                _ => R::I(i64::from(chr)),
            },
            "tok" => R::I(i64::from(t.0)),
            "id" => R::I(i64::from(t.0)),
            "csname" => {
                if t.is_cs() {
                    // The undefined control sequence has an empty name; an
                    // active character's name is the character.
                    if t.cs_id() == retex_engine::hash::UNDEFINED_CS {
                        R::S(Vec::new())
                    } else {
                        R::S(e.hash.name(t.cs_id()).to_vec())
                    }
                } else {
                    R::Nil
                }
            }
            "active" => R::B(t.is_cs() && matches!(e.hash.kind(t.cs_id()), CsKind::Active(_))),
            "expandable" => R::B(cmd.is_expandable()),
            "protected" => R::B(mac.is_some_and(|m| m.protected)),
            _ => R::Nil,
        }
    })?;
    Ok(match r {
        R::I(i) => Value::Integer(i),
        R::S(s) => Value::String(lua.create_string(&s)?),
        R::B(b) => Value::Boolean(b),
        R::Nil => Value::Nil,
    })
}

fn to_token(v: &Value) -> Option<Token> {
    match v {
        Value::UserData(u) => u.borrow::<LToken>().ok().map(|t| Token(t.0)),
        _ => None,
    }
}

fn wrap(lua: &Lua, t: Token) -> mlua::Result<Value> {
    Ok(Value::UserData(lua.create_userdata(LToken(t.0))?))
}

/// Looks up (creating if needed) a control sequence by name.
fn cs_by_name(e: &mut Engine, name: &[u8]) -> CsId {
    e.hash.lookup_or_insert(name)
}

/// Installs `token`.
pub fn install(lua: &Lua) -> mlua::Result<()> {
    let t = lua.create_table()?;
    t.set(
        "create",
        lua.create_function(|lua, (v, cat): (Value, Option<i64>)| {
            // Observed: a float is a character code when it is whole, and
            // 0 otherwise.
            let v = match v {
                Value::Number(x) => Value::Integer(if x.fract() == 0.0 { x as i64 } else { 0 }),
                v => v,
            };
            let tok = with_engine(lua, |e| match &v {
                // An unknown name gives the undefined control sequence
                // (LuaTeX does not enter the name in the hash).
                Value::String(s) => Some(Token::cs(
                    e.hash
                        .lookup(&s.as_bytes())
                        .unwrap_or(retex_engine::hash::UNDEFINED_CS),
                )),
                // token.create(chr [, cmd]): without cmd the current
                // catcode is used; 13 gives the active character, 0, 9, 14
                // and 15 an other character, and any other command code a
                // token with that command (e.g. char_given), as in LuaTeX.
                Value::Integer(c) => {
                    let c = valid_char(*c);
                    let cmd = cat.unwrap_or_else(|| i64::from(e.eqtb.catcode(c)));
                    Some(match cmd {
                        13 => Token::cs(
                            e.hash
                                .lookup_active(c)
                                .unwrap_or(retex_engine::hash::UNDEFINED_CS),
                        ),
                        0 | 9 | 14 | 15 => {
                            e.luatex_warning(
                                "token lib",
                                &format!(
                                    "not a good token, catcode {cmd} can not be returned, \
                                     so 12 will be used"
                                ),
                            );
                            Token::char(Cmd::OtherChar, c)
                        }
                        _ => match usize::try_from(cmd).ok().and_then(|i| Cmd::ALL.get(i)) {
                            Some(&cmd) => Token::char(cmd, c),
                            None => Token::char(Cmd::OtherChar, c),
                        },
                    })
                }
                _ => None,
            })?;
            match tok {
                Some(t) => wrap(lua, t),
                None => Ok(Value::Nil),
            }
        })?,
    )?;
    t.set(
        "new",
        lua.create_function(|lua, (chr, cmd): (i64, i64)| {
            // Out of range, a command is other_char and a character 0
            // (LuaTeX makes some token or other).
            let cmd = usize::try_from(cmd)
                .ok()
                .and_then(|i| Cmd::ALL.get(i))
                .copied()
                .unwrap_or(Cmd::OtherChar);
            wrap(lua, Token::char(cmd, valid_char(chr)))
        })?,
    )?;
    t.set(
        "is_token",
        lua.create_function(|_, v: Value| Ok(to_token(&v).is_some()))?,
    )?;
    t.set(
        "is_defined",
        lua.create_function(|lua, (name, _x): (mlua::String, Option<bool>)| {
            let n = name.as_bytes().to_vec();
            with_engine(lua, |e| {
                e.hash
                    .lookup(&n)
                    .is_some_and(|cs| e.eqtb.meaning(cs).cmd != Cmd::UndefinedCs)
            })
        })?,
    )?;
    // token.get_<field>(t) is t.<field> (LuaTeX manual §10.4; expl3's
    // \token_if_primitive:N uses get_command, get_index and get_mode).
    for field in [
        "active",
        "cmdname",
        "command",
        "csname",
        "expandable",
        "id",
        "index",
        "mode",
        "protected",
        "tok",
    ] {
        t.set(
            format!("get_{field}"),
            lua.create_function(move |lua, v: Value| match to_token(&v) {
                Some(tok) => token_field(lua, tok, field),
                None => Ok(Value::Nil),
            })?,
        )?;
    }
    t.set(
        "type",
        lua.create_function(|_, v: Value| Ok(to_token(&v).map(|_| "token")))?,
    )?;
    t.set("biggest_char", lua.create_function(|_, ()| Ok(0x10_FFFF))?)?;
    t.set(
        "commands",
        lua.create_function(|lua, ()| {
            let out = lua.create_table()?;
            for (i, c) in Cmd::ALL.iter().enumerate() {
                out.set(i, c.name())?;
            }
            Ok(out)
        })?,
    )?;
    t.set(
        "command_id",
        lua.create_function(|_, name: String| {
            Ok(Cmd::ALL
                .iter()
                .position(|c| c.name() == name)
                .map(|i| i as i64))
        })?,
    )?;
    // Scanners.
    t.set(
        "scan_int",
        lua.create_function(|lua, ()| {
            with_engine(lua, |e| {
                e.scan_int();
                i64::from(e.cur_val)
            })
        })?,
    )?;
    t.set(
        "scan_list",
        lua.create_function(|lua, ()| {
            let n = with_engine(lua, |e| e.scan_list())?;
            if n.is_null() {
                Ok(Value::Nil)
            } else {
                Ok(Value::UserData(
                    lua.create_userdata(crate::nodelib::LNode(n.0))?,
                ))
            }
        })?,
    )?;
    t.set(
        "scan_dimen",
        lua.create_function(|lua, (inf, mu): (Option<bool>, Option<bool>)| {
            with_engine(lua, |e| {
                let order = e.scan_dimen(mu.unwrap_or(false), inf.unwrap_or(false), false);
                (i64::from(e.cur_val), i64::from(order))
            })
        })?,
    )?;
    t.set(
        "scan_glue",
        lua.create_function(|lua, mu: Option<bool>| {
            let g = with_engine(lua, |e| {
                e.scan_glue(if mu.unwrap_or(false) {
                    retex_engine::scan::ValLevel::Mu
                } else {
                    retex_engine::scan::ValLevel::Glue
                });
                e.cur_glue
            })?;
            let t = lua.create_table()?;
            t.set("width", g.width)?;
            t.set("stretch", g.stretch)?;
            t.set("shrink", g.shrink)?;
            t.set("stretch_order", g.stretch_order)?;
            t.set("shrink_order", g.shrink_order)?;
            Ok(t)
        })?,
    )?;
    t.set(
        "scan_keyword",
        lua.create_function(|lua, k: String| with_engine(lua, |e| e.scan_keyword(&k)))?,
    )?;
    t.set(
        "scan_keyword_cs",
        lua.create_function(|lua, k: String| with_engine(lua, |e| e.scan_keyword(&k)))?,
    )?;
    t.set(
        "scan_toks",
        lua.create_function(|lua, (mac, expand): (Option<bool>, Option<bool>)| {
            let toks = with_engine(lua, |e| {
                if mac.unwrap_or(false) {
                    e.scan_toks(true, expand.unwrap_or(false))
                } else if expand.unwrap_or(false) {
                    e.scan_toks_expanded()
                } else {
                    e.scan_general_text().to_vec()
                }
            })?;
            let out = lua.create_table()?;
            for (i, tk) in toks.into_iter().enumerate() {
                out.set(i + 1, LToken(tk.0))?;
            }
            Ok(out)
        })?,
    )?;
    t.set(
        "scan_string",
        lua.create_function(|lua, ()| {
            let s = with_engine(lua, scan_string)?;
            lua.create_string(&s)
        })?,
    )?;
    t.set(
        "scan_argument",
        lua.create_function(|lua, expand: Option<bool>| {
            let s = with_engine(lua, |e| scan_argument(e, expand.unwrap_or(true)))?;
            lua.create_string(&s)
        })?,
    )?;
    t.set(
        "scan_word",
        lua.create_function(|lua, ()| {
            let s = with_engine(lua, |e| {
                let mut out = Vec::new();
                e.get_nonblank_noncall();
                loop {
                    if e.cur_cs == NO_CS
                        && (e.cur_cmd == Cmd::Letter || e.cur_cmd == Cmd::OtherChar)
                    {
                        retex_engine::utf8x::push(&mut out, e.cur_chr);
                        e.get_x_token();
                    } else {
                        if e.cur_cmd != Cmd::Spacer {
                            e.back_input();
                        }
                        break;
                    }
                }
                out
            })?;
            if s.is_empty() {
                Ok(Value::Nil)
            } else {
                Ok(Value::String(lua.create_string(&s)?))
            }
        })?,
    )?;
    t.set(
        "scan_csname",
        lua.create_function(|lua, ()| {
            let s = with_engine(lua, |e| {
                e.get_token();
                if e.cur_cs != NO_CS {
                    Some(e.hash.name(e.cur_cs).to_vec())
                } else {
                    None
                }
            })?;
            match s {
                Some(s) => Ok(Value::String(lua.create_string(&s)?)),
                None => Ok(Value::Nil),
            }
        })?,
    )?;
    t.set(
        "scan_token",
        lua.create_function(|lua, ()| {
            let tk = with_engine(lua, |e| {
                e.get_x_token();
                e.cur_tok
            })?;
            wrap(lua, tk)
        })?,
    )?;
    t.set(
        "get_next",
        lua.create_function(|lua, ()| {
            let tk = with_engine(lua, |e| {
                e.get_token();
                e.cur_tok
            })?;
            wrap(lua, tk)
        })?,
    )?;
    t.set(
        "put_next",
        lua.create_function(|lua, args: MultiValue| {
            let mut toks: Vec<Token> = Vec::new();
            // Observed: anything but tokens (or tables of tokens) is
            // LuaTeX's fatal error, with "(1)" for a table element.
            let fatal = |lua: &Lua, n: u8| -> mlua::Error {
                let msg = format!("lua <token> expected in put_next ({n})");
                match with_engine(lua, |e| e.luatex_fatal("token lib", &msg)) {
                    Err(e) => e,
                }
            };
            for v in args {
                match &v {
                    Value::Table(t) => {
                        for x in t.clone().sequence_values::<Value>() {
                            match to_token(&x?) {
                                Some(tk) => toks.push(tk),
                                None => return Err(fatal(lua, 1)),
                            }
                        }
                    }
                    other => match to_token(other) {
                        Some(tk) => toks.push(tk),
                        None => return Err(fatal(lua, 4)),
                    },
                }
            }
            // LuaTeX shows these tokens as `<inserted text>`.
            with_engine(lua, |e| e.ins_list(Rc::from(toks)))
        })?,
    )?;
    t.set(
        "expand",
        lua.create_function(|lua, ()| {
            with_engine(lua, |e| {
                e.get_next();
                if e.cur_cmd.is_expandable() {
                    e.expand();
                } else {
                    e.cur_tok = e.make_cur_tok();
                    e.back_input();
                }
            })
        })?,
    )?;
    t.set(
        "get_macro",
        lua.create_function(|lua, name: mlua::String| {
            let n = name.as_bytes().to_vec();
            let s = with_engine(lua, |e| {
                let cs = e.hash.lookup(&n)?;
                let m = e.eqtb.meaning(cs).mac.clone()?;
                let body: Vec<Token> = {
                    let i = m.tokens.iter().position(|t| *t == Token::END_MATCH)?;
                    m.tokens[i + 1..].to_vec()
                };
                Some(e.print_to_string(|e| e.show_token_list(&body, 0, body.len(), i64::MAX)))
            })?;
            match s {
                Some(s) => Ok(Value::String(lua.create_string(&s)?)),
                None => Ok(Value::Nil),
            }
        })?,
    )?;
    t.set(
        "get_meaning",
        lua.create_function(|lua, name: mlua::String| {
            let n = name.as_bytes().to_vec();
            let s = with_engine(lua, |e| {
                let cs = e.hash.lookup(&n)?;
                let m = e.eqtb.meaning(cs).mac.clone()?;
                let toks = m.tokens.clone();
                Some(e.print_to_string(|e| e.show_token_list(&toks, 0, toks.len(), i64::MAX)))
            })?;
            match s {
                Some(s) => Ok(Value::String(lua.create_string(&s)?)),
                None => Ok(Value::Nil),
            }
        })?,
    )?;
    t.set(
        "set_macro",
        lua.create_function(|lua, args: MultiValue| set_macro(lua, args))?,
    )?;
    t.set(
        "set_char",
        lua.create_function(|lua, (name, c, g): (mlua::String, i64, Option<String>)| {
            let n = name.as_bytes().to_vec();
            with_engine(lua, |e| {
                let cs = cs_by_name(e, &n);
                e.define(
                    g.as_deref() == Some("global"),
                    Loc::Cs(cs),
                    Val::Meaning(Meaning::prim(Cmd::CharGiven, c as u32)),
                );
            })
        })?,
    )?;
    t.set(
        "set_lua",
        lua.create_function(|lua, args: MultiValue| {
            let v: Vec<Value> = args.into_iter().collect();
            let name = match v.first() {
                Some(Value::String(s)) => s.as_bytes().to_vec(),
                _ => return Err(mlua::Error::runtime("name expected")),
            };
            let id = match v.get(1) {
                Some(Value::Integer(n)) => *n,
                _ => 0,
            };
            let mut global = false;
            let mut protected = false;
            for x in &v[2.min(v.len())..] {
                if let Value::String(s) = x {
                    match s.as_bytes().as_ref() {
                        b"global" => global = true,
                        b"protected" => protected = true,
                        _ => {}
                    }
                }
            }
            with_engine(lua, |e| {
                let cs = cs_by_name(e, &name);
                let cmd = if protected {
                    Cmd::LuaCall
                } else {
                    Cmd::LuaExpandableCall
                };
                e.define(
                    global,
                    Loc::Cs(cs),
                    Val::Meaning(Meaning::prim(cmd, id as u32)),
                );
            })
        })?,
    )?;
    lua.globals().set("token", t)?;
    let tok: Table = lua.globals().get("token")?;
    // Common aliases.
    for (alias, orig) in [("scan_code", "scan_int"), ("scan_real", "scan_int")] {
        let f: Value = tok.get(orig)?;
        tok.set(alias, f)?;
    }
    Ok(())
}

/// `token.scan_string`: a braced text (expanded) or a sequence of
/// non-blank characters.
fn scan_string(e: &mut Engine) -> Vec<u8> {
    e.get_nonblank_noncall();
    if e.cur_cmd == Cmd::LeftBrace {
        e.back_input();
        let toks = e.scan_toks_expanded();
        return e.print_to_string(|e| e.show_token_list(&toks, 0, toks.len(), i64::MAX));
    }
    let mut out = Vec::new();
    loop {
        if e.cur_cs == NO_CS && e.cur_cmd != Cmd::Spacer && e.cur_cmd != Cmd::Relax {
            retex_engine::utf8x::push(&mut out, e.cur_chr);
            e.get_x_token();
        } else {
            if e.cur_cmd != Cmd::Spacer {
                e.back_input();
            }
            break;
        }
    }
    out
}

/// `token.scan_argument`: a braced text or a single token, as a string.
///
/// Observed in LuaTeX: spaces and `\relax` are skipped; a macro gives its
/// replacement text, as if it were braced (expanded unless `expand` is
/// false); other expandable commands are expanded when `expand` is true.
fn scan_argument(e: &mut Engine, mut expand: bool) -> Vec<u8> {
    loop {
        e.get_token();
        if e.cur_cmd == Cmd::Spacer || e.cur_cmd == Cmd::Relax {
            continue;
        }
        if e.cur_cmd >= Cmd::Call
            && e.cur_cs != NO_CS
            && let Some(m) = e.eqtb.meaning(e.cur_cs).mac.clone()
        {
            let start = m
                .tokens
                .iter()
                .position(|t| *t == Token::END_MATCH)
                .map_or(0, |i| i + 1);
            if start > 1 {
                // A macro with parameters is called in braces, and reads
                // its arguments from what follows (observed: "Argument of
                // \z has an extra }"). Its replacement text alone would
                // refer to parameters that do not exist.
                let l = [Token::LEFT_BRACE, e.cur_tok, Token::RIGHT_BRACE];
                e.back_list(std::rc::Rc::from(l));
                e.get_token();
                expand = true;
                break;
            }
            let mut l = Vec::with_capacity(m.tokens.len() - start + 2);
            l.push(Token::LEFT_BRACE);
            l.extend_from_slice(&m.tokens[start..]);
            l.push(Token::RIGHT_BRACE);
            e.back_list(std::rc::Rc::from(l));
            e.get_token();
            break;
        }
        if expand && e.cur_cmd.is_expandable() {
            e.expand();
            continue;
        }
        break;
    }
    if e.cur_cmd == Cmd::LeftBrace {
        e.back_input();
        let toks = if expand {
            e.scan_toks_expanded()
        } else {
            e.scan_general_text().to_vec()
        };
        return e.print_to_string(|e| e.show_token_list(&toks, 0, toks.len(), i64::MAX));
    }
    let t = e.cur_tok;
    e.print_to_string(|e| e.show_token_list(&[t], 0, 1, i64::MAX))
}

/// `token.set_macro([catcodetable,] name [, value [, "global"]])`.
fn set_macro(lua: &Lua, args: MultiValue) -> mlua::Result<()> {
    let mut v: Vec<Value> = args.into_iter().collect();
    let mut cattable: Option<u32> = None;
    if let Some(Value::Integer(n)) = v.first() {
        cattable = Some(*n as u32);
        v.remove(0);
    }
    let name = match v.first() {
        Some(Value::String(s)) => s.as_bytes().to_vec(),
        _ => return Err(mlua::Error::runtime("token.set_macro: name expected")),
    };
    let body: Vec<u8> = match v.get(1) {
        Some(Value::String(s)) => s.as_bytes().to_vec(),
        _ => Vec::new(),
    };
    let global = v
        .iter()
        .skip(2)
        .any(|x| matches!(x, Value::String(s) if s.as_bytes().as_ref() == b"global"));
    with_engine(lua, |e| {
        let cs = cs_by_name(e, &name);
        // Tokenize the body with the given catcode table (or the current).
        let table = cattable.unwrap_or_else(|| e.eqtb.cur_cat_table());
        let mut toks = vec![Token::END_MATCH];
        toks.extend(tokenize_string(e, &body, table));
        let m = Meaning {
            cmd: Cmd::Call,
            chr: 0,
            mac: Some(Rc::new(Macro {
                protected: false,
                tokens: Rc::from(toks),
            })),
        };
        e.define(global, Loc::Cs(cs), Val::Meaning(m));
    })
}

/// Tokenizes a string with a catcode table (simplified scanner: control
/// words, single-character control sequences and characters).
pub fn tokenize_string(e: &mut Engine, s: &[u8], table: u32) -> Vec<Token> {
    let chars = retex_engine::utf8x::chars(s);
    let mut out = Vec::new();
    let mut i = 0;
    while i < chars.len() {
        let c = chars[i];
        let cat = e.eqtb.catcode_in(table, c);
        i += 1;
        match cat {
            0 => {
                let start = i;
                if i < chars.len() && e.eqtb.catcode_in(table, chars[i]) == 11 {
                    while i < chars.len() && e.eqtb.catcode_in(table, chars[i]) == 11 {
                        i += 1;
                    }
                } else if i < chars.len() {
                    i += 1;
                }
                let mut name = Vec::new();
                for &ch in &chars[start..i] {
                    retex_engine::utf8x::push(&mut name, ch);
                }
                out.push(Token::cs(e.hash.lookup_or_insert(&name)));
            }
            13 => out.push(Token::cs(e.hash.active(c))),
            9 | 14 | 15 => {}
            10 => out.push(Token::SPACE),
            _ => out.push(Token::char(Cmd::from_u8(cat), c)),
        }
    }
    out
}