#!/usr/bin/env node
// cross-reference of the compiler's global state: which function writes a global and which only read it.
//
// the compiler is the files of src/compiler/parts read as one unit. its state is global arrays and scalars
// that functions far apart read and write. this tool lists, for every top-level variable of the compiler,
// the functions that write it and the functions that only read it, and can fail ci when a global gains a
// writer that tests/xref.baseline does not list.
//
// how it works:
//   1. the parts are joined (the same text the language server serves as tinyvm:/compiler.zc) and lexed
//      here; the top-level declarations are found by parsing the lexed text, and so are the brace ranges
//      of the top-level functions.
//   2. exact scoping is not guessed: the joined text is opened as a document in build/bin/zclsp and one
//      textDocument/references request per global is sent, all in one server run. the server answers with
//      the uses the compiler itself resolved to that symbol, so a local that shadows a global is not
//      counted.
//   3. each reference is classified from its neighbouring tokens (see classify) and attributed to the
//      function whose range holds it.
const fs = require('fs');
const path = require('path');
const cp = require('child_process');
const os = require('os');

const ROOT = path.resolve(__dirname, '..', '..');

// ---- lexing ----

const OPS = ['<<=', '>>=', '...', '->', '++', '--', '+=', '-=', '*=', '/=', '%=', '&=', '^=', '|=', '==', '!=', '<=', '>=', '&&', '||', '<<', '>>', '::'];
const ASSIGN_OPS = new Set(['=', '+=', '-=', '*=', '/=', '%=', '&=', '^=', '|=', '<<=', '>>=']);
const IDSTART = /[A-Za-z_]/, IDCHAR = /[A-Za-z0-9_]/;

// tokens: {k: 'id'|'num'|'str'|'chr'|'op', s, off, line, col}; comments are skipped, and a backtick outside
// a comment or literal ends the source, as it does for the compiler.
function lex(text) {
    const toks = [];
    let i = 0, line = 0, lineStart = 0;
    const n = text.length;
    const push = (k, s, off) => toks.push({ k, s, off, line, col: off - lineStart });
    while (i < n) {
        const c = text[i];
        if (c === '\n') { line++; i++; lineStart = i; continue; }
        if (c === ' ' || c === '\t' || c === '\r') { i++; continue; }
        if (c === '/' && text[i + 1] === '/') { while (i < n && text[i] !== '\n') i++; continue; }
        if (c === '/' && text[i + 1] === '*') {
            i += 2;
            while (i < n && !(text[i] === '*' && text[i + 1] === '/')) { if (text[i] === '\n') { line++; lineStart = i + 1; } i++; }
            i += 2; continue;
        }
        if (c === '`') break;
        if (c === '"' || c === "'") {
            const start = i, startLine = line, startCol = i - lineStart;
            i++;
            while (i < n && text[i] !== c) { if (text[i] === '\\') i++; if (text[i] === '\n') { line++; lineStart = i + 1; } i++; }
            i++;
            toks.push({ k: c === '"' ? 'str' : 'chr', s: text.slice(start, i), off: start, line: startLine, col: startCol });
            continue;
        }
        if (IDSTART.test(c)) { const s = i; while (i < n && IDCHAR.test(text[i])) i++; push('id', text.slice(s, i), s); continue; }
        if (c >= '0' && c <= '9') {
            const s = i;
            // digits, hex, and a fraction: a dot only when a digit follows, so "1.u" style tags are not eaten
            while (i < n && (IDCHAR.test(text[i]) || (text[i] === '.' && text[i + 1] >= '0' && text[i + 1] <= '9'))) i++;
            push('num', text.slice(s, i), s); continue;
        }
        let op = null;
        for (const o of OPS) if (text.startsWith(o, i)) { op = o; break; }
        if (!op) op = c;
        push('op', op, i);
        i += op.length;
    }
    return toks;
}

// ---- top-level structure ----

const NOT_DECL = new Set(['const', 'struct', 'enum', 'namespace', 'include', 'embed', 'typedef', 'type', 'enhance', 'fn', 'macro', 'immediate', 'extern', 'cfunc', 'import', 'static', 'quote']);

