calckey/packages/client/src/scripts/katex-macro.ts

type KaTeXMacro = {
  args: number;
  rule: (string | number)[];
};

function parseSingleKaTeXMacro(src: string): [string, KaTeXMacro] {
	const invalid: [string, KaTeXMacro] = ["", { args: 0, rule: [] }];

	const skipSpaces = (pos: number): number => {
		while (src[pos] === " ")
			++pos;
		return pos;
	};

	if (!src.startsWith("\\newcommand") || src.slice(-1) !== "}")
		return invalid;

	// current index we are checking (= "\\newcommand".length)
	let currentPos: number = 11;
	currentPos = skipSpaces(currentPos);

	// parse {\name}, (\name), or [\name]
	let bracket: string;
	if (src[currentPos] === "{")
		bracket = "{}";
	else if (src[currentPos] === "(")
		bracket = "()";
	else if (src[currentPos] === "[")
		bracket = "[]";
	else
		return invalid;

	++currentPos;
	currentPos = skipSpaces(currentPos);

	if (src[currentPos] !== "\\")
		return invalid;

	const closeNameBracketPos: number = src.indexOf(bracket[1], currentPos);
	if (closeNameBracketPos === -1)
		return invalid;

	const name: string = src.slice(currentPos + 1, closeNameBracketPos).trim();
	if (!/^[a-zA-Z]+$/.test(name))
		return invalid;

	currentPos = skipSpaces(closeNameBracketPos + 1);

	let macro: KaTeXMacro = { args: 0, rule: [] };

	// parse [number of arguments] (optional)
	if (src[currentPos] === "[") {
		const closeArgsBracketPos: number = src.indexOf("]", currentPos);
		macro.args = Number(src.slice(currentPos + 1, closeArgsBracketPos).trim());
		currentPos = closeArgsBracketPos + 1;

		if (Number.isNaN(macro.args) || macro.args < 0)
			return invalid;
	} else if (src[currentPos] === "{") {
		macro.args = 0;
	} else {
		return invalid;
	}

	currentPos = skipSpaces(currentPos);

	// parse {rule}
	if (src[currentPos] !== "{")
		return invalid;

	++currentPos;
	currentPos = skipSpaces(currentPos);

	while (currentPos < src.length - 1) {
		let numbersignPos: number = -1;
		let isEscaped: boolean = false;

		for (let i = currentPos; i < src.length - 1; ++i) {
			if (src[i] !== "\\" && src[i] !== "#") {
				isEscaped = false;
				continue;
			}
			if (src[i] === "\\") {
				isEscaped = !isEscaped;
				continue;
			}
			if (!isEscaped && src[i] === "#") {
				numbersignPos = i;
				break;
			}
		}
		if (numbersignPos === -1) {
			macro.rule.push(src.slice(currentPos, -1));
			break;
		}

		const argIndexEndPos = src.slice(numbersignPos + 1).search(/[^\d]/) + numbersignPos;
		const argIndex: number = Number(src.slice(numbersignPos + 1, argIndexEndPos + 1));

		if (Number.isNaN(argIndex) || argIndex < 1 || macro.args < argIndex)
			return invalid;

		if (currentPos !== numbersignPos)
			macro.rule.push(src.slice(currentPos, numbersignPos));
		macro.rule.push(argIndex);

		currentPos = argIndexEndPos + 1;
	}

	if (macro.args === 0)
		return [name, macro];
	else
		return [name + bracket[0], macro];
}

export function parseKaTeXMacros(src: string): string {
	let result: { [name: string]: KaTeXMacro } = {};

	for (const s of src.split("\n")) {
		const [name, macro]: [string, KaTeXMacro] = parseSingleKaTeXMacro(s.trim());
		if (name !== "")
			result[name] = macro;
	}

	return JSON.stringify(result);
}

