#!/usr/bin/env python3
"""Generate NEON-only RTCD devirtualization headers.

On RTC guaranteed-Neon builds (CONFIG_ARM_NEON_IS_GUARANTEED=1) every function
wired with SET_NEON(public, c, neon) -- i.e. no dotprod/sve/i8mm/neoverse
variant -- always resolves its runtime pointer to the _neon impl. Redirecting
the dispatch name straight at the _neon symbol turns the indirect
function-pointer branch into a direct, inlinable call.

This script emits two headers containing ONLY `#define public neon` lines:
  Source/Lib/Codec/aom_dsp_rtcd_neon_devirt.h
  Source/Lib/Codec/common_dsp_rtcd_neon_devirt.h
The _neon symbols are already declared in the RTCD headers (under
#ifdef ARCH_AARCH64), so no forward declarations are needed here. The RTCD
headers include these files under
  #if CONFIG_ARM_NEON_IS_GUARANTEED && !defined(<setup-TU guard>)
so the setup TU keeps the indirect form and stays compilable.

Usage:
  python3 Build/gen_devirt.py            # dry run: print counts + sample
  python3 Build/gen_devirt.py --write    # write populated devirt headers
  python3 Build/gen_devirt.py --empty    # write banner-only headers (devirt OFF)
"""
import os
import re
import sys

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
SRC = os.path.normpath(os.path.join(SCRIPT_DIR, "..", "Source", "Lib", "Codec"))

# Guard tokens that are FALSE (or non-aarch64) in the RTC build -> skip
# SET_NEON entries inside them (dead in the RTC build, zero benefit, added risk).
SKIP_GUARDS = {
    "CONFIG_ENABLE_HIGH_BIT_DEPTH", "CONFIG_ENABLE_OBMC", "CONFIG_ENABLE_FILM_GRAIN",
    "CONFIG_ENABLE_QUANT_MATRIX", "OPT_TUNE_VMAF", "ARCH_X86_64", "EN_AVX512_SUPPORT",
}

# (output header basename, source .c, RTCD setup-TU guard symbol)
TARGETS = [
    ("aom_dsp_rtcd_neon_devirt.h", "aom_dsp_rtcd.c", "AOM_RTCD_C"),
    ("common_dsp_rtcd_neon_devirt.h", "common_dsp_rtcd.c", "RTCD_C"),
]


def parse_set_neon(cfile):
    """Return list of (public, neon) for SET_NEON not inside any SKIP_GUARD region."""
    out = []
    stack = []  # one bool per open #if frame: True means region is skipped
    skipdepth = 0
    with open(cfile) as f:
        for line in f:
            s = line.strip()
            m = re.match(r"#\s*if(n?def)?\s+(.*)", s)
            if m:
                cond = m.group(2)
                toks = set(re.findall(r"[A-Za-z_][A-Za-z0-9_]+", cond))
                skip = bool(toks & SKIP_GUARDS)
                stack.append(skip)
                if skip:
                    skipdepth += 1
                continue
            if re.match(r"#\s*elif", s):
                if stack:
                    prev = stack[-1]
                    if prev:
                        skipdepth -= 1
                    toks = set(re.findall(r"[A-Za-z_][A-Za-z0-9_]+", s))
                    skip = bool(toks & SKIP_GUARDS)
                    stack[-1] = skip
                    if skip:
                        skipdepth += 1
                continue
            if re.match(r"#\s*else", s):
                if stack:
                    prev = stack[-1]
                    if prev:
                        skipdepth -= 1
                    stack[-1] = False
                continue
            if re.match(r"#\s*endif", s):
                if stack:
                    prev = stack.pop()
                    if prev:
                        skipdepth -= 1
                continue
            if skipdepth > 0:
                continue
            mm = re.match(
                r"SET_NEON\(\s*([A-Za-z0-9_]+)\s*,\s*[A-Za-z0-9_]+\s*,\s*([A-Za-z0-9_]+)\s*\)", s)
            if mm:
                out.append((mm.group(1), mm.group(2)))
    return out


def parse_dispatch_names(hfile):
    """Return set of dispatch names declared as RTCD_EXTERN ... (*name)(...)."""
    txt = open(hfile).read()
    return set(re.findall(r"RTCD_EXTERN\s+.*?\(\s*\*\s*([A-Za-z0-9_]+)\s*\)\s*\(", txt, re.DOTALL))


def banner(guard_symbol, populated):
    lines = [
        "// Auto-generated by Build/gen_devirt.py. DO NOT EDIT -- regenerate instead.",
        "//",
        "// NEON-only RTCD devirtualization redirects (RTC guaranteed-Neon build).",
        "// Each macro points a dispatch name at its _neon impl so callers emit a",
        "// direct, inlinable call instead of an indirect function-pointer branch.",
        "// The _neon symbols are already declared in the RTCD header (ARCH_AARCH64).",
        f"// Included under \"#if CONFIG_ARM_NEON_IS_GUARANTEED && !defined({guard_symbol})\".",
    ]
    if not populated:
        lines.append("//")
        lines.append("// (devirt DISABLED: no redirects emitted)")
    lines.append("")
    return lines


def build_header(pairs, dispatch_names, guard_symbol, populated):
    body = []
    emitted = set()
    skipped = []
    if populated:
        for pub, neon in pairs:
            if pub in emitted:
                continue
            if pub not in dispatch_names:
                skipped.append(pub)
                continue
            emitted.add(pub)
            body.append(f"#define {pub} {neon}")
    lines = ["// clang-format off"] + banner(guard_symbol, populated) + body + ["// clang-format on"]
    text = "\n".join(lines) + "\n"
    return text, len(emitted), skipped


def main():
    write = "--write" in sys.argv
    empty = "--empty" in sys.argv
    populated = not empty

    aom_dispatch = parse_dispatch_names(os.path.join(SRC, "aom_dsp_rtcd.h"))
    com_dispatch = parse_dispatch_names(os.path.join(SRC, "common_dsp_rtcd.h"))
    all_dispatch = aom_dispatch | com_dispatch

    total = 0
    for out_name, cfile, guard in TARGETS:
        pairs = parse_set_neon(os.path.join(SRC, cfile))
        text, n, skipped = build_header(pairs, all_dispatch, guard, populated)
        total += n
        print(f"{out_name}: active SET_NEON={len(pairs)} devirtualized={n} "
              f"skipped(no dispatch decl)={len(skipped)}")
        if skipped:
            print("  skipped:", skipped[:20])
        if write or empty:
            with open(os.path.join(SRC, out_name), "w") as fh:
                fh.write(text)
            print(f"  WROTE {os.path.join(SRC, out_name)} "
                  f"({'empty' if empty else 'populated'})")
    if write or empty:
        print(f"total devirtualized: {total}")
    else:
        print(f"total devirtualized (dry run): {total}")
        print("\n--- sample (aom, first 8 lines) ---")
        sample_pairs = parse_set_neon(os.path.join(SRC, "aom_dsp_rtcd.c"))
        text, _, _ = build_header(sample_pairs, all_dispatch, "AOM_RTCD_C", True)
        print("\n".join(text.splitlines()[:8]))


if __name__ == "__main__":
    main()