// finds the top-level functions (name, token range, source range) and the top-level variable declarations.
// returns {functions, globals, problems}
function structure(toks) {
    const functions = [], globals = [], problems = [];
    let i = 0;
    const n = toks.length;
    const closeOf = { '(': ')', '[': ']', '{': '}' };
    // index of the token matching the opener at i
    const match = j => {
        let depth = 0;
        for (; j < n; j++) {
            const s = toks[j].k === 'op' ? toks[j].s : '';
            if (s === '(' || s === '[' || s === '{') depth++;
            else if (s === ')' || s === ']' || s === '}') { depth--; if (depth === 0) return j; }
        }
        return n - 1;
    };
    while (i < n) {
        const start = i;
        // scan the statement head: up to ';' or '{' at depth 0
        let j = i, paren = -1;
        while (j < n) {
            const t = toks[j];
            if (t.k === 'op') {
                if (t.s === ';' || t.s === '{') break;
                if (t.s === '(' && paren < 0) paren = j;
                if (t.s === '(' || t.s === '[') { j = match(j) + 1; continue; }
            }
            j++;
        }
        if (j >= n) break;
        const first = toks[start];
        if (toks[j].s === '{') {
            const end = match(j);
            const isFn = paren >= 0 && !NOT_DECL.has(first.s) || (first.s === 'fn' && paren >= 0);
            if (isFn) {
                // name: the identifier before the first '(' (a method is Type.name)
                let k = paren - 1, name = toks[k].s;
                while (k >= 2 && toks[k - 1].s === '.' && toks[k - 2].k === 'id') { name = toks[k - 2].s + '.' + name; k -= 2; }
                functions.push({ name, startTok: start, endTok: end, from: toks[start].off, to: toks[end].off + 1, line: toks[paren - 1].line });
            }
            i = end + 1;
            if (i < n && toks[i].s === ';') i++;
            continue;
        }
        // a ';' statement
        const head = toks.slice(start, j);
        i = j + 1;
        if (!head.length || NOT_DECL.has(first.s) || first.k !== 'id') continue;
        const decl = parseDecl(head);
        if (decl === null) continue;
        if (decl.error) { problems.push(decl.error + ' at line ' + (first.line + 1)); continue; }
        for (const d of decl.names) globals.push(d);
    }
    return { functions, globals, problems };
}

// a type (identifier, then * [..] .tag), then declarators: name [dims] [= init] separated by commas.
function parseDecl(head) {
    let i = 0;
    const n = head.length;
    if (head[i].k !== 'id') return null;
    i++;
    for (;;) {
        const t = head[i];
        if (!t) return null;
        if (t.s === '*') { i++; continue; }
        if (t.s === '.' && head[i + 1] && head[i + 1].k === 'id') { i += 2; continue; }
        if (t.s === '[') { i = skipGroup(head, i) + 1; continue; }
        break;
    }
    // a type needs a name after it; "a.b c" and "x = 1" and calls are not declarations
    const names = [];
    for (;;) {
        const t = head[i];
        if (!t || t.k !== 'id') return names.length ? { error: 'unparsed declarator' } : null;
        names.push({ name: t.s, tok: t });
        i++;
        while (head[i] && head[i].s === '[') i = skipGroup(head, i) + 1;
        if (head[i] && head[i].s === '=') {
            // initializer: to the next ',' at depth 0
            i++;
            while (i < n && head[i].s !== ',') { i = (head[i].s === '(' || head[i].s === '[' || head[i].s === '{') ? skipGroup(head, i) + 1 : i + 1; }
        }
        if (i >= n) return { names };
        if (head[i].s !== ',') return { error: 'unparsed declaration after ' + t.s };
        i++;
    }
}
function skipGroup(toks, i) {
    let depth = 0;
    for (; i < toks.length; i++) {
        const s = toks[i].k === 'op' ? toks[i].s : '';
        if (s === '(' || s === '[' || s === '{') depth++;
        else if (s === ')' || s === ']' || s === '}') { depth--; if (depth === 0) return i; }
    }
    return toks.length - 1;
}

