#!/usr/bin/env python3
"""Compare the code coverage of two PostgreSQL tests, per function.

usage: covdiff.py TEST_1 TEST_2 [-n RUNS] [--top N] [--format md|text]

--format md (default) prints a Markdown table.  --format text prints an
aligned plain text table that fits in 72 columns, for email.

Run inside a source tree built with -Db_coverage=true.  Each test runs
RUNS times from zeroed counters.  A line counts for a test only if every
run of that test hit it, and against the other test only if no run of
the other test hit it, so lines that come and go between runs are
ignored.  Only files in the source tree count: generated files under
build/ and system headers are left out.  Do not run other tests on the
same build at the same time, their coverage would be mixed in.
"""
import argparse, json, os, re, subprocess, sys, tempfile

# lcov 2.x stops on these two kinds of problems, which are harmless here:
#   range: files like copyfuncs.c #include generated code, so some line
#          numbers are past the end of the file
#   inconsistent: clang marks some lines as run without any branch data
LCOV_TOLERATE = "inconsistent,range"

# the program meson runs for a test tells what kind of test it is
TEST_KIND = {"perl": "perl", "pytest": "python", "pg_regress": "regress",
             "pg_isolation_regress": "isolation"}


def read_info(path):
    """Return (lines hit, function line ranges) from an lcov file."""
    lines, ranges = set(), {}
    src, fnl = None, {}
    for line in open(path):
        tag, _, rest = line.strip().partition(":")
        if tag == "SF":  # source file
            src, fnl = rest, {}
        elif tag == "DA":  # line,count
            n, count = rest.split(",")[:2]
            if int(count) > 0:
                lines.add((src, int(n)))
        elif tag == "FNL":  # function index,start line[,end line]
            i, start, *end = rest.split(",")
            fnl[i] = (int(start), int(end[0]) if end else int(start))
        elif tag == "FNA":  # function index,count,name
            i, _, name = rest.split(",", 2)
            if i in fnl:
                ranges[(src, name)] = fnl[i]
    return lines, ranges


def zero_counters(build):
    # os.walk, not glob: glob skips names starting with a dot, and objects
    # built from another directory get .gcda names like .._timezone_pgtz.c.gcda
    for d, _, files in os.walk(build):
        for f in files:
            if f.endswith(".gcda"):
                os.remove(os.path.join(d, f))


def run(build, test, runs):
    """Return (lines hit in every run, lines hit in any run, function ranges)."""
    every, anyrun, ranges = None, set(), {}
    for _ in range(runs):
        zero_counters(build)
        r = subprocess.run(["meson", "test", "-C", build, "--no-rebuild", test],
                           capture_output=True, text=True)
        if r.returncode != 0:
            sys.exit(f"{test} failed\n{r.stdout[-2000:]}")
        # a skipped test also exits 0, but covers almost nothing
        ok = re.search(r"^Ok:\s+(\d+)", r.stdout, re.M)
        if not ok or int(ok.group(1)) != 1:
            sys.exit(f"{test} did not run (skipped, or not a single test)\n{r.stdout[-2000:]}")
        with tempfile.NamedTemporaryFile(suffix=".info") as info:
            subprocess.run(["lcov", "--capture", "-d", build, "--base-directory", build,
                            "--ignore-errors", LCOV_TOLERATE, "--quiet", "-o", info.name],
                           capture_output=True, check=True)
            lines, rg = read_info(info.name)
        every = lines if every is None else every & lines
        anyrun |= lines
        ranges.update(rg)
    return every, anyrun, ranges


def kind_of(build, test):
    tests = json.loads(subprocess.run(["meson", "introspect", build, "--tests"],
                                      capture_output=True, text=True).stdout)
    for t in tests:
        if t["name"] == test:
            cmd = t["cmd"][t["cmd"].index("--") + 1:] if "--" in t["cmd"] else t["cmd"]
            program = os.path.basename(cmd[0])
            return TEST_KIND.get(program, program)
    sys.exit(f"no such test: {test}")


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("test_1")
    ap.add_argument("test_2")
    ap.add_argument("-n", "--runs", type=int, default=3)
    ap.add_argument("--top", type=int, default=30, help="functions to show in the table, 0 for all")
    ap.add_argument("--format", choices=["md", "text"], default="md",
                    help="md: Markdown table (default), text: plain text for email")
    args = ap.parse_args()
    if args.runs < 1:
        sys.exit("--runs must be at least 1")

    root = subprocess.run(["git", "rev-parse", "--show-toplevel"],
                          capture_output=True, text=True).stdout.strip()
    build = os.path.join(root, "build")
    if not root or not os.path.exists(os.path.join(build, "build.ninja")):
        sys.exit("run this inside a source tree with a coverage build in build/")

    kind_1, kind_2 = kind_of(build, args.test_1), kind_of(build, args.test_2)
    # short column labels: the test kinds, or 1 and 2 if both are the same kind
    label_1, label_2 = (kind_1, kind_2) if kind_1 != kind_2 else ("1", "2")
    every_1, any_1, ranges = run(build, args.test_1, args.runs)
    every_2, any_2, ranges_2 = run(build, args.test_2, args.runs)
    ranges.update(ranges_2)

    # keep only files in the source tree: no generated files under build/,
    # no system headers like the SDK's math.h
    def in_tree(src):
        rel = os.path.relpath(src, root)
        return not rel.startswith(("..", "build/"))
    source = lambda lines: {(src, n) for src, n in lines if in_tree(src)}
    every_1, any_1, every_2, any_2 = map(source, (every_1, any_1, every_2, any_2))

    # count the differing lines per function
    functions_in = {}
    for (src, name), (start, end) in ranges.items():
        functions_in.setdefault(src, []).append((start, end, name))
    rows = {}
    for which, diff in ((1, every_1 - any_2), (2, every_2 - any_1)):
        for src, n in diff:
            name = next((f for s, e, f in functions_in.get(src, ()) if s <= n <= e), "(no function)")
            rows.setdefault((src, name), {1: 0, 2: 0})[which] += 1

    ordered = sorted(rows.items(), key=lambda kv: -(kv[1][1] + kv[1][2]))
    shown = ordered if args.top == 0 else ordered[:args.top]
    more = ", use --top N for more, --top 0 for all." if len(shown) < len(ordered) else "."
    total_1, total_2 = sum(c[1] for c in rows.values()), sum(c[2] for c in rows.values())
    if args.format == "md":
        print_md(args, kind_1, kind_2, shown, ordered, more, root,
                 len(every_1), len(every_2), total_1, total_2)
    else:
        print_text(args, label_1, label_2, shown, ordered, more, root,
                   len(every_1), len(every_2), total_1, total_2)


def print_md(args, kind_1, kind_2, shown, ordered, more, root, hit_1, hit_2, only_1, only_2):
    name_1, name_2 = f"{args.test_1} ({kind_1})", f"{args.test_2} ({kind_2})"
    print(f"Lines only in one test, per function, largest first. "
          f"Showing {len(shown)} of {len(ordered)} functions{more}")
    print()
    print(f"| function | file | lines only in {name_1} | lines only in {name_2} |")
    print("|---|---|---:|---:|")
    for (src, name), count in shown:
        print(f"| {name} | {os.path.relpath(src, root)} | {count[1] or ''} | {count[2] or ''} |")
    print()
    print(f"- Lines hit by {name_1}: {hit_1}")
    print(f"- Lines hit by {name_2}: {hit_2}")
    print(f"- Lines only in {name_1}: {only_1}")
    print(f"- Lines only in {name_2}: {only_2}")


def print_text(args, label_1, label_2, shown, ordered, more, root, hit_1, hit_2, only_1, only_2):
    w1, w2 = max(len(label_1), 6), max(len(label_2), 6)
    print(f"{label_1} = {args.test_1}")
    print(f"{label_2} = {args.test_2}")
    print()
    print("Lines only in one test, per function, largest first.")
    print(f"Showing {len(shown)} of {len(ordered)} functions{more}")
    print()
    print(f"{label_1:>{w1}}  {label_2:>{w2}}  file:function")
    for (src, name), count in shown:
        where = "/".join(os.path.relpath(src, root).split("/")[-2:])
        print(f"{count[1] or '':>{w1}}  {count[2] or '':>{w2}}  {where}:{name}")
    print()
    print(f"Lines hit:        {label_1} {hit_1}, {label_2} {hit_2}")
    print(f"Lines only in it: {label_1} {only_1}, {label_2} {only_2}")


if __name__ == "__main__":
    main()