#!/usr/bin/env python3
# -*- coding: UTF-8 -*-

# Polly/LLVM update_check.py
# Update lit FileCheck files by replacing the 'CHECK:' lines by the actual output of the 'RUN:' command.

import argparse
import os
import subprocess
import shlex
import re


polly_src_dir = """@POLLY_SOURCE_DIR@"""
polly_lib_dir = """@POLLY_LIB_DIR@"""
shlibext = """@LLVM_SHLIBEXT@"""
llvm_tools_dir = """@LLVM_TOOLS_DIR@"""
llvm_polly_link_into_tools = not """@LLVM_POLLY_LINK_INTO_TOOLS@""".lower() in {
    "",
    "0",
    "n",
    "no",
    "off",
    "false",
    "notfound",
    "llvm_polly_link_into_tools-notfound",
}

runre = re.compile(r"\s*\;\s*RUN\s*\:(?P<tool>.*)")
filecheckre = re.compile(r"\s*(?P<tool>.*)\|\s*(?P<filecheck>FileCheck\s[^|]*)")
emptyline = re.compile(r"\s*(\;\s*)?")
commentline = re.compile(r"\s*(\;.*)?")


def ltrim_emptylines(lines, meta=None):
    while len(lines) and emptyline.fullmatch(lines[0]):
        del lines[0]
        if meta is not None:
            del meta[0]


def rtrim_emptylines(lines):
    while len(lines) and emptyline.fullmatch(lines[-1]):
        del lines[-1]


def trim_emptylines(lines):
    ltrim_emptylines(lines)
    rtrim_emptylines(lines)


def complete_exename(path, filename):
    complpath = os.path.join(path, filename)
    if os.path.isfile(complpath):
        return complpath
    elif os.path.isfile(complpath + ".exe"):
        return complpath + ".exe"
    return filename


def indention(line):
    for i, c in enumerate(line):
        if c != " " and c != "\t":
            return i
    return None


def common_indent(lines):
    indentions = (indention(line) for line in lines)
    indentions = (indent for indent in indentions if indent is not None)
    return min(indentions, default=0)


funcre = re.compile(r"^    Function: \S*$")
regionre = re.compile(r"^    Region: \S*$")
depthre = re.compile(r"^    Max Loop Depth: .*")
paramre = re.compile(r"    [0-9a-z-A-Z_]+\: .*")


def classyfier1(lines):
    i = iter(lines)
    line = i.__next__()
    while True:
        if line.startswith(
            "Printing analysis 'Polly - Calculate dependences' for region: "
        ):
            yield {"PrintingDependenceInfo"}
        elif line.startswith("remark: "):
            yield {"Remark"}
        elif funcre.fullmatch(line):
            yield {"Function"}
        elif regionre.fullmatch(line):
            yield {"Region"}
        elif depthre.fullmatch(line):
            yield {"MaxLoopDepth"}
        elif line == "    Invariant Accesses: {":
            while True:
                yield {"InvariantAccesses"}
                if line == "    }":
                    break
                line = i.__next__()
        elif line == "    Context:":
            yield {"Context"}
            line = i.__next__()
            yield {"Context"}
        elif line == "    Assumed Context:":
            yield {"AssumedContext"}
            line = i.__next__()
            yield {"AssumedContext"}
        elif line == "    Invalid Context:":
            yield {"InvalidContext"}
            line = i.__next__()
            yield {"InvalidContext"}
        elif line == "    Boundary Context:":
            yield {"BoundaryContext"}
            line = i.__next__()
            yield {"BoundaryContext"}
            line = i.__next__()
            while paramre.fullmatch(line):
                yield {"Param"}
                line = i.__next__()
            continue
        elif line == "    Arrays {":
            while True:
                yield {"Arrays"}
                if line == "    }":
                    break
                line = i.__next__()
        elif line == "    Arrays (Bounds as pw_affs) {":
            while True:
                yield {"PwAffArrays"}
                if line == "    }":
                    break
                line = i.__next__()
        elif line.startswith("    Alias Groups ("):
            while True:
                yield {"AliasGroups"}
                line = i.__next__()
                if not line.startswith("        "):
                    break
            continue
        elif line == "    Statements {":
            while True:
                yield {"Statements"}
                if line == "    }":
                    break
                line = i.__next__()
        elif line == "    RAW dependences:":
            yield {"RAWDep", "BasicDep", "Dep", "DepInfo"}
            line = i.__next__()
            while line.startswith("        "):
                yield {"RAWDep", "BasicDep", "Dep", "DepInfo"}
                line = i.__next__()
            continue
        elif line == "    WAR dependences:":
            yield {"WARDep", "BasicDep", "Dep", "DepInfo"}
            line = i.__next__()
            while line.startswith("        "):
                yield {"WARDep", "BasicDep", "Dep", "DepInfo"}
                line = i.__next__()
            continue
        elif line == "    WAW dependences:":
            yield {"WAWDep", "BasicDep", "Dep", "DepInfo"}
            line = i.__next__()
            while line.startswith("        "):
                yield {"WAWDep", "BasicDep", "Dep", "DepInfo"}
                line = i.__next__()
            continue
        elif line == "    Reduction dependences:":
            yield {"RedDep", "Dep", "DepInfo"}
            line = i.__next__()
            while line.startswith("        "):
                yield {"RedDep", "Dep", "DepInfo"}
                line = i.__next__()
            continue
        elif line == "    Transitive closure of reduction dependences:":
            yield {"TransitiveClosureDep", "DepInfo"}
            line = i.__next__()
            while line.startswith("        "):
                yield {"TransitiveClosureDep", "DepInfo"}
                line = i.__next__()
            continue
        elif line.startswith("New access function '"):
            yield {"NewAccessFunction"}
        elif line == "Schedule before flattening {":
            while True:
                yield {"ScheduleBeforeFlattening"}
                if line == "}":
                    break
                line = i.__next__()
        elif line == "Schedule after flattening {":
            while True:
                yield {"ScheduleAfterFlattening"}
                if line == "}":
                    break
                line = i.__next__()
        else:
            yield set()
        line = i.__next__()


def classyfier2(lines):
    i = iter(lines)
    line = i.__next__()
    while True:
        if funcre.fullmatch(line):
            while line.startswith("    "):
                yield {"FunctionDetail"}
                line = i.__next__()
            continue
        elif line.startswith(
            "Printing analysis 'Polly - Generate an AST from the SCoP (isl)' for region: "
        ):
            yield {"PrintingIslAst"}
            line = i.__next__()
            while not line.startswith("Printing analysis"):
                yield {"AstDetail"}
                line = i.__next__()
            continue
        else:
            yield set()
        line = i.__next__()


replrepl = {"{{": "{{[{][{]}}", "}}": "{{[}][}]}}", "[[": "{{\[\[}}", "]]": "{{\]\]}}"}
replre = re.compile("|".join(re.escape(k) for k in replrepl.keys()))


