[GlobalISel] Add a update_givaluetracking_test_checks.py script (#140296)

As with the other update scripts this takes the output of
-passes=print<gisel-value-tracking> and inserts the results into an
existing mir file. This means that the input is a lot like
update_analysis_test_checks.py, and the output needs to insert into a
mir file similarly to update_mir_test_checks.py. The code used to do the
inserting has been moved to common, to allow it to be reused. Otherwise
it tries to reuse the existing infrastructure, and
update_givaluetracking_test_checks is kept relatively short.

3 files changed, 399 insertions, 229 deletions

diff --git a/llvm/utils/UpdateTestChecks/common.py b/llvm/utils/UpdateTestChecks/common.py
index 0074858..178c623 100644
--- a/llvm/utils/UpdateTestChecks/common.py
+++ b/llvm/utils/UpdateTestChecks/common.py
@@ -2271,6 +2271,244 @@ def add_analyze_checks(
     )
 
 
+IR_FUNC_NAME_RE = re.compile(
+    r"^\s*define\s+(?:internal\s+)?[^@]*@(?P<func>[A-Za-z0-9_.]+)\s*\("
+)
+IR_PREFIX_DATA_RE = re.compile(r"^ *(;|$)")
+MIR_FUNC_NAME_RE = re.compile(r" *name: *(?P<func>[A-Za-z0-9_.-]+)")
+MIR_BODY_BEGIN_RE = re.compile(r" *body: *\|")
+MIR_BASIC_BLOCK_RE = re.compile(r" *bb\.[0-9]+.*:$")
+MIR_PREFIX_DATA_RE = re.compile(r"^ *(;|bb.[0-9].*: *$|[a-z]+:( |$)|$)")
+
+
+def find_mir_functions_with_one_bb(lines, verbose=False):
+    result = []
+    cur_func = None
+    bbs = 0
+    for line in lines:
+        m = MIR_FUNC_NAME_RE.match(line)
+        if m:
+            if bbs == 1:
+                result.append(cur_func)
+            cur_func = m.group("func")
+            bbs = 0
+        m = MIR_BASIC_BLOCK_RE.match(line)
+        if m:
+            bbs += 1
+    if bbs == 1:
+        result.append(cur_func)
+    return result
+
+
+def add_mir_checks_for_function(
+    test,
+    output_lines,
+    run_list,
+    func_dict,
+    func_name,
+    single_bb,
+    print_fixed_stack,
+    first_check_is_next,
+    at_the_function_name,
+):
+    printed_prefixes = set()
+    for run in run_list:
+        for prefix in run[0]:
+            if prefix in printed_prefixes:
+                break
+            if not func_dict[prefix][func_name]:
+                continue
+            if printed_prefixes:
+                # Add some space between different check prefixes.
+                indent = len(output_lines[-1]) - len(output_lines[-1].lstrip(" "))
+                output_lines.append(" " * indent + ";")
+            printed_prefixes.add(prefix)
+            add_mir_check_lines(
+                test,
+                output_lines,
+                prefix,
+                ("@" if at_the_function_name else "") + func_name,
+                single_bb,
+                func_dict[prefix][func_name],
+                print_fixed_stack,
+                first_check_is_next,
+            )
+            break
+        else:
+            warn(
+                "Found conflicting asm for function: {}".format(func_name),
+                test_file=test,
+            )
+    return output_lines
+
+
+def add_mir_check_lines(
+    test,
+    output_lines,
+    prefix,
+    func_name,
+    single_bb,
+    func_info,
+    print_fixed_stack,
+    first_check_is_next,
+):
+    func_body = str(func_info).splitlines()
+    if single_bb:
+        # Don't bother checking the basic block label for a single BB
+        func_body.pop(0)
+
+    if not func_body:
+        warn(
+            "Function has no instructions to check: {}".format(func_name),
+            test_file=test,
+        )
+        return
+
+    first_line = func_body[0]
+    indent = len(first_line) - len(first_line.lstrip(" "))
+    # A check comment, indented the appropriate amount
+    check = "{:>{}}; {}".format("", indent, prefix)
+
+    output_lines.append("{}-LABEL: name: {}".format(check, func_name))
+
+    if print_fixed_stack:
+        output_lines.append("{}: fixedStack:".format(check))
+        for stack_line in func_info.extrascrub.splitlines():
+            filecheck_directive = check + "-NEXT"
+            output_lines.append("{}: {}".format(filecheck_directive, stack_line))
+
+    first_check = not first_check_is_next
+    for func_line in func_body:
+        if not func_line.strip():
+            # The mir printer prints leading whitespace so we can't use CHECK-EMPTY:
+            output_lines.append(check + "-NEXT: {{" + func_line + "$}}")
+            continue
+        filecheck_directive = check if first_check else check + "-NEXT"
+        first_check = False
+        check_line = "{}: {}".format(filecheck_directive, func_line[indent:]).rstrip()
+        output_lines.append(check_line)
+
+
+def should_add_mir_line_to_output(input_line, prefix_set):
+    # Skip any check lines that we're handling as well as comments
+    m = CHECK_RE.match(input_line)
+    if (m and m.group(1) in prefix_set) or input_line.strip() == ";":
+        return False
+    return True
+
+
+def add_mir_checks(
+    input_lines,
+    prefix_set,
+    autogenerated_note,
+    test,
+    run_list,
+    func_dict,
+    print_fixed_stack,
+    first_check_is_next,
+    at_the_function_name,
+):
+    simple_functions = find_mir_functions_with_one_bb(input_lines)
+
+    output_lines = []
+    output_lines.append(autogenerated_note)
+
+    func_name = None
+    state = "toplevel"
+    for input_line in input_lines:
+        if input_line == autogenerated_note:
+            continue
+
+        if state == "toplevel":
+            m = IR_FUNC_NAME_RE.match(input_line)
+            if m:
+                state = "ir function prefix"
+                func_name = m.group("func")
+            if input_line.rstrip("| \r\n") == "---":
+                state = "document"
+            output_lines.append(input_line)
+        elif state == "document":
+            m = MIR_FUNC_NAME_RE.match(input_line)
+            if m:
+                state = "mir function metadata"
+                func_name = m.group("func")
+            if input_line.strip() == "...":
+                state = "toplevel"
+                func_name = None
+            if should_add_mir_line_to_output(input_line, prefix_set):
+                output_lines.append(input_line)
+        elif state == "mir function metadata":
+            if should_add_mir_line_to_output(input_line, prefix_set):
+                output_lines.append(input_line)
+            m = MIR_BODY_BEGIN_RE.match(input_line)
+            if m:
+                if func_name in simple_functions:
+                    # If there's only one block, put the checks inside it
+                    state = "mir function prefix"
+                    continue
+                state = "mir function body"
+                add_mir_checks_for_function(
+                    test,
+                    output_lines,
+                    run_list,
+                    func_dict,
+                    func_name,
+                    single_bb=False,
+                    print_fixed_stack=print_fixed_stack,
+                    first_check_is_next=first_check_is_next,
+                    at_the_function_name=at_the_function_name,
+                )
+        elif state == "mir function prefix":
+            m = MIR_PREFIX_DATA_RE.match(input_line)
+            if not m:
+                state = "mir function body"
+                add_mir_checks_for_function(
+                    test,
+                    output_lines,
+                    run_list,
+                    func_dict,
+                    func_name,
+                    single_bb=True,
+                    print_fixed_stack=print_fixed_stack,
+                    first_check_is_next=first_check_is_next,
+                    at_the_function_name=at_the_function_name,
+                )
+
+            if should_add_mir_line_to_output(input_line, prefix_set):
+                output_lines.append(input_line)
+        elif state == "mir function body":
+            if input_line.strip() == "...":
+                state = "toplevel"
+                func_name = None
+            if should_add_mir_line_to_output(input_line, prefix_set):
+                output_lines.append(input_line)
+        elif state == "ir function prefix":
+            m = IR_PREFIX_DATA_RE.match(input_line)
+            if not m:
+                state = "ir function body"
+                add_mir_checks_for_function(
+                    test,
+                    output_lines,
+                    run_list,
+                    func_dict,
+                    func_name,
+                    single_bb=False,
+                    print_fixed_stack=print_fixed_stack,
+                    first_check_is_next=first_check_is_next,
+                    at_the_function_name=at_the_function_name,
+                )
+
+            if should_add_mir_line_to_output(input_line, prefix_set):
+                output_lines.append(input_line)
+        elif state == "ir function body":
+            if input_line.strip() == "}":
+                state = "toplevel"
+                func_name = None
+            if should_add_mir_line_to_output(input_line, prefix_set):
+                output_lines.append(input_line)
+    return output_lines
+
+
 def build_global_values_dictionary(glob_val_dict, raw_tool_output, prefixes, ginfo):
     for nameless_value in ginfo.get_nameless_values():
         if nameless_value.global_ir_rhs_regexp is None:
diff --git a/llvm/utils/update_givaluetracking_test_checks.py b/llvm/utils/update_givaluetracking_test_checks.py
new file mode 100755
index 0000000..3a8f539
--- /dev/null
+++ b/llvm/utils/update_givaluetracking_test_checks.py
@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+
+"""Updates FileCheck checks in GlobalISel Known Bits tests.
+
+This script is a utility to update MIR based tests with new FileCheck
+patterns for GlobalISel Known Bits.
+
+The checks added by this script are similar to update_mir_test_checks, using
+the output of KnownBits and SignBits from -passes=print<gisel-value-tracking>.
+"""
+
+from __future__ import print_function
+
+from sys import stderr
+from traceback import print_exc
+import argparse
+import os
+import re
+import sys
+
+from UpdateTestChecks import common
+
+VT_FUNCTION_RE = re.compile(
+    r"\s*name:\s*@(?P<func>[A-Za-z0-9_-]+)"
+    r"(?P<body>(\s*%[0-9a-zA-Z_]+:_\s*KnownBits:[01?]+\sSignBits:[0-9]+$)+)",
+    flags=(re.X | re.M),
+)
+
+
+def update_test(ti: common.TestInfo):
+    run_list = []
+    for l in ti.run_lines:
+        if "|" not in l:
+            common.warn("Skipping unparsable RUN line: " + l)
+            continue
+
+        (llc_cmd, filecheck_cmd) = tuple([cmd.strip() for cmd in l.split("|", 1)])
+        common.verify_filecheck_prefixes(filecheck_cmd)
+
+        if not llc_cmd.startswith("llc "):
+            common.warn("Skipping non-llc RUN line: " + l)
+            continue
+
+        if not filecheck_cmd.startswith("FileCheck "):
+            common.warn("Skipping non-FileChecked RUN line: " + l)
+            continue
+
+        llc_cmd_args = llc_cmd[4:].strip()
+        llc_cmd_args = llc_cmd_args.replace("< %s", "").replace("%s", "").strip()
+        check_prefixes = common.get_check_prefixes(filecheck_cmd)
+
+        run_list.append((check_prefixes, llc_cmd_args))
+
+    ginfo = common.make_analyze_generalizer(version=1)
+    builder = common.FunctionTestBuilder(
+        run_list=run_list,
+        flags=type(
+            "",
+            (object,),
+            {
+                "verbose": ti.args.verbose,
+                "filters": ti.args.filters,
+                "function_signature": False,
+                "check_attributes": False,
+                "replace_value_regex": [],
+            },
+        ),
+        scrubber_args=[],
+        path=ti.path,
+        ginfo=ginfo,
+    )
+
+    for prefixes, llc_args in run_list:
+        common.debug("Extracted llc cmd:", "llc", llc_args)
+        common.debug("Extracted FileCheck prefixes:", str(prefixes))
+
+        if ti.path.endswith(".mir"):
+            llc_args += " -x mir"
+        raw_tool_output = common.invoke_tool(
+            ti.args.llc_binary or "llc", llc_args, ti.path, verbose=ti.args.verbose
+        )
+
+        builder.process_run_line(
+            VT_FUNCTION_RE,
+            common.scrub_body,
+            raw_tool_output,
+            prefixes,
+        )
+
+        builder.processed_prefixes(prefixes)
+
+    func_dict = builder.finish_and_get_func_dict()
+    prefix_set = set([prefix for p in run_list for prefix in p[0]])
+    common.debug("Rewriting FileCheck prefixes:", str(prefix_set))
+    output_lines = common.add_mir_checks(
+        ti.input_lines,
+        prefix_set,
+        ti.test_autogenerated_note,
+        ti.path,
+        run_list,
+        func_dict,
+        print_fixed_stack=False,
+        first_check_is_next=True,
+        at_the_function_name=True,
+    )
+
+    common.debug("Writing %d lines to %s..." % (len(output_lines), ti.path))
+
+    with open(ti.path, "wb") as f:
+        f.writelines(["{}\n".format(l).encode("utf-8") for l in output_lines])
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawTextHelpFormatter
+    )
+    parser.add_argument(
+        "--llc-binary",
+        default=None,
+        help='The "llc" binary to generate the test case with',
+    )
+    parser.add_argument("tests", nargs="+")
+    args = common.parse_commandline_args(parser)
+
+    script_name = os.path.basename(__file__)
+    returncode = 0
+    for ti in common.itertests(args.tests, parser, script_name="utils/" + script_name):
+        try:
+            update_test(ti)
+        except Exception:
+            stderr.write(f"Error: Failed to update test {ti.path}\n")
+            print_exc()
+            returncode = 1
+    return returncode
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/llvm/utils/update_mir_test_checks.py b/llvm/utils/update_mir_test_checks.py
index 8deaec7..8db46ad 100755
--- a/llvm/utils/update_mir_test_checks.py
+++ b/llvm/utils/update_mir_test_checks.py
@@ -32,9 +32,6 @@ import sys
 
 from UpdateTestChecks import common
 
