#!/usr/bin/env python3
"""Inventory of whole-file `&[u8]` parameters in the clawhdf5 workspace.

Used by docs/design/range-reads.md.  Run from the repository root:

    python3 docs/design/tools/inventory.py            # per-file table
    python3 docs/design/tools/inventory.py --list     # every signature
    python3 docs/design/tools/inventory.py --patterns # read patterns per crate

A function counts as taking "the whole file" when it has a parameter named
`file_data: &[u8]` (the repository convention), or a `data` / `file` / `buf` /
`bytes` / `mmap` parameter of type `&[u8]` *together with* a parameter whose
name says it is a file address (`*address*`, `*addr*`, `*offset*` of an
integer type).  The second rule is a heuristic; `--list` prints every match so
it can be checked by eye.  Code after the first `#[cfg(test)] mod ... {` in a
file is excluded (the repository keeps unit tests at the end of each file),
as are the tests/ and benches/ directories.
"""
import os
import re
import sys
from collections import defaultdict

ROOT = os.getcwd()
FN_RE = re.compile(r"\bfn\s+([A-Za-z_][A-Za-z0-9_]*)\s*(<[^()]*?>)?\s*\(", re.S)
PARAM_RE = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)\s*:\s*&(?:'[a-z_]+\s+)?\[u8\]")
ADDR_RE = re.compile(r"\b([a-z_]*(?:address|addr|offset)[a-z_]*)\s*:\s*(?:u64|usize|u32)")
WHOLE_NAMES = {"data", "file", "buf", "bytes", "file_bytes", "mmap"}


def signature(src, start):
    depth, i = 0, start
    while i < len(src):
        c = src[i]
        if c == "(":
            depth += 1
        elif c == ")":
            depth -= 1
            if depth == 0:
                return src[start + 1 : i]
        i += 1
    return ""


def non_test_source(src):
    m = re.search(r"#\[cfg\(test\)\]\s*mod\s+\w+\s*\{", src)
    return src[: m.start()] if m else src


PATTERNS = [
    ("`file_data` passed on (call sites)", re.compile(r"[(,]\s*&?file_data\s*[,)]")),
    ("`file_data[..]` slicing", re.compile(r"\bfile_data\s*\[")),
    ("open-ended `file_data[x..]`", re.compile(r"\bfile_data\s*\[[^\]]*\.\.\s*\]")),
    ("`file_data.get(..)`", re.compile(r"\bfile_data\s*\.\s*get\s*\(")),
    ("`file_data.len()`", re.compile(r"\bfile_data\s*\.\s*len\s*\(\)")),
    ("address/offset `as usize` casts", re.compile(r"\b[a-z_]*(?:addr|address|offset)[a-z_]*\s+as\s+usize")),
    ("`ObjectHeader::parse(` calls", re.compile(r"ObjectHeader::parse\s*\(")),
    ("`.as_bytes()` on a file/reader", re.compile(r"\b(?:file|reader|data|self\.file|self\.data|self\.file\.data|root\.file)\s*\.\s*as_bytes\s*\(\)")),
]


def patterns():
    """Per-crate counts of the read patterns (non-test code only)."""
    per_crate = defaultdict(lambda: [0] * len(PATTERNS))
    for crate in sorted(os.listdir(os.path.join(ROOT, "crates"))):
        srcdir = os.path.join(ROOT, "crates", crate, "src")
        for dp, _, fns in os.walk(srcdir):
            for fn in fns:
                if not fn.endswith(".rs"):
                    continue
                with open(os.path.join(dp, fn), encoding="utf-8") as fh:
                    src = non_test_source(fh.read())
                for i, (_, rx) in enumerate(PATTERNS):
                    per_crate[crate][i] += len(rx.findall(src))
    print("| crate | " + " | ".join(n for n, _ in PATTERNS) + " |")
    print("|---|" + "---:|" * len(PATTERNS))
    tot = [0] * len(PATTERNS)
    for c, v in sorted(per_crate.items()):
        if any(v):
            print("| %s | %s |" % (c, " | ".join(str(x) for x in v)))
            tot = [a + b for a, b in zip(tot, v)]
    print("| **total** | %s |" % " | ".join("**%d**" % x for x in tot))


def main():
    if "--patterns" in sys.argv:
        patterns()
        return
    per_file = defaultdict(lambda: [0, 0, 0])  # [file_data, heuristic, pub]
    rows = []
    for crate in sorted(os.listdir(os.path.join(ROOT, "crates"))):
        srcdir = os.path.join(ROOT, "crates", crate, "src")
        for dp, _, fns in os.walk(srcdir):
            for fn in sorted(fns):
                if not fn.endswith(".rs"):
                    continue
                path = os.path.join(dp, fn)
                with open(path, encoding="utf-8") as fh:
                    src = non_test_source(fh.read())
                for m in FN_RE.finditer(src):
                    sig = signature(src, m.end() - 1)
                    params = PARAM_RE.findall(sig)
                    kind = None
                    if "file_data" in params:
                        kind = "file_data"
                    elif any(p in WHOLE_NAMES for p in params) and ADDR_RE.search(sig):
                        kind = "heuristic"
                    if not kind:
                        continue
                    rel = os.path.relpath(path, ROOT)
                    line_start = src.rfind("\n", 0, m.start()) + 1
                    is_pub = src[line_start : m.start()].strip().startswith("pub")
                    per_file[rel][0 if kind == "file_data" else 1] += 1
                    per_file[rel][2] += int(is_pub)
                    line = src.count("\n", 0, m.start()) + 1
                    rows.append((rel, line, m.group(1), kind, is_pub))
    if "--list" in sys.argv:
        for r in rows:
            print("%s:%d %s [%s%s]" % (r[0], r[1], r[2], r[3], ", pub" if r[4] else ""))
        return
    print("| file | `file_data` fns | other whole-slice fns (heuristic) | of which `pub` |")
    print("|---|---:|---:|---:|")
    tot = [0, 0, 0]
    for f, (a, b, p) in sorted(per_file.items(), key=lambda kv: (-(kv[1][0] + kv[1][1]), kv[0])):
        print("| %s | %d | %d | %d |" % (f, a, b, p))
        tot = [tot[0] + a, tot[1] + b, tot[2] + p]
    print("| **total** | **%d** | **%d** | **%d** |" % tuple(tot))


if __name__ == "__main__":
    main()