// Generates engine/keywords.zena (keyword ids, sets, and the keywordAt lookup) and engine/keywords.wat (the
// lookup's table, a data segment scripts/build.sh merges in). Run: node scripts/gen/keywords.mjs
import fs from "node:fs";
import path from "node:path";
import { KW_BASE, TEXT_BASE, DATA_END, LEXC_BASE, LEXC_SIZE, dataModule } from "./layout.mjs";

// Reserved words first: everything below K_LET cannot be an identifier.
const reserved = `break case catch class const continue debugger default delete do else enum export extends false
finally for function if import in instanceof new null return super switch this throw true try typeof var void while with`;
// Every word in a set below must get an id under 64 (the sets are i64 masks), so they come first.
const contextual = `let static yield await async of get set as satisfies type interface declare namespace module abstract
readonly public private protected override accessor global using implements is keyof infer unique asserts require from out
any unknown string number bigint symbol boolean never object undefined constructor target meta assert defer`;

// Keyword sets the engine tests with one mask lookup (keywordIn), named for TypeScript's predicates.
const sets = {
  // parseStatement's tokens that go through isStartOfDeclaration, plus using and await (isUsingDeclaration).
  HEAD: "async interface type module namespace declare const enum export import private protected public abstract accessor static readonly global using await",
  // isModifierKind, less in/out, which cannot follow a statement's first token.
  STMT_MOD: "abstract accessor async const declare default export override private protected public readonly static",
  // Class and object member modifiers, with get/set.
  MEMBER_MOD: "static public private protected readonly abstract override declare accessor async get set export",
  // The modifiers ts-blank-space removes (classElementModifiersToRemove).
  REMOVED: "private protected public readonly override abstract declare",
  // isStartOfExpression's keywords (contextual ones are identifiers there).
  EXPR_START: "this super null true false function class new delete typeof void import",
  // isDeclAhead's declaration keywords that commit at once.
  DECL: "var let const function class enum",
  // isTypeStart's reserved words.
  TYPE_START: "void null true false typeof this new import",
  // types.zena's type-member modifiers.
  TYPE_MEMBER_MOD: "readonly public private protected static declare accessor override get set",
  // downlevel.zena's reservedName: the contextual words reserved in strict code.
  STRICT_RESERVED: "let static yield await implements interface private protected public",
  // downlevel.zena's declSpans: keywords that start a declaration whose names follow.
  DL_DECL: "var let const function class namespace module",
  // isBinaryOperatorToken's keyword operators.
  BINARY_WORD: "in instanceof as satisfies",
  // Statements whose keyword is followed by a parenthesized head.
  PAREN_HEAD: "if while with switch",
  // A constructor parameter's property modifiers, less override.
  PARAM_PROP: "public private protected readonly",
  // The keyword types that never take a .Qualifier (parseKeywordAndNoDot).
  NO_DOT_TYPE: "void null true false",
};
// Sets with ids of 64 and up: KS_<name> holds the ids under 64, KS_<name>_HI the rest less 64 (keywordInHi).
const wideSets = {
  // isKeywordType: the keyword types that take no type arguments.
  KEYWORD_TYPE: "any unknown string number bigint symbol boolean never object undefined void null true false",
};
// Token-type sets (lex.zena's T_ values, all under 64), tested with tokenIn.
const tokenSets = {
  EXPR_START: "NUM STR TEMPLATE TEMPLATE_HEAD LPAREN LBRACK LBRACE SLASH PLUS MINUS TILDE BANG INC DEC LT PRIVATE AT",
  TYPE_START: "STR NUM TEMPLATE TEMPLATE_HEAD LBRACE LBRACK LT BAR AMP STAR QUESTION BANG ELLIPSIS LPAREN MINUS",
  PROPERTY_NAME: "NAME STR NUM LBRACK PRIVATE",
  CAN_FOLLOW_MOD: "LBRACK LBRACE STAR ELLIPSIS NAME STR NUM",
  PREFIX_OP: "BANG TILDE PLUS MINUS INC DEC",
  LITERAL: "NUM STR TEMPLATE PRIVATE REGEX",
  OPERAND_END: "NAME NUM STR TEMPLATE PRIVATE INC DEC",
  CLOSE_OR_SEP: "RPAREN RBRACK RBRACE COMMA SEMI",
  OPEN: "LPAREN LBRACK LBRACE",
  CLOSE: "RPAREN RBRACK RBRACE",
  BINDING_START: "NAME LBRACK LBRACE",
  PARAM_START: "NAME LBRACK LBRACE ELLIPSIS",
  STR_NUM_TPL: "STR NUM TEMPLATE",
  JSDOC_NULLABLE_END: "COMMA RBRACE RPAREN GT EQ BAR",
  PARAM_FOLLOW: "COLON COMMA QUESTION EQ",
  OPTIONAL_FOLLOW: "COLON COMMA RPAREN EQ",
  MEMBER_NAME: "NAME STR NUM PRIVATE",
  POSTFIX_UPDATE: "BANG INC DEC",
};

const words = [...reserved.split(/\s+/), ...contextual.split(/\s+/)].filter(Boolean);
const firstContextual = reserved.split(/\s+/).filter(Boolean).length + 1;

