Support generating checks for multiple modules

scripts/update_lit_checks.py

@@ -27,13 +27,15 @@
 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' +
           ' be edited.')
 
 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 emit_checks(indent, prefix, lines):
         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))

test/lit/passes/extract-function.wast

@@ -1,71 +1,84 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
 ;; RUN: foreach %s %t wasm-opt --extract-function=foo -S -o - | filecheck %s
 ;; RUN: foreach %s %t wasm-opt --extract-function --pass-arg=extract-function@foo -S -o - | filecheck %s
 ;; RUN: foreach %s %t wasm-opt --extract-function-index=0 -S -o - | filecheck %s
 ;; RUN: foreach %s %t wasm-opt --extract-function-index --pass-arg=extract-function-index@0 -S -o - | filecheck %s
 
-;; CHECK:      (module
-;; CHECK-NEXT:  (type $none_=>_none (func))
-;; CHECK-NEXT:  (import "env" "bar" (func $bar))
-;; CHECK-NEXT:  (export "foo" (func $foo))
-;; CHECK-NEXT:  (func $foo
-;; CHECK-NEXT:   (call $bar)
-;; CHECK-NEXT:  )
-;; CHECK-NEXT: )
 (module
+  ;; CHECK:      (type $none_=>_none (func))
+
+  ;; CHECK:      (import "env" "bar" (func $bar))
+
+  ;; CHECK:      (export "foo" (func $foo))
+
+  ;; CHECK:      (func $foo
+  ;; CHECK-NEXT:  (call $bar)
+  ;; CHECK-NEXT: )
   (func $foo
     (call $bar)
   )
+
   (func $bar
     (call $foo)
   )
+
   (func $other
     (drop (i32.const 1))
   )
 )
 
-;; CHECK:      (module
-;; CHECK-NEXT:  (type $none_=>_none (func))
-;; CHECK-NEXT:  (import "env" "other" (func $other))
-;; CHECK-NEXT:  (export "foo" (func $foo))
-;; CHECK-NEXT:  (func $foo
-;; CHECK-NEXT:   (nop)
-;; CHECK-NEXT:  )
-;; CHECK-NEXT: )
 (module
   ;; Use another function in the table, but the table is not used in the
   ;; extracted function
   (table $t 10 funcref)
+
   (elem $0 (table $t) (i32.const 0) func $other)
+
+  ;; CHECK:      (type $none_=>_none (func))
+
+  ;; CHECK:      (import "env" "other" (func $other))
+
+  ;; CHECK:      (export "foo" (func $foo))
+
+  ;; CHECK:      (func $foo
+  ;; CHECK-NEXT:  (nop)
+  ;; CHECK-NEXT: )
   (func $foo
     (nop)
   )
+
   (func $other
     (drop (i32.const 1))
   )
 )
 
-;; CHECK:      (module
-;; CHECK-NEXT:  (type $none (func))
-;; CHECK-NEXT:  (import "env" "other" (func $other))
-;; CHECK-NEXT:  (table $t 10 funcref)
-;; CHECK-NEXT:  (elem $0 (i32.const 0) $other)
-;; CHECK-NEXT:  (export "foo" (func $foo))
-;; CHECK-NEXT:  (func $foo
-;; CHECK-NEXT:   (call_indirect (type $none)
-;; CHECK-NEXT:    (i32.const 10)
-;; CHECK-NEXT:   )
-;; CHECK-NEXT:  )
-;; CHECK-NEXT: )
 (module
   ;; Use another function in the table, and the table *is* used. As a result,
   ;; the table and its elements will remain. The called function, $other, will
   ;; remain as an import that is placed in the table.
+
+  ;; CHECK:      (type $none (func))
   (type $none (func))
+
+  ;; CHECK:      (import "env" "other" (func $other))
+
+  ;; CHECK:      (table $t 10 funcref)
   (table $t 10 funcref)
+
+  ;; CHECK:      (elem $0 (i32.const 0) $other)
   (elem $0 (table $t) (i32.const 0) func $other)
+
+  ;; CHECK:      (export "foo" (func $foo))
+
+  ;; CHECK:      (func $foo
+  ;; CHECK-NEXT:  (call_indirect (type $none)
+  ;; CHECK-NEXT:   (i32.const 10)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
   (func $foo
    (call_indirect (type $none) (i32.const 10))
   )
+
   (func $other
     (drop (i32.const 1))
   )