Support generating checks for multiple modules (#3962)

In conjunction with the `foreach` tool, allows autogenerating checks for lit
tests containing multiple modules. Supporting this will help automatically port
existing bespoke wast tests to be lit tests, since many of those tests contain
multiple modules per file.

1 files changed, 94 insertions, 52 deletions

diff --git a/scripts/update_lit_checks.py b/scripts/update_lit_checks.py
index e6b5435b3..9288aa3ca 100755
--- a/scripts/update_lit_checks.py
+++ b/scripts/update_lit_checks.py
@@ -27,6 +27,7 @@ import subprocess
 import sys
 import tempfile
 
+script_dir = os.path.dirname(__file__)
 script_name = os.path.basename(__file__)
 
 NOTICE = (';; NOTE: Assertions have been generated by {script} and should not' +
@@ -34,6 +35,7 @@ NOTICE = (';; NOTE: Assertions have been generated by {script} and should not' +
 
 RUN_LINE_RE = re.compile(r'^\s*;;\s*RUN:\s*(.*)$')
 CHECK_PREFIX_RE = re.compile(r'.*--check-prefix[= ](\S+).*')
+MODULE_RE = re.compile(r'^\(module.*$', re.MULTILINE)
 
 items = ['type', 'import', 'global', 'memory', 'data', 'table', 'elem', 'tag',
          'export', 'start', 'func']
@@ -85,6 +87,7 @@ def run_command(args, test, tmp, command):
     env['PATH'] = args.binaryen_bin + os.pathsep + env['PATH']
     command = command.replace('%s', test)
     command = command.replace('%t', tmp)
+    command = command.replace('foreach', os.path.join(script_dir, 'foreach.py'))
     return subprocess.check_output(command, shell=True, env=env).decode('utf-8')
 
 
@@ -103,20 +106,40 @@ def find_end(module, start):
     return end
 
 
-def parse_output(module):
-    # Return a list of (name, [lines]) for module items
-    out = []
-    for match in ITEM_RE.finditer(module):
-        kind, name = match[2], match[3]
-        end = find_end(module, match.end(1))
-        lines = module[match.start():end].split('\n')
-        out.append(((kind, name), lines))
-    return out
+def split_modules(text):
+    # Return a list of strings; one for each module
+    module_starts = [match.start() for match in MODULE_RE.finditer(text)]
+    if len(module_starts) < 2:
+        return [text]
+    first_module = text[:module_starts[1]]
+    modules = [first_module]
+    for i in range(1, len(module_starts) - 1):
+        module = text[module_starts[i]:module_starts[i + 1]]
+        modules.append(module)
+    last_module = text[module_starts[-1]:]
+    modules.append(last_module)
+    return modules
+
+
+def parse_output(text):
+    # Return a list containing, for each module in the text, a list of
+    # (name, [line]) for module items.
+    modules = []
+    for module in split_modules(text):
+        items = []
+        for match in ITEM_RE.finditer(module):
+            kind, name = match[2], match[3]
+            end = find_end(module, match.end(1))
+            lines = module[match.start():end].split('\n')
+            items.append(((kind, name), lines))
+        modules.append(items)
+    return modules
 
 
 def get_command_output(args, test, lines, tmp):
-    # Map check prefixes to lists of ((kind, name), [line])
-    command_output = {}
+    # Return list of maps from prefixes to lists of module items of the form
+    # ((kind, name), [line]). The outer list has an entry for each module.
+    command_output = []
     for line in find_run_lines(test, lines):
         commands = [cmd.strip() for cmd in line.rsplit('|', 1)]
         filecheck_cmd = ''
@@ -134,7 +157,11 @@ def get_command_output(args, test, lines, tmp):
 
         output = run_command(args, test, tmp, commands[0])
         if prefix:
-            command_output[prefix] = parse_output(output)
+            module_outputs = parse_output(output)
+            for i in range(len(module_outputs)):
+                if len(command_output) == i:
+                    command_output.append({})
+                command_output[i][prefix] = module_outputs[i]
 
     return command_output
 
@@ -150,8 +177,10 @@ def update_test(args, test, lines, tmp):
 
     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 +
+    prefixes = set(prefix
+                   for module_output in command_output
+                   for prefix in module_output.keys())
+    check_line_re = re.compile(r'^\s*;;\s*(' + '|'.join(prefixes) +
                                r')(?:-NEXT|-LABEL|-NOT)?:.*$')
 
     # Filter out whitespace between check blocks
@@ -181,48 +210,61 @@ def update_test(args, test, lines, tmp):
         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
+    input_modules = [m.split('\n') for m in split_modules('\n'.join(lines))]
+    if len(input_modules) > len(command_output):
+        warn('Fewer output modules than input modules:'
+             'not all modules will get checks.')
 
-        match = ITEM_RE.match(line)
-        if not match:
-            output_lines.append(line)
-            continue
+    # Remove extra newlines at the end of modules
+    input_modules = [m[:-1] for m in input_modules[:-1]] + [input_modules[-1]]
 
-        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)
+    for module_idx in range(len(input_modules)):
+        output = command_output[module_idx] \
+            if module_idx < len(command_output) else {}
+
+        for line in input_modules[module_idx]:
+            # 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 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
-        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 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.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))