// returns [expanded text, whether something is expanded, how many times we can expand more]
// the boolean value is used for multi-pass expansions (macros can expand to other macros)
function expandKaTeXMacroOnce(src: string, macros: { [name: string]: KaTeXMacro }, maxNumberOfExpansions: number)
	: [string, boolean, number] {
	const bracketKinds = 3;
	const openBracketId: { [bracket: string]: number } = {"(": 0, "{": 1, "[": 2};
	const closeBracketId: { [bracket: string]: number } = {")": 0, "}": 1, "]": 2};
	const openBracketFromId = ["(", "{", "["];
	const closeBracketFromId = [")", "}", "]"];

	// mappings from open brackets to their corresponding close brackets
	type BracketMapping = { [openBracketPos: number]: number };

	const bracketMapping = ((): BracketMapping => {
		let result: BracketMapping = {};
		const n = src.length;

		let depths = new Array<number>(bracketKinds).fill(0);  // current bracket depth for "()", "{}", and "[]"
		let buffer = Array.from(Array<number[]>(bracketKinds), () => Array<number>(n));

		let isEscaped = false;

		for (let i = 0; i < n; ++i) {
			if (!isEscaped && src[i] === "\\" && i + 1 < n && ["{", "}", "\\"].includes(src[i+1])) {
				isEscaped = true;
				continue;
			}
			if (isEscaped
				|| (src[i] !== "\\"
							&& !openBracketFromId.includes(src[i])
							&& !closeBracketFromId.includes(src[i])))
			{
				isEscaped = false;
				continue;
			}
			isEscaped = false;

			if (openBracketFromId.includes(src[i])) {
				const id: number = openBracketId[src[i]];
				buffer[id][depths[id]] = i;
				++depths[id];
			} else if (closeBracketFromId.includes(src[i])) {
				const id: number = closeBracketId[src[i]];
				if (depths[id] > 0) {
					--depths[id];
					result[buffer[id][depths[id]]] = i;
				}
			}
		}

		return result;
	})();

	function expandSingleKaTeXMacro(expandedArgs: string[], macroName: string): string {
		let result = "";
		for (const block of macros[macroName].rule) {
			if (typeof block === "string")
				result += block;
			else
				result += expandedArgs[block - 1];
		}
		return result;
	}

	// only expand src.slice(beginPos, endPos)
	function expandKaTeXMacroImpl(beginPos: number, endPos: number): [string, boolean] {
		if (endPos <= beginPos)
			return ["", false];

		const raw: string = src.slice(beginPos, endPos);
		const fallback: string = raw;  // returned for invalid inputs or too many expansions

		if (maxNumberOfExpansions <= 0)
			return [fallback, false];
		--maxNumberOfExpansions;

		// search for a custom macro
		let checkedPos = beginPos - 1;
		let macroName = "";
		let macroBackslashPos = 0;

		// for macros w/o args: unused
		//            w/  args: the first open bracket ("(", "{", or "[") after cmd name
		let macroArgBeginPos = 0;

		// for macros w/o args: the end of cmd name
		//            w/  args: the closing bracket of the last arg
		let macroArgEndPos = 0;

		while (checkedPos < endPos) {
			checkedPos = src.indexOf("\\", checkedPos + 1);

			// there is no macro to expand
			if (checkedPos === -1)
				return [raw, false];

			// is it a custom macro?
			let nonAlphaPos = src.slice(checkedPos + 1).search(/[^A-Za-z]/) + checkedPos + 1;

			if (nonAlphaPos === checkedPos)
				nonAlphaPos = endPos;

			let macroNameCandidate = src.slice(checkedPos + 1, nonAlphaPos);
			if (macros.hasOwnProperty(macroNameCandidate)) {
				// this is a custom macro without args
				macroBackslashPos = checkedPos;
				macroArgEndPos = nonAlphaPos - 1;
				macroName = macroNameCandidate;
				break;
			}

			let nextOpenBracketPos = endPos;
			for (let i = 0; i < bracketKinds; ++i) {
				const pos = src.indexOf(openBracketFromId[i], checkedPos + 1);
				if (pos !== -1 && pos < nextOpenBracketPos)
					nextOpenBracketPos = pos;
			}

			if (nextOpenBracketPos === endPos)
				continue;  // there is no open bracket

			macroNameCandidate += src[nextOpenBracketPos];

			if (macros.hasOwnProperty(macroNameCandidate)) {
				macroBackslashPos = checkedPos;
				macroArgBeginPos = nextOpenBracketPos;
				macroArgEndPos = nextOpenBracketPos;  // to search the first arg from here
				macroName = macroNameCandidate;
				break;
			}
		}

		const numArgs: number = macros[macroName].args;
		const openBracket: string = macroName.slice(-1);

		let expandedArgs = new Array<string>(numArgs);

		for (let i = 0; i < numArgs; ++i) {
			// find the first open bracket after what we've searched
			const nextOpenBracketPos = src.indexOf(openBracket, macroArgEndPos);
			if (nextOpenBracketPos === -1)
				return [fallback, false];  // not enough arguments are provided
			if (!bracketMapping[nextOpenBracketPos])
				return [fallback, false];  // found open bracket doesn't correspond to any close bracket

			macroArgEndPos = bracketMapping[nextOpenBracketPos];
			expandedArgs[i] = expandKaTeXMacroImpl(nextOpenBracketPos + 1, macroArgEndPos)[0];
		}

		return [src.slice(beginPos, macroBackslashPos)
						+ expandSingleKaTeXMacro(expandedArgs, macroName)
						+ expandKaTeXMacroImpl(macroArgEndPos + 1, endPos)[0], true];
	}

	const [expandedText, expandedFlag]: [string, boolean] = expandKaTeXMacroImpl(0, src.length);
	return [expandedText, expandedFlag, maxNumberOfExpansions];
}

export function expandKaTeXMacro(src: string, macrosAsJSONString: string, maxNumberOfExpansions: number): string {
	const macros = JSON.parse(macrosAsJSONString);

	let expandMore = true;

	while (expandMore && (0 < maxNumberOfExpansions))
		[src, expandMore, maxNumberOfExpansions] = expandKaTeXMacroOnce(src, macros, maxNumberOfExpansions);

	return src;
}