scriptorium/internal/mathml/macros.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package mathml

import "strings"

// macro is a user-defined command: the number of arguments its body takes
// and the body itself as tokens.
type macro struct {
	args int
	body []token
}

// The depth limits one call site from nesting forever; the budget bounds
// the whole parse, so a self-splicing macro can never outrun the parser.
const (
	macroExpansionDepth  = 64
	macroExpansionBudget = 10000
)

// expandMacros replaces the macro call at the cursor with its expanded
// body, ready for the parser to read. A call beyond the limits degrades
// to its own source.
func (p *parser) expandMacros() {
	for range macroExpansionDepth {
		t := p.toks[p.pos]
		if t.kind != tokCommand {
			return
		}
		if p.expansions >= macroExpansionBudget {
			p.degradeAt(p.pos, t)
			return
		}
		m, ok := p.macros[t.text[1:]]
		if !ok {
			return
		}
		j := p.pos + 1
		args := make([][]token, 0, m.args)
		for range m.args {
			arg, next, ok := p.argTokens(j)
			if !ok {
				break
			}
			args = append(args, arg)
			j = next
		}
		if len(args) != m.args {
			p.degradeAt(p.pos, t)
			return
		}
		var body []token
		for i := 0; i < len(m.body); i++ {
			bt := m.body[i]
			if bt.kind == tokChar && bt.text == "#" && i+1 < len(m.body) &&
				m.body[i+1].kind == tokChar && len(m.body[i+1].text) == 1 &&
				isDigitByte(m.body[i+1].text[0]) {
				k := int(m.body[i+1].text[0] - '0')
				if k >= 1 && k <= len(args) {
					body = append(body, args[k-1]...)
				}
				i++
				continue
			}
			body = append(body, bt)
		}
		// The spliced tokens carry the call site as their position, so a
		// construct that fails inside a macro degrades at the call.
		for i := range body {
			body[i].start = t.start
			body[i].end = t.end
		}
		spliced := make([]token, 0, len(p.toks)-(j-p.pos)+len(body))
		spliced = append(spliced, p.toks[:p.pos]...)
		spliced = append(spliced, body...)
		spliced = append(spliced, p.toks[j:]...)
		p.toks = spliced
		p.expansions++
	}
	t := p.toks[p.pos]
	p.degradeAt(p.pos, t)
}

// degradeAt replaces one token with a degraded token.
func (p *parser) degradeAt(i int, t token) {
	p.toks[i] = token{kind: tokDegraded, text: t.text, start: t.start, end: t.end}
}

// argTokens reads one macro argument from position j: a braced group with
// its braces, or a single token.
func (p *parser) argTokens(j int) ([]token, int, bool) {
	if j >= len(p.toks) || p.toks[j].kind == tokEOF {
		return nil, j, false
	}
	if p.toks[j].kind != tokLBrace {
		return []token{p.toks[j]}, j + 1, true
	}
	depth := 0
	for k := j; k < len(p.toks); k++ {
		switch p.toks[k].kind {
		case tokLBrace:
			depth++
		case tokRBrace:
			depth--
			if depth == 0 {
				group := make([]token, k+1-j)
				copy(group, p.toks[j:k+1])
				return group, k + 1, true
			}
		case tokEOF:
			return nil, j, false
		}
	}
	return nil, j, false
}

// macroDefinition registers a \newcommand or \def style definition and
// produces no output. The cursor sits just after the definition command.
func (p *parser) macroDefinition(kind string) *node {
	source := `\` + kind
	if kind == "DeclareMathOperator" || kind == "DeclareMathOperator*" {
		name := p.defName()
		if name == "" {
			return errorNode(source)
		}
		body, ok := p.rawBraced()
		if !ok {
			return errorNode(source)
		}
		wrap := `\operatorname{` + body + `}`
		if strings.HasSuffix(kind, "*") {
			wrap = `\operatorname*{` + body + `}`
		}
		p.macros[name] = macro{body: tokenise([]byte(wrap))}
		return nil
	}
	if kind == "def" || kind == "gdef" {
		return p.tecDefinition(source)
	}
	name := p.defName()
	if name == "" {
		return errorNode(source)
	}
	args := 0
	if count, ok := p.bracketArg(); ok && count != "" {
		n := 0
		for i := 0; i < len(count); i++ {
			if !isDigitByte(count[i]) {
				return errorNode(source)
			}
			n = n*10 + int(count[i]-'0')
		}
		if n > 9 {
			return errorNode(source)
		}
		args = n
	}
	body, ok := p.rawBraced()
	if !ok {
		return errorNode(source)
	}
	p.macros[name] = macro{args: args, body: tokenise([]byte(body))}
	return nil
}

// tecDefinition registers a \def, whose parameter text names undelimited
// arguments with #1 up to #9.
func (p *parser) tecDefinition(source string) *node {
	name := p.defName()
	if name == "" {
		return errorNode(source)
	}
	args := 0
	for {
		t := p.cur()
		if t.kind == tokChar && t.text == "#" {
			p.pos++
			d := p.cur()
			if d.kind != tokChar || len(d.text) != 1 || !isDigitByte(d.text[0]) {
				return errorNode(source)
			}
			if int(d.text[0]-'0') != args+1 {
				return errorNode(source)
			}
			args++
			p.pos++
			continue
		}
		break
	}
	body, ok := p.rawBraced()
	if !ok {
		return errorNode(source)
	}
	p.macros[name] = macro{args: args, body: tokenise([]byte(body))}
	return nil
}

// defName reads the name a definition declares: a braced command or a
// bare command.
func (p *parser) defName() string {
	if p.at(tokLBrace) {
		p.pos++
		if p.at(tokCommand) {
			name := p.cur().text[1:]
			p.pos++
			if p.at(tokRBrace) {
				p.pos++
				return name
			}
		}
		return ""
	}
	if p.at(tokCommand) {
		name := p.cur().text[1:]
		p.pos++
		return name
	}
	return ""
}