#!/usr/bin/env node

/*
 * Copyright (c) 2026 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Inlinability census of a `.abc`: which methods the Panda AOT/JIT inliner can
// inline, and what stops the rest. The gates and the allow-list below are read
// from the compiler that ships with the pinned SDK
// (compiler/optimizer/optimizations/inlining.cpp and ir_builder.cpp) and from
// its option defaults (compiler/compiler.yaml); changing one without the other
// makes this census a lie, so keep the constants next to their names.
//
// Usage:
//   node scripts/inline-census.mjs <file.abc> [--json] [--filter <substr>] [--top N] [--keep] [--disasm <path>]

import { spawnSync } from "node:child_process";
import { createReadStream, existsSync, mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { basename, dirname, isAbsolute, join, resolve } from "node:path";
import { createInterface } from "node:readline";
import { fileURLToPath } from "node:url";

const REPO_ROOT = resolve(dirname(fileURLToPath(import.meta.url)), "..");

// compiler/compiler.yaml, as the shipped ark_aot reports them in --help
const ALWAYS_INLINE_BC_SIZE = 12; // Inlining::CheckMethodSize: past the inst and depth limits
const MAX_BC_SIZE = 200; // Inlining::CheckMethodSize: never inlined past this
const MAX_INSTS = 512; // compiler-inlining-max-insts: the caller's whole IR graph

// The host-tools directory is named for the RUNNING host *and* arch:
// `<platform>[_arm64]_host_tools`. This is the same derivation
// `scripts/panda-driver-alias.cjs` (hostToolsDir) and smart-arkts's
// getArchTools() use — `scripts/panda-layout.cjs` is where the string is meant
// to live once, and a platform-only map is wrong on linux/arm64 and darwin/x64,
// whose tools the SDK ships under a different name (or not at all).
function hostToolsDir() {
    const platform = { linux: "linux", darwin: "macos", win32: "windows" }[process.platform];
    if (platform === undefined) return undefined;
    const arch = process.arch === "arm64" ? "arm64" : "";
    return [platform, arch, "host_tools"].filter(Boolean).join("_");
}

// IrBuilderExternalInliningAnalysis::IsSuitableForInline (a cross-file callee).
// A method whose every instruction is accepted here can be inlined out of
// another .abc; one rejected instruction disqualifies the whole method.
const EXTERNAL_ACCEPTED = [
    /^(f?)mov/, /^(f?)lda/, /^(f?)sta/, /^return/, /^(f?)add/, /^(f?)sub/, /^(f?)mul/,
    /^(x?)or/, /^and/, /^(f?)neg/, /^not/, /^(a?)shr/, /^shl/, /^inc/, /^(u|f)?cmp/,
    /^[ifu][0-9]+to[ifu][0-9]+$/,
];
// ...unless the opcode is also one of these, which the same function refuses
const EXTERNAL_REJECTED = [
    /^(div|mod)/, /^call/, /(ld|st|len|new)(arr|obj)/, /throw/, /^initobj/, /^checkcast/,
    /^ldstatic/, /^ststatic/, /^lda\.(type|str|const)/, /^isinstance/, /^builtin/, /^monitor/, /launch/,
];

function externalEligibleInstruction(op) {
    if (/^j/.test(op)) return true; // Inlining treats every jump as suitable
    if (/^ldobj/.test(op) || /^stobj/.test(op)) return true;
    if (!EXTERNAL_ACCEPTED.some((re) => re.test(op))) return false;
    return !EXTERNAL_REJECTED.some((re) => re.test(op));
}

function findDisassembler(explicit) {
    // `ARKUI3_ARK_DISASM` is the override `arktsx/tools/api-index` already
    // honours, so a checkout with a relocated SDK resolves the same way for both
    // tools.
    const override = explicit ?? process.env.ARKUI3_ARK_DISASM;
    if (override !== undefined) {
        // This value reaches `spawnSync` as argv[0]. Accept only an absolute path
        // that exists, so a wrapper forwarding a variable into argv cannot turn
        // the census into an arbitrary-execution primitive.
        if (!isAbsolute(override)) {
            throw new Error(`--disasm must be an absolute path: ${override}`);
        }
        if (!existsSync(override)) {
            throw new Error(`--disasm not found: ${override}`);
        }
        return override;
    }
    const host = hostToolsDir();
    if (host === undefined) throw new Error(`unsupported platform: ${process.platform}`);
    const candidates = [
        join(REPO_ROOT, "node_modules/@panda/sdk", host, "bin/ark_disasm_static"),
        join(REPO_ROOT, "node_modules/@panda/sdk", host, "bin/ark_disasm"),
    ];
    for (const candidate of candidates) if (existsSync(candidate)) return candidate;
    throw new Error(`no disassembler under node_modules/@panda/sdk/${host}/bin — run npm install first`);
}

function parseArgs(argv) {
    const options = { file: undefined, json: false, keep: false, top: 0, filter: undefined, disasm: undefined };
    for (let i = 0; i < argv.length; i++) {
        const arg = argv[i];
        if (arg === "--json") options.json = true;
        else if (arg === "--keep") options.keep = true;
        else if (arg === "--top") options.top = Number(argv[++i] ?? 0);
        else if (arg === "--filter") options.filter = argv[++i];
        else if (arg === "--disasm") options.disasm = argv[++i];
        else if (arg.startsWith("-")) throw new Error(`unknown option: ${arg}`);
        else if (options.file === undefined) options.file = arg;
        else throw new Error(`unexpected argument: ${arg}`);
    }
    if (options.file === undefined) throw new Error("usage: inline-census.mjs <file.abc> [--json] [--filter s] [--top N] [--keep] [--disasm path]");
    return options;
}

// One `ark_disasm_static --verbose` run, parsed as it streams: the text form of
// an application .abc is tens of times the file, so it is never held whole.
function disassemble(disassembler, abcPath, outPath) {
    const result = spawnSync(disassembler, ["--verbose", abcPath, outPath], { stdio: ["ignore", "ignore", "pipe"] });
    if (result.status !== 0) {
        throw new Error(`disassembly failed (${result.status}): ${result.stderr?.toString().slice(0, 400) ?? ""}`);
    }
}

function emptyMethod() {
    return { bytes: 0, instructions: 0, tries: 0, calls: 0, anyCalls: 0, externalEligible: true, largest: [] };
}

function census(paPath, options) {
    return new Promise((resolveCensus, rejectCensus) => {
        const totals = {
            methods: 0,
            instructions: 0,
            bytes: 0,
            withTry: 0,
            withTryInstructions: 0,
            withCalls: 0,
            withAnyCalls: 0,
            externalEligible: 0,
            externalEligibleSmall: 0,
            buckets: { small: 0, near: 0, medium: 0, large: 0, over: 0 },
            bucketBytes: { small: 0, near: 0, medium: 0, large: 0, over: 0 },
        };
        const top = [];
        let current = null;

        const finish = () => {
            if (current === null) return;
            const method = current;
            current = null;
            totals.methods++;
            totals.instructions += method.instructions;
            totals.bytes += method.bytes;
            if (method.tries > 0) {
                totals.withTry++;
                totals.withTryInstructions += method.instructions;
            }
            if (method.calls > 0) totals.withCalls++;
            if (method.anyCalls > 0) totals.withAnyCalls++;
            if (method.externalEligible) {
                totals.externalEligible++;
                if (method.bytes <= ALWAYS_INLINE_BC_SIZE) totals.externalEligibleSmall++;
            }
            const bucket =
                method.bytes <= ALWAYS_INLINE_BC_SIZE ? "small"
                    : method.bytes <= 2 * ALWAYS_INLINE_BC_SIZE ? "near"
                        : method.bytes <= 45 ? "medium"
                            : method.bytes <= MAX_BC_SIZE ? "large" : "over";
            totals.buckets[bucket]++;
            totals.bucketBytes[bucket] += method.bytes;
            if (options.top > 0) {
                top.push({ name: method.name, bytes: method.bytes, instructions: method.instructions });
                top.sort((a, b) => b.bytes - a.bytes);
                if (top.length > options.top) top.length = options.top;
            }
        };

        const stream = createReadStream(paPath, { encoding: "utf8" });
        const lines = createInterface({ input: stream, crlfDelay: Infinity });
        lines.on("line", (line) => {
            if (line.startsWith(".function ")) {
                finish();
                const name = line.slice(10).match(/([^\s(]+)\s*\(/);
                current = emptyMethod();
                current.name = name === null ? line.slice(10).trim() : name[1];
                if (options.filter !== undefined && !current.name.includes(options.filter)) current = null;
                return;
            }
            if (current === null) return;
            if (line.startsWith(".catch")) {
                current.tries++;
                return;
            }
            if (line === "}") {
                finish();
                return;
            }
            if (!line.startsWith("\t")) return;
            const offset = line.match(/# offset: 0x([0-9a-f]+), /);
            if (offset === null) return;
            const op = line.slice(1).match(/^([^\s]+)/)?.[1];
            if (op === undefined) return;
            // The last bracketed group on the line is the encoding, one `0xNN` per byte.
            const encodings = line.match(/\[((?:0x[0-9a-f]{2} ?)+)\]/g);
            const width = encodings === null ? 0 : (encodings[encodings.length - 1].match(/0x[0-9a-f]{2}/g) ?? []).length;
            current.instructions++;
            current.bytes = Math.max(current.bytes, parseInt(offset[1], 16) + width);
            if (/^call/.test(op) || op === "initobj" || op === "initobj.short") current.calls++;
            if (/^any\./.test(op)) current.anyCalls++;
            if (!externalEligibleInstruction(op)) current.externalEligible = false;
        });
        lines.on("close", () => {
            finish();
            totals.top = top;
            resolveCensus(totals);
        });
        lines.on("error", rejectCensus);
        stream.on("error", rejectCensus);
    });
}

function percent(part, whole) {
    return whole === 0 ? "0.0%" : `${((100 * part) / whole).toFixed(1)}%`;
}

function report(totals, options) {
    const { methods, instructions, bytes } = totals;
    const lines = [];
    lines.push(`file                 : ${options.file}`);
    lines.push(`methods with code    : ${methods}`);
    lines.push(`instructions / bytes : ${instructions} / ${bytes} (${(bytes / Math.max(1, instructions)).toFixed(2)} B per instruction)`);
    lines.push("");
    lines.push("never inlinable at all:");
    lines.push(`  try table present    : ${totals.withTry} (${percent(totals.withTry, methods)}) — ${percent(totals.withTryInstructions, instructions)} of all instructions`);
    lines.push(`  bytecode > ${MAX_BC_SIZE} B      : ${totals.buckets.over} (${percent(totals.buckets.over, methods)})`);
    lines.push(`  any.* call site      : ${totals.withAnyCalls} (${percent(totals.withAnyCalls, methods)})`);
    lines.push("");
    lines.push("bytecode size (compiler-inlining-always-inline-bc-size = 12 B, max-bc-size = 200 B):");
    const bucketNames = {
        small: `<= ${ALWAYS_INLINE_BC_SIZE} B  (inlined past the inst and depth limits)`,
        near: `${ALWAYS_INLINE_BC_SIZE + 1}-${2 * ALWAYS_INLINE_BC_SIZE} B  (one peephole away from the cliff)`,
        medium: `${2 * ALWAYS_INLINE_BC_SIZE + 1}-45 B`,
        large: `46-${MAX_BC_SIZE} B`,
        over: `> ${MAX_BC_SIZE} B   (never inlined)`,
    };
    for (const key of ["small", "near", "medium", "large", "over"]) {
        const count = totals.buckets[key];
        lines.push(`  ${bucketNames[key].padEnd(52)} : ${String(count).padStart(6)} (${percent(count, methods).padStart(6)})  ${totals.bucketBytes[key]} B`);
    }
    lines.push("");
    lines.push(`callers of the ${MAX_INSTS}-instruction budget: a method whose own IR graph exceeds it can inline only the <= 12 B class.`);
    lines.push(`methods with calls   : ${totals.withCalls} (${percent(totals.withCalls, methods)})`);
    lines.push("");
    lines.push("cross-file (external) inlining, IrBuilderExternalInliningAnalysis allow-list:");
    lines.push(`  eligible methods     : ${totals.externalEligible} (${percent(totals.externalEligible, methods)})`);
    lines.push(`  of which <= ${ALWAYS_INLINE_BC_SIZE} B      : ${totals.externalEligibleSmall}`);
    if (totals.top.length > 0) {
        lines.push("");
        lines.push(`largest methods (top ${totals.top.length}):`);
        for (const method of totals.top) lines.push(`  ${String(method.bytes).padStart(7)} B  ${String(method.instructions).padStart(6)} inst  ${method.name}`);
    }
    return lines.join("\n");
}

async function main() {
    const options = parseArgs(process.argv.slice(2));
    options.file = resolve(options.file);
    if (!existsSync(options.file)) throw new Error(`no such file: ${options.file}`);
    const disassembler = findDisassembler(options.disasm);
    const scratch = mkdtempSync(join(tmpdir(), "arkui3-inline-census-"));
    const paPath = join(scratch, `${basename(options.file)}.pa`);
    try {
        disassemble(disassembler, options.file, paPath);
        const totals = await census(paPath, options);
        if (options.keep) process.stderr.write(`pandasm kept at ${paPath}\n`);
        process.stdout.write(options.json ? `${JSON.stringify({ ...totals, file: options.file }, null, 2)}\n` : `${report(totals, options)}\n`);
    } finally {
        if (!options.keep) rmSync(scratch, { recursive: true, force: true });
    }
}

main().catch((error) => {
    process.stderr.write(`${error.message}\n`);
    process.exit(1);
});