update_mir_test_checks: Accept IR as input as well as MIR

We need to handle IR for tests that want to do lowering (or just
-stop-after with IR as input). I've run this on one AArch64 test to
demonstrate what it looks like.

llvm-svn: 321048

1 files changed, 29 insertions, 2 deletions

diff --git a/llvm/utils/update_mir_test_checks.py b/llvm/utils/update_mir_test_checks.py
index f0b28124262..3756af1b517 100755
--- a/llvm/utils/update_mir_test_checks.py
+++ b/llvm/utils/update_mir_test_checks.py
@@ -43,6 +43,10 @@ VREG_DEF_RE = re.compile(
 MIR_PREFIX_DATA_RE = re.compile(r'^ *(;|bb.[0-9].*: *$|[a-z]+:( |$)|$)')
 VREG_CLASS_RE = re.compile(r'^ *- *{ id: ([0-9]+), class: ([a-z0-9_]+)', re.M)
 
+IR_FUNC_NAME_RE = re.compile(
+    r'^\s*define\s+(?:internal\s+)?[^@]*@(?P<func>\w+)\s*\(')
+IR_PREFIX_DATA_RE = re.compile(r'^ *(;|$)')
+
 MIR_FUNC_RE = re.compile(
     r'^---$'
     r'\n'
@@ -340,8 +344,10 @@ def update_test_file(llc, test, remove_common_prefixes=False,
         warn('Ignoring common prefixes: {}'.format(common_prefixes),
              test_file=test)
 
-    autogenerated_note = ('# NOTE: Assertions have been autogenerated by '
-                          'utils/{}'.format(os.path.basename(__file__)))
+    comment_char = '#' if test.endswith('.mir') else ';'
+    autogenerated_note = ('{} NOTE: Assertions have been autogenerated by '
+                          'utils/{}'.format(comment_char,
+                                            os.path.basename(__file__)))
     output_lines = []
     output_lines.append(autogenerated_note)
 
@@ -350,6 +356,10 @@ def update_test_file(llc, test, remove_common_prefixes=False,
             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.strip() == '---':
                 state = 'document'
             output_lines.append(input_line)
@@ -392,6 +402,23 @@ def update_test_file(llc, test, remove_common_prefixes=False,
                 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, add_vreg_checks,
+                                        single_bb=False, verbose=verbose)
+
+            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)
+
 
     log('Writing {} lines to {}...'.format(len(output_lines), test), verbose)