-MIR_FUNC_NAME_RE = re.compile(r" *name: *(?P<func>[A-Za-z0-9_.-]+)")
-MIR_BODY_BEGIN_RE = re.compile(r" *body: *\|")
-MIR_BASIC_BLOCK_RE = re.compile(r" *bb\.[0-9]+.*:$")
 VREG_RE = re.compile(r"(%[0-9]+)(?:\.[a-z0-9_]+)?(?::[a-z0-9_]+)?(?:\([<>a-z0-9 ]+\))?")
 MI_FLAGS_STR = (
     r"(frame-setup |frame-destroy |nnan |ninf |nsz |arcp |contract |afn "
@@ -47,12 +44,6 @@ VREG_DEF_RE = re.compile(
         VREG_RE.pattern, MI_FLAGS_STR, VREG_DEF_FLAGS_STR
     )
 )
-MIR_PREFIX_DATA_RE = re.compile(r"^ *(;|bb.[0-9].*: *$|[a-z]+:( |$)|$)")
-
-IR_FUNC_NAME_RE = re.compile(
-    r"^\s*define\s+(?:internal\s+)?[^@]*@(?P<func>[A-Za-z0-9_.]+)\s*\("
-)
-IR_PREFIX_DATA_RE = re.compile(r"^ *(;|$)")
 
 MIR_FUNC_RE = re.compile(
     r"^---$"
@@ -88,16 +79,6 @@ class LLC:
         return stdout
 
 
-class Run:
-    def __init__(self, prefixes, cmd_args, triple):
-        self.prefixes = prefixes
-        self.cmd_args = cmd_args
-        self.triple = triple
-
-    def __getitem__(self, index):
-        return [self.prefixes, self.cmd_args, self.triple][index]
-
-
 def log(msg, verbose=True):
     if verbose:
         print(msg, file=sys.stderr)
@@ -147,46 +128,16 @@ def build_run_list(test, run_lines, verbose=False):
         check_prefixes = common.get_check_prefixes(filecheck_cmd)
         all_prefixes += check_prefixes
 
-        run_list.append(Run(check_prefixes, cmd_args, triple))
+        run_list.append((check_prefixes, cmd_args, triple))
 
     # Sort prefixes that are shared between run lines before unshared prefixes.
     # This causes us to prefer printing shared prefixes.
     for run in run_list:
-        run.prefixes.sort(key=lambda prefix: -all_prefixes.count(prefix))
+        run[0].sort(key=lambda prefix: -all_prefixes.count(prefix))
 
     return run_list
 
 
-def find_functions_with_one_bb(lines, verbose=False):
-    result = []
-    cur_func = None
-    bbs = 0
-    for line in lines:
-        m = MIR_FUNC_NAME_RE.match(line)
-        if m:
-            if bbs == 1:
-                result.append(cur_func)
-            cur_func = m.group("func")
-            bbs = 0
-        m = MIR_BASIC_BLOCK_RE.match(line)
-        if m:
-            bbs += 1
-    if bbs == 1:
-        result.append(cur_func)
-    return result
-
-
-class FunctionInfo:
-    def __init__(self, body, fixedStack):
-        self.body = body
-        self.fixedStack = fixedStack
-
-    def __eq__(self, other):
-        if not isinstance(other, FunctionInfo):
-            return False
-        return self.body == other.body and self.fixedStack == other.fixedStack
-
-
 def build_function_info_dictionary(
     test, raw_tool_output, triple, prefixes, func_dict, verbose
 ):
@@ -222,88 +173,20 @@ def build_function_info_dictionary(
         body = "".join(mangled)
 
         for prefix in prefixes:
-            info = FunctionInfo(body, fixedStack)
+            info = common.function_body(
+                body, fixedStack, None, None, None, None, ginfo=None
+            )
             if func in func_dict[prefix]:
-                if func_dict[prefix][func] != info:
+                if (
+                    not func_dict[prefix][func]
+                    or func_dict[prefix][func].scrub != info.scrub
+                    or func_dict[prefix][func].extrascrub != info.extrascrub
+                ):
                     func_dict[prefix][func] = None
             else:
                 func_dict[prefix][func] = info
 
 
-def add_checks_for_function(
-    test, output_lines, run_list, func_dict, func_name, single_bb, args
-):
-    printed_prefixes = set()
-    for run in run_list:
-        for prefix in run.prefixes:
-            if prefix in printed_prefixes:
-                break
-            if not func_dict[prefix][func_name]:
-                continue
-            if printed_prefixes:
-                # Add some space between different check prefixes.
-                indent = len(output_lines[-1]) - len(output_lines[-1].lstrip(" "))
-                output_lines.append(" "*indent + ";")
-            printed_prefixes.add(prefix)
-            log("Adding {} lines for {}".format(prefix, func_name), args.verbose)
-            add_check_lines(
-                test,
-                output_lines,
-                prefix,
-                func_name,
-                single_bb,
-                func_dict[prefix][func_name],
-                args,
-            )
-            break
-        else:
-            common.warn(
-                "Found conflicting asm for function: {}".format(func_name),
-                test_file=test,
-            )
-    return output_lines
-
-
-def add_check_lines(
-    test, output_lines, prefix, func_name, single_bb, func_info: FunctionInfo, args
-):
-    func_body = func_info.body.splitlines()
-    if single_bb:
-        # Don't bother checking the basic block label for a single BB
-        func_body.pop(0)
-
-    if not func_body:
-        common.warn(
-            "Function has no instructions to check: {}".format(func_name),
-            test_file=test,
-        )
-        return
-
-    first_line = func_body[0]
-    indent = len(first_line) - len(first_line.lstrip(" "))
-    # A check comment, indented the appropriate amount
-    check = "{:>{}}; {}".format("", indent, prefix)
-
-    output_lines.append("{}-LABEL: name: {}".format(check, func_name))
-
-    if args.print_fixed_stack:
-        output_lines.append("{}: fixedStack:".format(check))
-        for stack_line in func_info.fixedStack.splitlines():
-            filecheck_directive = check + "-NEXT"
-            output_lines.append("{}: {}".format(filecheck_directive, stack_line))
-
-    first_check = True
-    for func_line in func_body:
-        if not func_line.strip():
-            # The mir printer prints leading whitespace so we can't use CHECK-EMPTY:
-            output_lines.append(check + "-NEXT: {{" + func_line + "$}}")
-            continue
-        filecheck_directive = check if first_check else check + "-NEXT"
-        first_check = False
-        check_line = "{}: {}".format(filecheck_directive, func_line[indent:]).rstrip()
-        output_lines.append(check_line)
-
-
 def mangle_vreg(opcode, current_names):
     base = opcode
     # Simplify some common prefixes and suffixes
@@ -338,14 +221,6 @@ def mangle_vreg(opcode, current_names):
     return base
 
 
-def should_add_line_to_output(input_line, prefix_set):
-    # Skip any check lines that we're handling as well as blank comment.
-    m = common.CHECK_RE.match(input_line)
-    if (m and m.group(1) in prefix_set) or input_line.strip() == ";":
-        return False
-    return True
-
-
 def update_test_file(args, test, autogenerated_note):
     with open(test) as fd:
         input_lines = [l.rstrip() for l in fd]
@@ -354,11 +229,9 @@ def update_test_file(args, test, autogenerated_note):
     run_lines = common.find_run_lines(test, input_lines)
     run_list = build_run_list(test, run_lines, args.verbose)
 
-    simple_functions = find_functions_with_one_bb(input_lines, args.verbose)
-
     func_dict = {}
     for run in run_list:
-        for prefix in run.prefixes:
+        for prefix in run[0]:
             func_dict.update({prefix: dict()})
     for prefixes, llc_args, triple_in_cmd in run_list:
         log("Extracted LLC cmd: llc {}".format(llc_args), args.verbose)
@@ -378,99 +251,20 @@ def update_test_file(args, test, autogenerated_note):
             args.verbose,
         )
 
-    state = "toplevel"
-    func_name = None
-    prefix_set = set([prefix for run in run_list for prefix in run.prefixes])
+    prefix_set = set([prefix for run in run_list for prefix in run[0]])
     log("Rewriting FileCheck prefixes: {}".format(prefix_set), args.verbose)
 
-    output_lines = []
-    output_lines.append(autogenerated_note)
-
-    for input_line in input_lines:
-        if input_line == autogenerated_note:
-            continue
-
-        if state == "toplevel":
-            m = IR_FUNC_NAME_RE.match(input_line)
-            if m:
-                state = "ir function prefix"
-                func_name = m.group("func")
-            if input_line.rstrip("| \r\n") == "---":
-                state = "document"
-            output_lines.append(input_line)
-        elif state == "document":
-            m = MIR_FUNC_NAME_RE.match(input_line)
-            if m:
-                state = "mir function metadata"
-                func_name = m.group("func")
-            if input_line.strip() == "...":
-                state = "toplevel"
-                func_name = None
-            if should_add_line_to_output(input_line, prefix_set):
-                output_lines.append(input_line)
-        elif state == "mir function metadata":
-            if should_add_line_to_output(input_line, prefix_set):
-                output_lines.append(input_line)
-            m = MIR_BODY_BEGIN_RE.match(input_line)
-            if m:
-                if func_name in simple_functions:
-                    # If there's only one block, put the checks inside it
-                    state = "mir function prefix"
-                    continue
-                state = "mir function body"
-                add_checks_for_function(
-                    test,
-                    output_lines,
-                    run_list,
-                    func_dict,
-                    func_name,
-                    single_bb=False,
-                    args=args,
-                )
-        elif state == "mir function prefix":
-            m = MIR_PREFIX_DATA_RE.match(input_line)
-            if not m:
-                state = "mir function body"
-                add_checks_for_function(
-                    test,
-                    output_lines,
-                    run_list,
-                    func_dict,
-                    func_name,
-                    single_bb=True,
-                    args=args,
-                )
-
-            if should_add_line_to_output(input_line, prefix_set):
-                output_lines.append(input_line)
-        elif state == "mir function body":
-            if input_line.strip() == "...":
-                state = "toplevel"
-                func_name = None
-            if should_add_line_to_output(input_line, prefix_set):
-                output_lines.append(input_line)
-        elif state == "ir function prefix":
-            m = IR_PREFIX_DATA_RE.match(input_line)
-            if not m:
-                state = "ir function body"
-                add_checks_for_function(
-                    test,
-                    output_lines,
-                    run_list,
-                    func_dict,
-                    func_name,
-                    single_bb=False,
-                    args=args,
-                )
-
-            if should_add_line_to_output(input_line, prefix_set):
-                output_lines.append(input_line)
-        elif state == "ir function body":
-            if input_line.strip() == "}":
-                state = "toplevel"
-                func_name = None
-            if should_add_line_to_output(input_line, prefix_set):
-                output_lines.append(input_line)
+    output_lines = common.add_mir_checks(
+        input_lines,
+        prefix_set,
+        autogenerated_note,
+        test,
+        run_list,
+        func_dict,
+        args.print_fixed_stack,
+        first_check_is_next=False,
+        at_the_function_name=False,
+    )
 
     log("Writing {} lines to {}...".format(len(output_lines), test), args.verbose)