// ---- classification ----

const KEYWORD_BEFORE_OPERAND = new Set(['return', 'else', 'case', 'do', 'in', 'as', 'new', 'sizeof', 'print']);

// is the token at index i the end of an operand (so that a following & is binary)?
function endsOperand(t) {
    if (!t) return false;
    if (t.k === 'num' || t.k === 'str' || t.k === 'chr') return true;
    if (t.k === 'id') return !KEYWORD_BEFORE_OPERAND.has(t.s);
    return t.s === ')' || t.s === ']';
}

// 'w' when the reference at token index i is stored to: the target of = or a compound assignment, the
// operand of ++ or --, an element or field of it stored the same way, or its address is taken; else 'r'.
function classify(toks, i) {
    // the lvalue chain around the name: a.b.name[..].f
    let a = i;
    while (a >= 2 && (toks[a - 1].s === '.' || toks[a - 1].s === '->') && toks[a - 2].k === 'id') a -= 2;
    const before = toks[a - 1], before2 = toks[a - 2];
    if (before && before.k === 'op') {
        if (before.s === '++' || before.s === '--') return 'w';
        if (before.s === '&' && !endsOperand(before2)) return 'w';
    }
    let j = i + 1;
    for (;;) {
        const t = toks[j];
        if (!t || t.k !== 'op') break;
        if (t.s === '[') { j = skipGroup(toks, j) + 1; continue; }
        if ((t.s === '.' || t.s === '->') && toks[j + 1] && toks[j + 1].k === 'id') { j += 2; continue; }
        break;
    }
    const t = toks[j];
    if (t && t.k === 'op' && (ASSIGN_OPS.has(t.s) || t.s === '++' || t.s === '--')) return 'w';
    return 'r';
}

// ---- the language server ----

const frame = m => { const b = Buffer.from(JSON.stringify(m)); return Buffer.concat([Buffer.from('Content-Length: ' + b.length + '\r\n\r\n'), b]); };

// one server run: open text, ask for the references of every position in `positions` ({line, character}),
// resolve to the array of location lists, in the same order. the server is spawned, not run synchronously, so
// several runs can share the work.
function referencesRun(server, text, positions, timeoutMs) {
    return new Promise((resolve, reject) => {
        const uri = 'file:///xref/compiler.zc';
        const msgs = [
            { jsonrpc: '2.0', id: 1, method: 'initialize', params: {} },
            { jsonrpc: '2.0', method: 'textDocument/didOpen', params: { textDocument: { uri, version: 1, languageId: 'zc', text } } }];
        positions.forEach((position, k) => msgs.push({ jsonrpc: '2.0', id: 2 + k, method: 'textDocument/references', params: { textDocument: { uri }, position, context: { includeDeclaration: true } } }));
        msgs.push({ jsonrpc: '2.0', method: 'exit' });
        const child = cp.spawn(server, [], { stdio: ['pipe', 'pipe', 'pipe'] });
        const out = [], err = [];
        const timer = setTimeout(() => { child.kill(); reject(new Error('language server timed out after ' + Math.round(timeoutMs / 1000) + 's (ZC_TEST_TIMEOUT_SCALE stretches it)')); }, timeoutMs);
        child.stdout.on('data', c => out.push(c));
        child.stderr.on('data', c => err.push(c));
        child.stdin.on('error', () => {});
        child.on('error', e => { clearTimeout(timer); reject(new Error('language server ' + server + ': ' + e.message)); });
        child.on('close', status => {
            clearTimeout(timer);
            if (status !== 0) return reject(new Error('language server exited with status ' + status + ': ' + Buffer.concat(err).toString().slice(0, 400)));
            const result = new Array(positions.length).fill(undefined);
            let d = Buffer.concat(out);
            while (d.length) {
                const e = d.indexOf('\r\n\r\n');
                if (e < 0) break;
                const len = Number(/Content-Length:\s*(\d+)/i.exec(d.subarray(0, e).toString())[1]);
                const m = JSON.parse(d.subarray(e + 4, e + 4 + len));
                d = d.subarray(e + 4 + len);
                if (m.id >= 2) result[m.id - 2] = m.result || [];
            }
            resolve(result);
        });
        child.stdin.end(Buffer.concat(msgs.map(frame)));
    });
}
// the requests are dealt out to `jobs` server processes (each checks the document once, about half a
// minute interpreted); one process per core is about as fast as it gets.
async function references(server, text, positions, timeoutMs, jobs) {
    jobs = Math.max(1, Math.min(jobs, positions.length));
    const slices = Array.from({ length: jobs }, () => []);
    positions.forEach((p, k) => slices[k % jobs].push(k));
    const parts = await Promise.all(slices.map(ks => referencesRun(server, text, ks.map(k => positions[k]), timeoutMs)));
    const out = new Array(positions.length);
    slices.forEach((ks, j) => ks.forEach((k, n) => { out[k] = parts[j][n]; }));
    return out;
}