def main():
    parser = argparse.ArgumentParser(description="Update CHECK lines")
    parser.add_argument(
        "testfile", help="File to update (absolute or relative to --testdir)"
    )
    parser.add_argument(
        "--check-style",
        choices=["CHECK", "CHECK-NEXT"],
        default="CHECK-NEXT",
        help="What kind of checks lines to generate",
    )
    parser.add_argument(
        "--check-position",
        choices=["end", "before-content", "autodetect"],
        default="autodetect",
        help="Where to add the CHECK lines into the file; 'autodetect' searches for the first 'CHECK' line ind inserts it there",
    )
    parser.add_argument(
        "--check-include",
        action="append",
        default=[],
        help="What parts of the output lines to check; use syntax 'CHECK=include' to apply to one CHECK-prefix only (by default, everything)",
    )
    parser.add_argument(
        "--check-label-include",
        action="append",
        default=[],
        help="Use CHECK-LABEL for these includes",
    )
    parser.add_argument(
        "--check-part-newline",
        action="store_true",
        help="Add empty line between different check parts",
    )
    parser.add_argument(
        "--prefix-only",
        action="append",
        default=None,
        help="Update only these prefixes (default: all)",
    )
    parser.add_argument("--bindir", help="Location of the opt program")
    parser.add_argument("--testdir", help="Root dir for unit tests")
    parser.add_argument(
        "--inplace", "-i", action="store_true", help="Replace input file"
    )
    parser.add_argument("--output", "-o", help="Write changed input to this file")
    known = parser.parse_args()

    if not known.inplace and known.output is None:
        print("Must specify what to do with output (--output or --inplace)")
        exit(1)
    if known.inplace and known.output is not None:
        print("--inplace and --output are mutually exclusive")
        exit(1)

    outfile = known.output

    filecheckparser = argparse.ArgumentParser(add_help=False)
    filecheckparser.add_argument("-check-prefix", "--check-prefix", default="CHECK")

    filename = known.testfile
    for dir in [".", known.testdir, os.path.join(polly_src_dir, "test"), polly_src_dir]:
        if not dir:
            continue
        testfilename = os.path.join(dir, filename)
        if os.path.isfile(testfilename):
            filename = testfilename
            break

    if known.inplace:
        outfile = filename

    allchecklines = []
    checkprefixes = []

    with open(filename, "r") as file:
        oldlines = [line.rstrip("\r\n") for line in file.readlines()]

    runlines = []
    for line in oldlines:
        m = runre.match(line)
        if m:
            runlines.append(m.group("tool"))

    continuation = ""
    newrunlines = []
    for line in runlines:
        if line.endswith("\\"):
            continuation += line[:-2] + " "
        else:
            newrunlines.append(continuation + line)
            continuation = ""
    if continuation:
        newrunlines.append(continuation)

    for line in newrunlines:
        m = filecheckre.match(line)
        if not m:
            continue

        tool, filecheck = m.group("tool", "filecheck")
        filecheck = shlex.split(filecheck)
        tool = shlex.split(tool)
        if known.bindir is not None:
            tool[0] = complete_exename(known.bindir, tool[0])
        if os.path.isdir(llvm_tools_dir):
            tool[0] = complete_exename(llvm_tools_dir, tool[0])
        check_prefix = filecheckparser.parse_known_args(filecheck)[0].check_prefix
        if known.prefix_only is not None and not check_prefix in known.prefix_only:
            continue
        if check_prefix in checkprefixes:
            continue
        checkprefixes.append(check_prefix)

        newtool = []
        optstderr = None
        for toolarg in tool:
            toolarg = toolarg.replace("%s", filename)
            toolarg = toolarg.replace("%S", os.path.dirname(filename))
            if toolarg == "%loadPolly":
                if not llvm_polly_link_into_tools:
                    newtool += [
                        "-load",
                        os.path.join(polly_lib_dir, "LLVMPolly" + shlibext),
                    ]
                newtool.append("-polly-process-unprofitable")
                newtool.append("-polly-remarks-minimal")
            elif toolarg == "2>&1":
                optstderr = subprocess.STDOUT
            else:
                newtool.append(toolarg)
        tool = newtool

        inpfile = None
        i = 1
        while i < len(tool):
            if tool[i] == "<":
                inpfile = tool[i + 1]
                del tool[i : i + 2]
                continue
            i += 1
        if inpfile:
            with open(inpfile) as inp:
                retlines = subprocess.check_output(
                    tool, universal_newlines=True, stdin=inp, stderr=optstderr
                )
        else:
            retlines = subprocess.check_output(
                tool, universal_newlines=True, stderr=optstderr
            )
        retlines = [line.replace("\t", "    ") for line in retlines.splitlines()]
        check_include = []
        for checkme in known.check_include + known.check_label_include:
            parts = checkme.split("=")
            if len(parts) == 2:
                if parts[0] == check_prefix:
                    check_include.append(parts[1])
            else:
                check_include.append(checkme)

        if check_include:
            filtered_retlines = []
            classified_retlines = []
            lastmatch = None
            for line, kind in (
                (line, class1.union(class2))
                for line, class1, class2 in zip(
                    retlines, classyfier1(retlines), classyfier2(retlines)
                )
            ):
                match = kind.intersection(check_include)
                if match:
                    if lastmatch != match:
                        filtered_retlines.append("")
                        classified_retlines.append({"Separator"})
                    filtered_retlines.append(line)
                    classified_retlines.append(kind)
                lastmatch = match

            retlines = filtered_retlines
        else:
            classified_retlines = (set() for line in retlines)

        rtrim_emptylines(retlines)
        ltrim_emptylines(retlines, classified_retlines)
        retlines = [
            replre.sub(lambda m: replrepl[m.group(0)], line) for line in retlines
        ]
        indent = common_indent(retlines)
        retlines = [line[indent:] for line in retlines]
        checklines = []
        previous_was_empty = True
        for line, kind in zip(retlines, classified_retlines):
            if line:
                if known.check_style == "CHECK" and known.check_label_include:
                    if not kind.isdisjoint(known.check_label_include):
                        checklines.append("; " + check_prefix + "-LABEL: " + line)
                    else:
                        checklines.append("; " + check_prefix + ":       " + line)
                elif known.check_style == "CHECK":
                    checklines.append("; " + check_prefix + ": " + line)
                elif known.check_label_include and known.check_label_include:
                    if not kind.isdisjoint(known.check_label_include):
                        checklines.append("; " + check_prefix + "-LABEL: " + line)
                    elif previous_was_empty:
                        checklines.append("; " + check_prefix + ":       " + line)
                    else:
                        checklines.append("; " + check_prefix + "-NEXT:  " + line)
                else:
                    if previous_was_empty:
                        checklines.append("; " + check_prefix + ":      " + line)
                    else:
                        checklines.append("; " + check_prefix + "-NEXT: " + line)
                previous_was_empty = False
            else:
                if not "Separator" in kind or known.check_part_newline:
                    checklines.append(";")
                previous_was_empty = True
        allchecklines.append(checklines)

    if not checkprefixes:
        return

    checkre = re.compile(
        r"^\s*\;\s*("
        + "|".join([re.escape(s) for s in checkprefixes])
        + ")(\-NEXT|\-DAG|\-NOT|\-LABEL|\-SAME)?\s*\:"
    )
    firstcheckline = None
    firstnoncommentline = None
    headerlines = []
    newlines = []
    uptonowlines = []
    emptylines = []
    lastwascheck = False
    for line in oldlines:
        if checkre.match(line):
            if firstcheckline is None:
                firstcheckline = len(newlines) + len(emptylines)
            if not lastwascheck:
                uptonowlines += emptylines
            emptylines = []
            lastwascheck = True
        elif emptyline.fullmatch(line):
            emptylines.append(line)
        else:
            newlines += uptonowlines
            newlines += emptylines
            newlines.append(line)
            emptylines = []
            uptonowlines = []
            lastwascheck = False

    for i, line in enumerate(newlines):
        if not commentline.fullmatch(line):
            firstnoncommentline = i
            break

    with open(outfile, "w", newline="") as file:

        def writelines(lines):
            for line in lines:
                file.write(line)
                file.write("\n")

        if firstcheckline is not None and known.check_position == "autodetect":
            writelines(newlines[:firstcheckline])
            writelines(uptonowlines)
            for i, checklines in enumerate(allchecklines):
                if i != 0:
                    file.write("\n")
                writelines(checklines)
            writelines(newlines[firstcheckline:])
            writelines(emptylines)
        elif (
            firstnoncommentline is not None and known.check_position == "before-content"
        ):
            headerlines = newlines[:firstnoncommentline]
            rtrim_emptylines(headerlines)
            contentlines = newlines[firstnoncommentline:]
            ltrim_emptylines(contentlines)

            writelines(headerlines)
            for checklines in allchecklines:
                file.write("\n")
                writelines(checklines)
            file.write("\n")
            writelines(contentlines)
            writelines(uptonowlines)
            writelines(emptylines)
        else:
            writelines(newlines)
            rtrim_emptylines(newlines)
            for checklines in allchecklines:
                file.write("\n\n")
                writelines(checklines)


if __name__ == "__main__":
    main()