#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# name_candidates.py - propose API-implementation names for the hand-written
# functions called by the generated CUDA API wrappers.
#
# Background: every exported cu* entry is a generated stub that jumps through
# an interposition slot to a generated wrapper (tracing/tool dispatch); the
# wrapper then calls the hand-written implementation of that API. Export
# families have several variants of one API (cuMemAlloc_v2, cuMemAlloc_ptsz,
# cuMemAlloc_ptds, ...), all of which must reach the same implementation.
#
# Method (robust, structure-only):
#   1. Group exported wrappers by BASE NAME: strip the variant suffixes
#      `_v<N>`, `_ptsz`, `_ptds` (in any order/combination) from the export
#      name; all exports in a group are variants of one API.
#   2. The implementation of a base name is a DIRECT CALLEE shared by ALL of
#      its variants' wrappers, i.e. the intersection of the wrappers' direct
#      callee sets.
#   3. Filter out global helpers: a candidate is dropped when it is called by
#      >= --max-outside wrappers OUTSIDE its group (trace/init/lock helpers
#      are called by dozens to hundreds of unrelated wrappers). Wrappers are
#      never candidates for themselves.
#   4. Confidence:
#        high   - exactly one candidate survives, and either the group has
#                 >= 2 variants (the survivor is shared by all of them) or the
#                 survivor is a unique callee (called by exactly one wrapper
#                 in the whole library);
#        medium - exactly one survivor but it is also called by another
#                 wrapper (possible shared implementation), or several
#                 survivors remain (the best-ranked one is proposed).
#      Ranking among several survivors: fewest outside wrappers first, then
#      largest direct-callee count (implementations do work; helpers do not),
#      then lowest address, for a deterministic result.
#
# Inputs:
#   <exports.tsv>    elf-index exports table: addr file_addr size name
#   <callgraph.tsv>  elf-index callgraph: caller callee caller_file callee_file
#   <names.tsv>      name map: file_vaddr<TAB>name (wrapper_<export> entries
#                    provide the wrapper address of each export)
#   <out.tsv>        output (comments + one row per proposal)
#
# Output rows (TSV):
#   file_vaddr  proposed_name  confidence  evidence  status
# `proposed_name` is `impl_<base>`; `evidence` carries the wrapper group,
# candidate callees before/after filtering and the outside-wrapper count;
# `status` is always `unverified` (the tool proposes, it does not prove).
# All addresses are FILE virtual addresses (libcuda+0x...).
#
# Usage:
#   name_candidates.py <exports.tsv> <callgraph.tsv> <names.tsv> <out.tsv>
#                      [--max-outside N] [--quiet]

import argparse
import re
import sys
from collections import defaultdict


SUFFIX_RE = re.compile(r"(_ptsz|_ptds)$|_v\d+$")


def base_name(name):
    """Strip CUDA variant suffixes (_v<N>, _ptsz, _ptds) from an export name."""
    prev = None
    while prev != name:
        prev = name
        name = SUFFIX_RE.sub("", name)
    return name


def load_exports(path):
    rows = []
    with open(path) as f:
        f.readline()
        for line in f:
            p = line.rstrip("\n").split("\t")
            if len(p) < 4:
                continue
            try:
                rows.append(dict(addr=int(p[0], 16), file=int(p[1], 16),
                                 size=int(p[2]), name=p[3]))
            except ValueError:
                continue
    return rows


def load_names(path):
    names = {}
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            p = line.split("\t")
            if len(p) >= 2:
                try:
                    names[int(p[0], 16)] = p[1]
                except ValueError:
                    pass
    return names


def load_callgraph(path):
    callees = defaultdict(set)
    callers = defaultdict(set)
    with open(path) as f:
        f.readline()
        for line in f:
            p = line.rstrip("\n").split("\t")
            if len(p) < 4:
                continue
            try:
                cf, kf = int(p[2], 16), int(p[3], 16)
            except ValueError:
                continue
            callees[cf].add(kf)
            callers[kf].add(cf)
    return callees, callers


def hx(a):
    return f"0x{a:x}"


def main():
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("exports")
    ap.add_argument("callgraph")
    ap.add_argument("names")
    ap.add_argument("out")
    ap.add_argument("--max-outside", type=int, default=2,
                    help="drop a candidate called by >= N wrappers outside "
                         "its base-name group (default 2)")
    ap.add_argument("--quiet", action="store_true")
    args = ap.parse_args()

    exports = load_exports(args.exports)
    names = load_names(args.names)
    callees, callers = load_callgraph(args.callgraph)

    wrapper_by_export = {}
    for a, n in names.items():
        if n.startswith("wrapper_"):
            wrapper_by_export[n[len("wrapper_"):]] = a
    wrapper_set = set(wrapper_by_export.values())

    groups = defaultdict(list)
    for e in exports:
        w = wrapper_by_export.get(e["name"])
        if w is not None:
            groups[base_name(e["name"])].append((e["name"], w))

    rows = []
    n_high = n_med = 0
    for group in sorted(groups):
        variants = groups[group]
        if not variants:
            continue
        variant_wrappers = {w for _, w in variants}
        inter = None
        for _, w in variants:
            cs = set(callees.get(w, ()))
            inter = cs if inter is None else (inter & cs)
        if not inter:
            continue
        before = sorted(inter)
        survivors = []
        for c in before:
            if c in wrapper_set:
                continue
            outside = [w for w in callers.get(c, ())
                       if w in wrapper_set and w not in variant_wrappers]
            if len(outside) >= args.max_outside:
                continue
            survivors.append((len(outside), c))
        if not survivors:
            continue
        # rank: fewest outside wrappers, then most direct callees, then address
        survivors.sort(key=lambda t: (t[0], -len(callees.get(t[1], ())), t[1]))
        best_outside, best = survivors[0]
        after = [c for _, c in survivors]
        if len(survivors) == 1 and (len(variants) >= 2 or
                                    len(callers.get(best, ()) & wrapper_set) == 1):
            confidence = "high"
            n_high += 1
        else:
            confidence = "medium"
            n_med += 1
        evidence = (
            f"group={group} variants=" + ",".join(n for n, _ in variants) +
            "; before=" + ",".join(hx(c) for c in before) +
            "; after=" + ",".join(hx(c) for c in after) +
            f"; outside_wrappers={best_outside}")
        rows.append((best, f"impl_{group}", confidence, evidence))

    rows.sort()
    with open(args.out, "w") as f:
        f.write("# Open-IOV CUDA API-implementation candidates (our own work product; tracked).\n")
        f.write("# Generated by tools/name_candidates.py from the elf-index exports,\n")
        f.write("# callgraph and the version's names.tsv; see that script for the method.\n")
        f.write("#\n")
        f.write("# file_vaddr<TAB>proposed_name<TAB>confidence<TAB>evidence<TAB>status\n")
        f.write("# confidence: high = exactly one filtered candidate, shared by >=2 variants\n")
        f.write("#             or a unique callee; medium = one survivor also called by another\n")
        f.write("#             wrapper, or several survivors (best-ranked proposed).\n")
        f.write("# status: unverified (tool output). Verification = read the body in the\n")
        f.write("#         decompiled corpus and mark verified/corrected here; only verified\n")
        f.write("#         names may be copied into names.tsv.\n")
        f.write("#\n")
        for a, name, conf, ev in rows:
            f.write(f"0x{a:08x}\t{name}\t{conf}\t{ev}\tunverified\n")

    if not args.quiet:
        print(f"[name_candidates] groups={len(groups)} proposals={len(rows)} "
              f"(high={n_high} medium={n_med}) -> {args.out}",
              file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())