// ---- the analysis ----

// parts: [{name, text}] in order. returns {globals: [{name, part, line, kind, writes: {fn: n}, reads: {fn: n}}], notes}
async function analyze(compilerDir, server, timeoutMs, jobs) {
    const names = fs.readFileSync(path.join(compilerDir, 'parts'), 'utf8').split(/\r?\n/).filter(Boolean);
    const parts = names.map(name => ({ name, text: fs.readFileSync(path.join(compilerDir, name), 'utf8') }));
    const text = parts.map(p => p.text).join('');
    const toks = lex(text);
    const { functions, globals, problems } = structure(toks);
    // where each part starts, in lines of the joined text
    const starts = []; let acc = 0;
    for (const p of parts) { starts.push(acc); acc += p.text.split('\n').length - 1; }
    const partOf = line => { let k = starts.length - 1; while (k > 0 && starts[k] > line) k--; return k; };
    const byPos = new Map(toks.map((t, k) => [t.line + ':' + t.col, k]));
    const notes = [...problems.map(p => 'declaration not parsed: ' + p)];

    // duplicates would mean one name is declared twice; the second is ignored with a note
    const seen = new Map(), decls = [];
    for (const g of globals) {
        if (seen.has(g.name)) { notes.push('declared twice, the first is used: ' + g.name); continue; }
        seen.set(g.name, g); decls.push(g);
    }
    const answers = await references(server, text, decls.map(g => ({ line: g.tok.line, character: g.tok.col })), timeoutMs, jobs);

    // a function of the token index, by binary search over the sorted, non-overlapping ranges
    const fnAt = k => {
        let lo = 0, hi = functions.length - 1;
        while (lo <= hi) {
            const mid = (lo + hi) >> 1, f = functions[mid];
            if (k < f.startTok) hi = mid - 1; else if (k > f.endTok) lo = mid + 1; else return f.name;
        }
        return '(top level)';
    };
    const result = [];
    let unmapped = 0, noDecl = [];
    decls.forEach((g, gi) => {
        const locs = answers[gi];
        const entry = { name: g.name, part: parts[partOf(g.tok.line)].name, line: g.tok.line - starts[partOf(g.tok.line)] + 1, writes: {}, reads: {} };
        const declKey = g.tok.line + ':' + g.tok.col;
        let sawDecl = false;
        for (const loc of locs || []) {
            const key = loc.range.start.line + ':' + loc.range.start.character;
            if (key === declKey) { sawDecl = true; continue; }
            const k = byPos.get(key);
            if (k === undefined || toks[k].s !== g.name) { unmapped++; continue; }
            const into = classify(toks, k) === 'w' ? entry.writes : entry.reads;
            const f = fnAt(k);
            into[f] = (into[f] || 0) + 1;
        }
        if (!sawDecl) noDecl.push(g.name);
        // a function that writes a global is a writer; it is a reader only when it never writes
        for (const f of Object.keys(entry.writes)) delete entry.reads[f];
        result.push(entry);
    });
    if (noDecl.length) notes.push('the server did not return the declaration for ' + noDecl.length + ' names (not symbols it knows?): ' + noDecl.join(' '));
    if (unmapped) notes.push(unmapped + ' references did not land on an identifier of that name and were ignored');
    result.sort((a, b) => a.name < b.name ? -1 : a.name > b.name ? 1 : 0);
    return { globals: result, notes, functions: functions.length, tokens: toks.length };
}

