#!/usr/bin/env python3
"""
hls_pragma_lint.py - check Xilinx Vitis HLS pragmas.

Why this exists
---------------
`#pragma HLS ...` is not C++. Stock clang and gcc do not know these pragmas
and simply ignore them, so a typo like

    #pragma HLS PIPLINE II=2

compiles perfectly happily. Vitis HLS may also accept it silently, meaning the
optimisation you asked for just never happens - which is a genuinely nasty way
to lose an afternoon.

This linter knows the pragma names, the options each one takes, and the values
those options allow. It reports problems in the standard

    file:line:col: error: message

form, so it works with nvim's errorformat, nvim-lint, efm-langserver, `:make`,
or any CI.

Usage
-----
    ./hls_pragma_lint.py file.cpp [more.cpp ...]
    ./hls_pragma_lint.py ../practical-3b/game_of_life_practical3b/gameoflife.cpp

Exit status is 0 when nothing is wrong, 1 when there is at least one error.
"""

import difflib
import re
import sys

# ---------------------------------------------------------------------------
# The pragma reference
# ---------------------------------------------------------------------------
# "FLAG"  -> a bare keyword with no value, e.g. `rewind`
# "NUM"   -> an integer
# set     -> a value that must be one of these
# "ANY"   -> anything (identifier, number, ...)
# ---------------------------------------------------------------------------
FLAG, NUM, ANY = "FLAG", "NUM", "ANY"

INTERFACE_MODES = {
    "ap_none", "ap_vld", "ap_ack", "ap_hs", "ap_ovld", "ap_fifo",
    "ap_memory", "ap_bram", "bram", "ap_ctrl_none", "ap_ctrl_hs",
    "ap_ctrl_chain", "m_axi", "s_axilite", "axis", "s_axis", "m_axis",
    "ap_stable", "ap_bus", "ap_array", "ap_array_rtl",
}

PRAGMAS = {
    "PIPELINE": {
        "options": {
            "II": NUM, "rewind": FLAG, "off": FLAG, "enable_flush": FLAG,
            "style": {"stp", "frp", "flp"},
        },
    },
    "UNROLL": {
        "options": {
            "factor": NUM, "region": FLAG, "skip_exit_check": FLAG, "off": FLAG,
        },
    },
    "ARRAY_PARTITION": {
        "options": {
            "variable": ANY, "type": {"block", "cyclic", "complete"},
            "factor": NUM, "dim": NUM, "recursive": FLAG, "off": FLAG,
        },
        "required": ["variable"],
    },
    "ARRAY_RESHAPE": {
        "options": {
            "variable": ANY, "type": {"block", "cyclic", "complete"},
            "factor": NUM, "dim": NUM, "off": FLAG,
        },
        "required": ["variable"],
    },
    "ARRAY_MAP": {
        "options": {
            "variable": ANY, "instance": ANY,
            "mode": {"horizontal", "vertical"}, "offset": NUM, "off": FLAG,
        },
        "required": ["variable"],
    },
    "DATA_PACK": {
        "options": {
            "variable": ANY, "instance": ANY, "offset": NUM, "off": FLAG,
        },
    },
    "AGGREGATE": {
        "options": {
            "variable": ANY, "compact": {"bit", "byte", "none", "auto"},
            "dim": NUM, "off": FLAG,
        },
        "required": ["variable"],
    },
    "STREAM": {
        "options": {
            "variable": ANY, "depth": NUM, "dim": NUM, "off": FLAG,
        },
        "required": ["variable"],
    },
    "INTERFACE": {
        "positional": INTERFACE_MODES,
        "options": {
            "mode": INTERFACE_MODES, "port": ANY, "bundle": ANY,
            "register": FLAG, "depth": NUM, "offset": ANY, "clock": ANY,
            "name": ANY, "num_read_outstanding": NUM,
            "num_write_outstanding": NUM, "max_read_burst_length": NUM,
            "max_write_burst_length": NUM, "latency": NUM,
            "m_axi_offset": ANY, "storage_type": ANY, "storage_impl": ANY,
            "sync": FLAG, "max_widen_bitwidth": NUM,
        },
        "required": ["port"],
    },
    "DATAFLOW": {
        "options": {"off": FLAG, "disable_start_propagation": FLAG},
    },
    "INLINE": {
        "options": {"off": FLAG, "recursive": FLAG, "region": FLAG},
    },
    "LATENCY": {
        "options": {"min": NUM, "max": NUM, "off": FLAG},
    },
    "DEPENDENCE": {
        "options": {
            "variable": ANY, "type": {"inter", "intra"},
            "direction": {"RAW", "WAR", "WAW"}, "distance": NUM, "off": FLAG,
        },
        "required": ["variable"],
    },
    "ALLOCATION": {
        "options": {
            "limit": NUM, "instances": ANY, "operation": ANY, "off": FLAG,
        },
    },
    "OCCURRENCE": {
        "options": {"cycle": NUM, "operation": ANY, "off": FLAG},
    },
    "LOOP_FLATTEN": {"options": {"off": FLAG}},
    "LOOP_MERGE": {"options": {"force": FLAG, "off": FLAG}},
    "LOOP_TRIPCOUNT": {
        "options": {"min": NUM, "max": NUM, "avg": NUM},
        "required": ["min"],
    },
    "FUNCTION_INSTANTIATE": {
        "options": {"variable": ANY}, "required": ["variable"],
    },
    "RESET": {"options": {"off": FLAG}},
    "PROTOCOL": {"positional": {"fixed", "float"}},
    "EXPRESSION_BALANCE": {"options": {"off": FLAG}},
    "REQUIREMENTS": {"options": {"min": NUM, "max": NUM}},
    "DISAGGREGATE": {
        "options": {"variable": ANY, "dim": NUM, "type": ANY},
        "required": ["variable"],
    },
    "BIND_OP": {
        "options": {
            "op": ANY, "impl": ANY, "variable": ANY, "latency": NUM,
        },
        "required": ["op", "impl"],
    },
    "BIND_STORAGE": {
        "options": {
            "type": ANY, "impl": ANY, "variable": ANY, "latency": NUM,
        },
        "required": ["type", "impl"],
    },
    "SHARED": {"options": {"off": FLAG}},
}

