[NFC] Refactor update_lit_checks.py (#3964)

Decompose the code into more functions and make other simplifying changes to
prepare for multi-module support introduced in #3962.

1 files changed, 120 insertions, 111 deletions

diff --git a/scripts/update_lit_checks.py b/scripts/update_lit_checks.py
index 083298f57..e6b5435b3 100755
--- a/scripts/update_lit_checks.py
+++ b/scripts/update_lit_checks.py
@@ -27,7 +27,6 @@ import subprocess
 import sys
 import tempfile
 
-
 script_name = os.path.basename(__file__)
 
 NOTICE = (';; NOTE: Assertions have been generated by {script} and should not' +
@@ -104,7 +103,7 @@ def find_end(module, start):
     return end
 
 
-def split_output(module):
+def parse_output(module):
     # Return a list of (name, [lines]) for module items
     out = []
     for match in ITEM_RE.finditer(module):
@@ -115,6 +114,124 @@ def split_output(module):
     return out
 
 
+def get_command_output(args, test, lines, tmp):
+    # Map check prefixes to lists of ((kind, name), [line])
+    command_output = {}
+    for line in find_run_lines(test, lines):
+        commands = [cmd.strip() for cmd in line.rsplit('|', 1)]
+        filecheck_cmd = ''
+        if len(commands) > 1 and commands[1].startswith('filecheck '):
+            filecheck_cmd = commands[1]
+            commands = commands[:1]
+
+        prefix = ''
+        if filecheck_cmd.startswith('filecheck '):
+            prefix_match = CHECK_PREFIX_RE.match(filecheck_cmd)
+            if prefix_match:
+                prefix = prefix_match.group(1)
+            else:
+                prefix = 'CHECK'
+
+        output = run_command(args, test, tmp, commands[0])
+        if prefix:
+            command_output[prefix] = parse_output(output)
+
+    return command_output
+
+
+def update_test(args, test, lines, tmp):
+    all_items = args.all_items
+    if lines and script_name in lines[0]:
+        # Apply previously used options for this file
+        if '--all-items' in lines[0]:
+            all_items = True
+        # Skip the notice if it is already in the output
+        lines = lines[1:]
+
+    command_output = get_command_output(args, test, lines, tmp)
+
+    any_prefix = '|'.join(command_output.keys())
+    check_line_re = re.compile(r'^\s*;;\s*(' + any_prefix +
+                               r')(?:-NEXT|-LABEL|-NOT)?:.*$')
+
+    # Filter out whitespace between check blocks
+    if lines:
+        filtered = [lines[0]]
+        for i in range(1, len(lines) - 1):
+            if lines[i] or not check_line_re.match(lines[i - 1]) or \
+               not check_line_re.match(lines[i + 1]):
+                filtered.append(lines[i])
+        filtered.append(lines[-1])
+        lines = filtered
+
+    named_items = []
+    for line in lines:
+        match = ITEM_RE.match(line)
+        if match:
+            kind, name = match[2], match[3]
+            named_items.append((kind, name))
+
+    script = script_name
+    if all_items:
+        script += ' --all-items'
+    output_lines = [NOTICE.format(script=script)]
+
+    def emit_checks(indent, prefix, lines):
+        output_lines.append(f'{indent};; {prefix}:     {lines[0]}')
+        for line in lines[1:]:
+            output_lines.append(f'{indent};; {prefix}-NEXT:{line}')
+
+    for line in lines:
+        # Skip pre-existing check lines; we will regenerate them.
+        if check_line_re.match(line):
+            continue
+
+        match = ITEM_RE.match(line)
+        if not match:
+            output_lines.append(line)
+            continue
+
+        indent, kind, name = match.groups()
+        for prefix, items in command_output.items():
+            # If the output for this prefix contains an item with this
+            # name, emit all the items up to and including the matching
+            # item
+            has_item = False
+            for kind_name, lines in items:
+                if name and (kind, name) == kind_name:
+                    has_item = True
+                    break
+            if has_item:
+                first = True
+                while True:
+                    kind_name, lines = items.pop(0)
+                    if all_items or kind_name in named_items:
+                        if not first:
+                            output_lines.append('')
+                        first = False
+                        emit_checks(indent, prefix, lines)
+                    if name and (kind, name) == kind_name:
+                        break
+        output_lines.append(line)
+
+    # Output any remaining checks for each prefix
+    first = True
+    for prefix, items in command_output.items():
+        for kind_name, lines in items:
+            if all_items or kind_name in named_items:
+                if not first:
+                    output_lines.append('')
+                first = False
+                emit_checks('', prefix, lines)
+
+    if args.dry_run:
+        print('\n'.join(output_lines))
+    else:
+        with open(test, 'w') as f:
+            for line in output_lines:
+                f.write(line + '\n')
+
+
 def main():
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(
@@ -141,115 +258,7 @@ def main():
     tmp = tempfile.mktemp()
 
     for test, lines in itertests(args):
-        all_items = args.all_items
-        if lines and script_name in lines[0]:
-            # Apply previously used options for this file
-            if '--all-items' in lines[0]:
-                all_items = True
-            # Skip the notice if it is already in the output
-            lines = lines[1:]
-
-        # List of (prefix, command)
-        run_list = []
-        for line in find_run_lines(test, lines):
-            commands = [cmd.strip() for cmd in line.rsplit('|', 1)]
-            filecheck_cmd = ''
-            if len(commands) > 1 and commands[1].startswith('filecheck '):
-                filecheck_cmd = commands[1]
-                commands = commands[:1]
-
-            check_prefix = ''
-            if filecheck_cmd.startswith('filecheck '):
-                prefix_match = CHECK_PREFIX_RE.match(filecheck_cmd)
-                if prefix_match:
-                    check_prefix = prefix_match.group(1)
-                else:
-                    check_prefix = 'CHECK'
-
-            run_list.append((check_prefix, commands[0]))
-
-        # Map check prefixes to lists of ((kind, name), [line])
-        output_modules = {}
-        for prefix, command, in run_list:
-            output = run_command(args, test, tmp, command)
-            if prefix:
-                output_modules[prefix] = split_output(output)
-
-        any_prefix = '|'.join(output_modules.keys())
-        check_line_re = re.compile(r'^\s*;;\s*(' + any_prefix +
-                                   r')(?:-NEXT|-LABEL|-NOT)?:.*$')
-
-        # Filter out whitespace between check blocks
-        if lines:
-            filtered = [lines[0]]
-            for i in range(1, len(lines) - 1):
-                if lines[i] or not check_line_re.match(lines[i - 1]) or \
-                   not check_line_re.match(lines[i + 1]):
-                    filtered.append(lines[i])
-            filtered.append(lines[-1])
-            lines = filtered
-
-        named_items = []
-        for line in lines:
-            match = ITEM_RE.match(line)
-            if match:
-                kind, name = match[2], match[3]
-                named_items.append((kind, name))
-
-        script = script_name
-        if all_items:
-            script += ' --all-items'
-        output_lines = [NOTICE.format(script=script)]
-
-        def emit_checks(indent, prefix, lines):
-            output_lines.append(f'{indent};; {prefix}:     {lines[0]}')
-            for line in lines[1:]:
-                output_lines.append(f'{indent};; {prefix}-NEXT:{line}')
-
-        for line in lines:
-            # Skip pre-existing check lines; we will regenerate them.
-            if check_line_re.match(line):
-                continue
-            match = ITEM_RE.match(line)
-            if match:
-                indent, kind, name = match.groups()
-                for prefix, items in output_modules.items():
-                    # If the output for this prefix contains an item with this
-                    # name, emit all the items up to and including the matching
-                    # item
-                    has_item = False
-                    for kind_name, lines in items:
-                        if name and (kind, name) == kind_name:
-                            has_item = True
-                            break
-                    if has_item:
-                        first = True
-                        while True:
-                            kind_name, lines = items.pop(0)
-                            if all_items or kind_name in named_items:
-                                if not first:
-                                    output_lines.append('')
-                                first = False
-                                emit_checks(indent, prefix, lines)
-                            if name and (kind, name) == kind_name:
-                                break
-            output_lines.append(line)
-        # Output any remaining checks for each prefix
-        first = True
-        for prefix, items in output_modules.items():
-            for kind_name, lines in items:
-                if all_items or kind_name in named_items:
-                    if not first:
-                        output_lines.append('')
-                    first = False
-                    emit_checks('', prefix, lines)
-
-        if args.dry_run:
-            print('\n'.join(output_lines))
-        else:
-            with open(test, 'w') as f:
-                for line in output_lines:
-                    f.write(line + '\n')
+        update_test(args, test, lines, tmp)
 
 
 if __name__ == '__main__':