// ---- output ----

const sortedKeys = o => Object.keys(o).sort();
function reportText(res) {
    const lines = [];
    for (const g of res.globals) {
        lines.push(g.name + '  (' + g.part + ':' + g.line + ')');
        lines.push('  writes: ' + (sortedKeys(g.writes).map(f => f + ' (' + g.writes[f] + ')').join(', ') || '-'));
        lines.push('  reads: ' + (sortedKeys(g.reads).map(f => f + ' (' + g.reads[f] + ')').join(', ') || '-'));
        lines.push('');
    }
    return lines.join('\n');
}
function reportJson(res) {
    const globals = {};
    for (const g of res.globals) {
        const w = {}, r = {};
        for (const f of sortedKeys(g.writes)) w[f] = g.writes[f];
        for (const f of sortedKeys(g.reads)) r[f] = g.reads[f];
        globals[g.name] = { part: g.part, line: g.line, writers: sortedKeys(g.writes), writes: w, readers: sortedKeys(g.reads), reads: r };
    }
    return JSON.stringify({ generated_by: 'src/tools/xref.js', globals }, null, 1) + '\n';
}
// the baseline: one global per line, the name then its sorted writer functions
function baselineText(res) {
    return res.globals.map(g => [g.name, ...sortedKeys(g.writes)].join(' ')).join('\n') + '\n';
}
function parseBaseline(text) {
    const m = new Map();
    for (const line of text.split(/\r?\n/)) {
        const w = line.trim().split(/\s+/).filter(Boolean);
        if (w.length) m.set(w[0], new Set(w.slice(1)));
    }
    return m;
}
// {added: [[global, fn]], removed: [[global, fn]], newGlobals, goneGlobals}
function compare(res, base) {
    const added = [], removed = [], newGlobals = [], goneGlobals = [];
    const now = new Set(res.globals.map(g => g.name));
    for (const g of res.globals) {
        const had = base.get(g.name);
        if (!had) { newGlobals.push(g.name); }
        for (const f of sortedKeys(g.writes)) if (!had || !had.has(f)) added.push([g.name, f]);
        if (had) for (const f of [...had].sort()) if (!g.writes[f]) removed.push([g.name, f]);
    }
    for (const name of base.keys()) if (!now.has(name)) goneGlobals.push(name);
    return { added, removed, newGlobals, goneGlobals };
}

// ---- command line ----

const USAGE = `usage: src/tools/xref.sh [options]
cross-reference of the compiler's global variables: for each top-level variable of src/compiler/parts,
the functions that write it and the functions that only read it. exact scoping comes from the language
server (textDocument/references), so a local that shadows a global is not counted. a write is the target
of = or a compound assignment, a ++ or --, an indexed or field store, or an address taken (&x).
options:
  (none)          write build/xref/report.txt and report.json
  --check         also compare the writers with tests/xref.baseline: exit 1 and print the difference when
                  a global gains a writer the baseline does not list (a new reader is fine; a removed
                  writer is only noted)
  --update        rewrite tests/xref.baseline from the current writers
  --server PATH   language server to use (default build/bin/zclsp; build it with make)
  --compiler-dir DIR   compiler parts directory (default src/compiler)
  --baseline FILE baseline file (default tests/xref.baseline)
  --out DIR       directory for the reports (default build/xref)
  --jobs N        language server processes to share the work (default: one per core, or XREF_JOBS)
  --help, -h      show this help and exit
exit status: 0 ok, 1 --check found a new writer, 2 bad usage, 3 the analysis failed.
the language server is interpreted on hosts without the arm64 jit, so a run takes tens of seconds.`;