# Options whose value is a name declared in the code.
NAME_OPTIONS = ("variable", "port", "op")

PRAGMA_RE = re.compile(r"#\s*pragma\s+HLS\b(?P<rest>.*)$")
TOKEN_RE = re.compile(r"(?P<key>[A-Za-z_]\w*)\s*=\s*(?P<val>[^\s=]+)|(?P<bare>[^\s=]+)")


def find_pragma_name(rest):
    """Pull the pragma name off the front of a pragma's argument text."""
    m = re.match(r"\s*([A-Za-z_]\w*)", rest)
    if not m:
        return None, rest
    return m.group(1), rest[m.end():]


def parse_tokens(text):
    """Return (options, bare_words) where options is [(key, val, col), ...]."""
    options, bare = [], []
    for m in TOKEN_RE.finditer(text):
        if m.group("key"):
            options.append((m.group("key"), m.group("val"), m.start("key")))
        else:
            bare.append((m.group("bare"), m.start("bare")))
    return options, bare


def lint_file(path):
    """Return a list of (line, col, severity, message)."""
    out = []
    try:
        with open(path, "r", errors="replace") as fh:
            lines = fh.readlines()
    except OSError as exc:
        return [(1, 1, "error", "cannot read file: %s" % exc)]

    whole_file = "".join(lines)

    for lineno, line in enumerate(lines, start=1):
        m = PRAGMA_RE.search(line)
        if not m:
            continue
        rest = m.group("rest")

        name, args = find_pragma_name(rest)
        col = m.start() + m.group(0).index("HLS") + 1

        if name is None:
            out.append((lineno, col, "error",
                        "HLS pragma has no name (expected e.g. PIPELINE, UNROLL)"))
            continue

        if name not in PRAGMAS:
            hint = difflib.get_close_matches(name, PRAGMAS.keys(), n=1, cutoff=0.6)
            msg = "unknown HLS pragma '%s'" % name
            if hint:
                msg += " - did you mean '%s'?" % hint[0]
            else:
                msg += " (see the Vitis HLS pragma reference)"
            out.append((lineno, col, "error", msg))
            continue

        spec = PRAGMAS[name]
        known = spec.get("options", {})
        positional = spec.get("positional")

        options, bare = parse_tokens(args)
        seen = {k for k, _, _ in options}

        # Bare words: either a positional value or a flag.
        for word, off in bare:
            if word in known and known[word] is FLAG:
                seen.add(word)
                continue
            if positional and word in positional:
                continue
            if positional and word not in positional:
                out.append((lineno, col, "error",
                            "%s: '%s' is not a valid mode (expected one of: %s)"
                            % (name, word, ", ".join(sorted(positional)))))
                continue
            if word in known:
                out.append((lineno, col, "error",
                            "%s: '%s' needs a value, e.g. %s=..." % (name, word, word)))
                continue
            hint = difflib.get_close_matches(word, list(known) + list(positional or ()),
                                             n=1, cutoff=0.6)
            msg = "%s: unknown option '%s'" % (name, word)
            if hint:
                msg += " - did you mean '%s'?" % hint[0]
            out.append((lineno, col, "error", msg))

        # key=value options
        for key, val, off in options:
            if key not in known:
                hint = difflib.get_close_matches(key, known.keys(), n=1, cutoff=0.6)
                msg = "%s: unknown option '%s'" % (name, key)
                if hint:
                    msg += " - did you mean '%s'?" % hint[0]
                out.append((lineno, col, "error", msg))
                continue

            expect = known[key]
            if expect is FLAG:
                out.append((lineno, col, "warning",
                            "%s: '%s' is a flag and takes no value" % (name, key)))
            elif expect is NUM:
                if not re.fullmatch(r"\d+", val):
                    out.append((lineno, col, "error",
                                "%s: %s must be an integer, got '%s'"
                                % (name, key, val)))
            elif isinstance(expect, set):
                if val not in expect:
                    out.append((lineno, col, "error",
                                "%s: %s='%s' is invalid (expected one of: %s)"
                                % (name, key, val, ", ".join(sorted(expect)))))
            if key in NAME_OPTIONS and val != "return":
                # Count occurrences across the file. A real object is declared
                # at least once and named in at least one pragma, so a single
                # hit means the only mention is this pragma itself.
                hits = len(re.findall(r"\b%s\b" % re.escape(val), whole_file))
                if hits <= 1:
                    out.append((lineno, col, "warning",
                                "%s: %s='%s' is not declared anywhere in this file"
                                % (name, key, val)))

        # Required options
        for req in spec.get("required", []):
            if req not in seen and not (positional and req in positional):
                out.append((lineno, col, "error",
                            "%s: missing required option '%s='" % (name, req)))

    return out


def main(argv):
    if len(argv) < 2:
        print(__doc__.strip(), file=sys.stderr)
        return 2

    errors = 0
    for path in argv[1:]:
        for line, col, sev, msg in lint_file(path):
            print("%s:%d:%d: %s: %s" % (path, line, col, sev, msg))
            if sev == "error":
                errors += 1
    return 1 if errors else 0


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