function pack(word) {
  const bytes = Buffer.from(word, "ascii");
  let w0 = 0n;
  for (let i = 0; i < Math.min(8, bytes.length); i++) w0 |= BigInt(bytes[i]) << BigInt(8 * i);
  let w1 = 0;
  for (let i = 8; i < Math.min(12, bytes.length); i++) w1 |= bytes[i] << (8 * (i - 8));
  return { w0, w1, n: bytes.length };
}

let out = `// GENERATED by scripts/gen/keywords.mjs. Do not edit by hand.
// Keyword ids for engine/lex.zena. Ids below K_LET are reserved words.
import { Memory } from 'zena:memory';

export let K_NONE: i32 = 0;
`;
words.forEach((w, i) => { out += `export let K_${w.toUpperCase()}: i32 = ${i + 1};\n`; });
out += `export let K_FIRST_CONTEXTUAL: i32 = ${firstContextual};\n\n`;
for (const [name, list] of Object.entries(sets)) {
  let mask = 0n;
  for (const w of list.split(" ")) {
    const id = words.indexOf(w) + 1;
    if (id < 1 || id > 63) throw new Error(`${name}: ${w} has id ${id}`);
    mask |= 1n << BigInt(id);
  }
  out += `export let KS_${name}: i64 = 0x${mask.toString(16)};\n`;
}
out += `
// Whether keyword id k is in the set mask s.
export function keywordIn(k: i32, s: i64): boolean {
  return k < 64 && ((s >> (k as i64)) & (1 as i64)) != (0 as i64);
}

// Whether keyword id k, 64 or more, is in the set mask s of ids less 64.
export function keywordInHi(k: i32, s: i64): boolean {
  return k >= 64 && ((s >> ((k - 64) as i64)) & (1 as i64)) != (0 as i64);
}

// Whether token type t is in the set mask s (lex.zena's tokenIn tests the current token).
export function tokenTypeIn(t: i32, s: i64): boolean {
  return t < 64 && ((s >> (t as i64)) & (1 as i64)) != (0 as i64);
}

`;
const lexSrc = fs.readFileSync(path.resolve(import.meta.dirname, "../../engine/lex.zena"), "utf8");
const TT = Object.fromEntries([...lexSrc.matchAll(/export let T_(\w+): i32 = (\d+);/g)].map((m) => [m[1], +m[2]]));
for (const [name, list] of Object.entries(wideSets)) {
  let lo = 0n, hi = 0n;
  for (const w of list.split(" ")) {
    const id = words.indexOf(w) + 1;
    if (id < 1 || id > 127) throw new Error(`${name}: ${w} has id ${id}`);
    if (id < 64) lo |= 1n << BigInt(id);
    else hi |= 1n << BigInt(id - 64);
  }
  out += `export let KS_${name}: i64 = 0x${lo.toString(16)};\nexport let KS_${name}_HI: i64 = 0x${hi.toString(16)};\n`;
}
for (const [name, list] of Object.entries(tokenSets)) {
  let mask = 0n;
  for (const t of list.split(" ")) {
    const v = TT[t];
    if (v === undefined || v > 63) throw new Error(`token set ${name}: T_${t} is ${v}`);
    mask |= 1n << BigInt(v);
  }
  out += `export let TS_${name}: i64 = 0x${mask.toString(16)};\n`;
}

// keywordAt finds a candidate with one multiplicative hash of the first 8 bytes (masked to the length) into
// a 256-slot table, then checks it against the keyword's own bytes. The multiplier is any odd one that sends
// every keyword to its own slot; findMultiplier only runs when the word list changes and it stops doing so.
const SLOTS = 256;
const M64 = (1n << 64n) - 1n;
const slotOf = (w0, mul) => Number(((w0 * mul) & M64) >> 56n);
const packed = words.map((w) => pack(w).w0);
function collisionFree(mul) {
  const seen = new Set(packed.map((w0) => slotOf(w0, mul)));
  return seen.size === packed.length;
}
function findMultiplier() {
  let seed = 0x9e3779b97f4a7c15n;
  for (;;) {
    seed = (seed * 6364136223846793005n + 1442695040888963407n) & M64;
    const mul = (seed >> 1n) | 1n; // under 2^63, so it is a plain i64 literal
    if (collisionFree(mul)) return mul;
  }
}
let mul = 0x3678c5607077029dn;
if (!collisionFree(mul)) {
  mul = findMultiplier();
  console.log(`new multiplier 0x${mul.toString(16)}: update keywords.mjs`);
}

// Layout at KW_BASE: slot -> id (u8 x 256), id -> start offset in the text (u16 x (ids + 1), so that
// the next id's start is this one's end), then every keyword's bytes in id order.
const text = Buffer.from(words.join(""), "ascii");
const slot = Buffer.alloc(SLOTS);
words.forEach((w, i) => { slot[slotOf(packed[i], mul)] = i + 1; });
const offs = Buffer.alloc(2 * (words.length + 2));
let at = 0;
words.forEach((w, i) => { at += w.length; offs.writeUInt16LE(at, 2 * (i + 2)); });
const KW_SLOT = KW_BASE;
const KW_OFF = KW_SLOT + SLOTS;
const KW_TEXT = KW_OFF + offs.length;