async function main(argv) {
    const opt = { check: false, update: false, server: path.join(ROOT, 'build', 'bin', 'zclsp'), dir: path.join(ROOT, 'src', 'compiler'), baseline: path.join(ROOT, 'tests', 'xref.baseline'), out: path.join(ROOT, 'build', 'xref'), jobs: Number(process.env.XREF_JOBS) || Math.min(os.cpus().length, 8) };
    const needs = (k) => { if (k + 1 >= argv.length) { process.stderr.write('xref: ' + argv[k] + ' needs an argument; use --help for usage\n'); process.exit(2); } return argv[k + 1]; };
    for (let k = 0; k < argv.length; k++) {
        const a = argv[k];
        if (a === '--help' || a === '-h') { process.stdout.write(USAGE + '\n'); return 0; }
        else if (a === '--check') opt.check = true;
        else if (a === '--update') opt.update = true;
        else if (a === '--server') opt.server = path.resolve(needs(k++));
        else if (a === '--compiler-dir') opt.dir = path.resolve(needs(k++));
        else if (a === '--baseline') opt.baseline = path.resolve(needs(k++));
        else if (a === '--out') opt.out = path.resolve(needs(k++));
        else if (a === '--jobs') { opt.jobs = Number(needs(k++)); if (!(opt.jobs >= 1)) { process.stderr.write('xref: --jobs needs a positive number\n'); return 2; } }
        else { process.stderr.write('xref: unknown option ' + a + '; use --help for usage\n'); return 2; }
    }
    if (opt.check && opt.update) { process.stderr.write('xref: --check and --update exclude each other\n'); return 2; }
    let res;
    try {
        if (!fs.existsSync(opt.server)) throw new Error('no language server at ' + opt.server + ' (run make)');
        const scale = Number(process.env.ZC_TEST_TIMEOUT_SCALE) || 1;
        res = await analyze(opt.dir, opt.server, 3600000 * scale, opt.jobs);
    } catch (e) { process.stderr.write('xref: ' + e.message + '\n'); return 3; }
    fs.mkdirSync(opt.out, { recursive: true });
    fs.writeFileSync(path.join(opt.out, 'report.txt'), reportText(res));
    fs.writeFileSync(path.join(opt.out, 'report.json'), reportJson(res));
    for (const n of res.notes) process.stderr.write('xref: note: ' + n + '\n');
    const rel = p => path.relative(process.cwd(), p) || p;
    process.stdout.write('xref: ' + res.globals.length + ' globals, ' + res.functions + ' functions; wrote ' + rel(path.join(opt.out, 'report.txt')) + ' and report.json\n');
    if (opt.update) {
        fs.writeFileSync(opt.baseline, baselineText(res));
        process.stdout.write('xref: wrote ' + rel(opt.baseline) + '\n');
        return 0;
    }
    if (opt.check) {
        let base;
        try { base = parseBaseline(fs.readFileSync(opt.baseline, 'utf8')); } catch (e) { process.stderr.write('xref: cannot read the baseline ' + opt.baseline + ' (use --update)\n'); return 3; }
        const d = compare(res, base);
        for (const [g, f] of d.removed) process.stdout.write('note: ' + g + ' is no longer written by ' + f + '\n');
        for (const g of d.goneGlobals) process.stdout.write('note: global ' + g + ' is gone\n');
        if (d.added.length) {
            process.stdout.write('FAIL: new writers of compiler globals (a new reader is fine; review these, then run src/tools/xref.sh --update):\n');
            for (const [g, f] of d.added) process.stdout.write('  ' + g + ' is now written by ' + f + (d.newGlobals.includes(g) ? ' (new global)' : '') + '\n');
            return 1;
        }
        process.stdout.write('xref: check ok, no new writers against ' + rel(opt.baseline) + '\n');
    }
    return 0;
}

if (require.main === module) main(process.argv.slice(2)).then(c => { process.exitCode = c; });
else module.exports = { lex, structure, classify, analyze, main };