#!/usr/bin/env python3
"""Every model-checker harness this crate hosts runs in exactly one shard.

`.github/workflows/wai-formal-verification.yml` splits the harnesses into shards, each
a list of `--harness` filters (a harness runs when its full name contains a filter).
This finds every `#[kani::proof]` function in the sources `src/lib.rs` includes by
`#[path]` (with the modules those files declare in turn), gives each its full name,
and fails unless each name is matched by the filters of exactly one shard, and each
filter matches at least one harness. A harness no shard runs would be compiled and
never checked; one in two shards would be checked twice. A harness gated by
`cfg(kani_open)` (one whose window the checker did not finish within a shard's
budget) is compiled only on request, so no shard runs it; it is counted apart.

Usage: python3 -I check_shards.py <repository root>
Standard library only.
"""

import os
import re
import sys

PATH_MOD = re.compile(r'#\[path = "([^"]+)"\]\s*(?:pub(?:\(crate\))? )?mod (\w+);')
PLAIN_MOD = re.compile(r'^(?:pub(?:\(crate\))? )?mod (\w+);', re.M)
PROOF = re.compile(r'(#\[cfg\(kani_open\)\]\s*)?#\[kani::proof\]\s*(?:#\[[^\]]*\]\s*)*(?:pub(?:\(crate\))? )?fn (\w+)')
# A matrix line names its shard first and its filters last; between them it may give
# other plain fields, such as the runner (`os: macos`).
SHARD = re.compile(r'^\s*- \{ shard: ([\w-]+),(?: *[\w-]+: [\w-]+,)* *harness: "([^"]*)"', re.M)
MATRIX_LINE = re.compile(r'^\s*- \{ shard: ', re.M)


OPEN = []


def harnesses(path, prefix, out):
    """The harnesses of the module in `path` and of the modules it declares. A harness
    gated by `cfg(kani_open)` is compiled only on request and runs in no shard; it is
    listed apart."""
    src = open(path, encoding="utf-8").read()
    for gated, name in PROOF.findall(src):
        (OPEN if gated else out).append(prefix + name)
    base = os.path.dirname(path)
    stem = os.path.splitext(os.path.basename(path))[0]
    # A `mod x;` in a.rs is a/x.rs or a/x/mod.rs; in mod.rs or lib.rs, x.rs beside it.
    here = base if stem in ("mod", "lib") else os.path.join(base, stem)
    for name in PLAIN_MOD.findall(src):
        for cand in (os.path.join(here, name + ".rs"), os.path.join(here, name, "mod.rs")):
            if os.path.exists(cand):
                harnesses(cand, prefix + name + "::", out)
                break


def main(root):
    lib = os.path.join(root, "wai", "wai-verify", "src", "lib.rs")
    src = open(lib, encoding="utf-8").read()
    names = []
    for rel, name in PATH_MOD.findall(src):
        harnesses(os.path.normpath(os.path.join(os.path.dirname(lib), rel)), name + "::", names)
    wf = open(os.path.join(root, ".github", "workflows", "wai-formal-verification.yml"), encoding="utf-8").read()
    shards = [(s, f.split()) for s, f in SHARD.findall(wf)]
    failures = 0
    if not names or not shards:
        print("FAIL no harnesses or no shards found")
        return 1
    # A shard line this cannot read would drop its filters unseen.
    if len(shards) != len(MATRIX_LINE.findall(wf)):
        print("FAIL %d shard lines, %d read" % (len(MATRIX_LINE.findall(wf)), len(shards)))
        return 1
    for n in names:
        hit = [s for s, fs in shards if any(f in n for f in fs)]
        if len(hit) != 1:
            print("FAIL harness %s is in %d shards: %s" % (n, len(hit), ", ".join(hit) or "none"))
            failures += 1
    for s, fs in shards:
        for f in fs:
            if not any(f in n for n in names):
                print("FAIL shard %s: the filter %s matches no harness" % (s, f))
                failures += 1
    print("%d harnesses in %d shards; %d gated by cfg(kani_open), in none" % (len(names), len(shards), len(OPEN)))
    for s, fs in shards:
        print("  %-14s %d" % (s, sum(1 for n in names if any(f in n for f in fs))))
    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1] if len(sys.argv) > 1 else "."))