// Punctuators, for lex.zena's scan: PUNCT maps an ASCII byte to an entry of 5 token types (0 for none):
// the byte alone, followed by `=`, followed by `==`, doubled, doubled then `=`. Entry 0 is the stray
// byte's one-byte token. Token type values are read from lex.zena's T_ constants.
const lex = fs.readFileSync(path.resolve(import.meta.dirname, "../../engine/lex.zena"), "utf8");
const T = Object.fromEntries([...lex.matchAll(/export let T_(\w+): i32 = (\d+);/g)].map((m) => [m[1], +m[2]]));
const OPS = [
  ["", "TILDE"],
  ["(", "LPAREN"], [")", "RPAREN"], ["{", "LBRACE"], ["}", "RBRACE"], ["[", "LBRACK"], ["]", "RBRACK"],
  [";", "SEMI"], [",", "COMMA"], [":", "COLON"], ["@", "AT"], [">", "GT"],
  ["?", "QUESTION", "", "", "QQ", "QQ_EQ"],
  ["=", "EQ", "EQEQ", "EQEQEQ"],
  ["!", "BANG", "NEQ", "NEQEQ"],
  ["+", "PLUS", "PLUS_EQ", "", "INC"],
  ["-", "MINUS", "MINUS_EQ", "", "DEC"],
  ["*", "STAR", "STAR_EQ", "", "STARSTAR", "STARSTAR_EQ"],
  ["/", "SLASH", "SLASH_EQ"],
  ["%", "PERCENT", "PERCENT_EQ"],
  ["&", "AMP", "AMP_EQ", "", "AMPAMP", "AMPAMP_EQ"],
  ["|", "BAR", "BAR_EQ", "", "BARBAR", "BARBAR_EQ"],
  ["^", "CARET", "CARET_EQ"],
  ["<", "LT", "LTE", "", "SHL", "SHL_EQ"],
];
const pidx = Buffer.alloc(128);
const pent = Buffer.alloc(5 * OPS.length);
OPS.forEach(([ch, ...ts], i) => {
  if (ch) pidx[ch.charCodeAt(0)] = i;
  ts.forEach((t, j) => {
    if (!t) return;
    if (!(t in T) || T[t] > 255) throw new Error(`punctuator token T_${t}`);
    pent[5 * i + j] = T[t];
  });
});
const PUNCT = KW_TEXT + text.length + 16; // the keyword check reads up to 12 bytes past a word's start
const PUNCT_ENT = PUNCT + pidx.length;
const data = Buffer.concat([slot, offs, text, Buffer.alloc(16), pidx, pent]);

out += `// The static data below the frame stack: strip.zena starts the stack at DATA_END.
export let DATA_END: i32 = ${DATA_END};
// The lexer's token replay cache: LEXC_MASK + 1 entries of 16 bytes at LEXC.
export let LEXC: i32 = ${LEXC_BASE};
export let LEXC_MASK: i32 = ${LEXC_SIZE - 1};
let KW_SLOT: i32 = ${KW_SLOT};
let KW_OFF: i32 = ${KW_OFF};
let KW_TEXT: i32 = ${KW_TEXT};
let KW_MUL: i64 = 0x${mul.toString(16)};
export let PUNCT: i32 = ${PUNCT};
export let PUNCT_ENT: i32 = ${PUNCT_ENT};

function lowMask64(n: i32): i64 {
  if (n >= 8) return -1;
  return ((1 as i64) << ((n * 8) as i64)) - (1 as i64);
}

function lowMask32(n: i32): i32 {
  if (n >= 4) return -1;
  return (1 << (n * 8)) - 1;
}

// The keyword id of the n-byte identifier at memory address p, zero when it is not a keyword. Reads 12 bytes
// at p whatever n is, so p needs that much addressable memory after it.
export function keywordAt(p: i32, n: i32): i32 {
  if (n > 11) return 0;
  let m = lowMask64(n);
  let w0 = Memory.default.getI64(p) & m;
  let k = Memory.default.getU8(KW_SLOT + ((((w0 * KW_MUL) >> (56 as i64)) as i32) & 255));
  let o = Memory.default.getI32(KW_OFF + k * 2);
  let s = o & 0xffff;
  if ((o >> 16) - s != n) return 0;
  if ((Memory.default.getI64(KW_TEXT + s) & m) != w0) return 0;
  if (n > 8 && ((Memory.default.getI32(p + 8) ^ Memory.default.getI32(KW_TEXT + s + 8)) & lowMask32(n - 8)) != 0) return 0;
  return k;
}
`;

const dest = path.resolve(import.meta.dirname, "../../engine/keywords.zena");
fs.writeFileSync(dest, out);
const wat = path.resolve(import.meta.dirname, "../../engine/keywords.wat");
fs.writeFileSync(wat, dataModule("keywords.mjs", KW_BASE, data, TEXT_BASE));
console.log(`wrote ${dest} and ${wat}: ${words.length} keywords, ${data.length} bytes of table`);