[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.

1 files changed, 238 insertions, 0 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: