Handle annotations in update_lit_checks.py

scripts/update_lit_checks.py

@@ -57,6 +57,8 @@
 
 FUZZ_EXEC_FUNC = re.compile(r'^\[fuzz-exec\] calling (?P<name>\S*)$')
 
+ANNOTATION_RE = re.compile(r'^\s*\(\@.*')
+
 
 def indentKindName(match):
     # Return the indent, kind, and name from an ITEM_RE match
@@ -127,6 +129,31 @@ def find_end(module, start):
     return end
 
 
+def find_annotations(module, start):
+    # Search backward to find the start of the line containing the first of the
+    # annotations preceding `start`, if any.
+    depth = 0
+    annotation = start
+    for i in range(start - 1, -1, -1):
+        if module[i] == ')':
+            depth += 1
+        elif module[i] == '(':
+            depth -= 1
+            if depth == 0:
+                if module[i + 1] == '@':
+                    # Found the start of a new annotation.
+                    annotation = i
+                else:
+                    # Found something that isn't an annotation.
+                    break
+    # Look for the start of the line containin the first annoation.
+    for i in range(annotation - 1, -1, -1):
+        if module[i] == '\n':
+            return i + 1
+    # The annotation should not have been on the first line of the module
+    assert False
+
+
 def split_modules(text):
     # Return a list of strings; one for each module
     module_starts = [match.start() for match in MODULE_RE.finditer(text)]
@@ -150,8 +177,9 @@ def parse_output_modules(text):
         items = []
         for match in ITEM_RE.finditer(module):
             _, kind, name = indentKindName(match)
+            start = find_annotations(module, match.start())
             end = find_end(module, match.end(1))
-            lines = module[match.start():end].split('\n')
+            lines = module[start:end].split('\n')
             items.append(((kind, name), lines))
         modules.append(items)
     return modules
@@ -277,7 +305,12 @@ def pad(line):
     # Remove extra newlines at the end of modules
     input_modules = [m[:-1] for m in input_modules[:-1]] + [input_modules[-1]]
 
+    # Collect annotation lines as we see them so we can put them after any
+    # checks we generate.
+    annotation_lines = []
+
     for module_idx in range(len(input_modules)):
+        assert len(annotation_lines) == 0
         output = command_output[module_idx] \
             if module_idx < len(command_output) else {}
 
@@ -286,8 +319,16 @@ def pad(line):
             if check_line_re.match(line):
                 continue
 
+            # Collect annotations to emit later once we know what they should
+            # attach to.
+            if ANNOTATION_RE.match(line):
+                annotation_lines.append(line)
+                continue
+
             match = ITEM_RE.match(line)
             if not match:
+                output_lines.extend(annotation_lines)
+                annotation_lines = []
                 output_lines.append(line)
                 continue
 
@@ -313,6 +354,8 @@ def pad(line):
                             emit_checks(indent, prefix, lines)
                         if name and (kind, name) == kind_name:
                             break
+            output_lines.extend(annotation_lines)
+            annotation_lines = []
             output_lines.append(line)
 
         # Output any remaining checks for each prefix

test/lit/inline-hints-func.wast

@@ -1,7 +1,14 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
 ;; RUN: wasm-opt -all             %s -S -o - | filecheck %s
 ;; RUN: wasm-opt -all --roundtrip %s -S -o - | filecheck %s
 
 (module
+  ;; CHECK:      (@metadata.code.inline "\12")
+  ;; CHECK-NEXT: (func $func-annotation (type $0)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
   (@metadata.code.inline "\12")
   (func $func-annotation
     ;; The annotation here is on the function.
@@ -11,13 +18,4 @@
   )
 )
 
-;; CHECK: (module
-;; CHECK-NEXT:  (type $0 (func))
-;; CHECK-NEXT:  (@metadata.code.inline "\12")
-;; CHECK-NEXT:  (func $func-annotation
-;; CHECK-NEXT:   (drop
-;; CHECK-NEXT:    (i32.const 0)
-;; CHECK-NEXT:   )
-;; CHECK-NEXT:  )
-;; CHECK-NEXT: )
 

test/lit/passes/strip-toolchain-annotations-func.wast

@@ -1,41 +1,39 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
 ;; RUN: wasm-opt -all --strip-toolchain-annotations %s -S -o - | filecheck %s
 
 (module
+  ;; CHECK:      (@metadata.code.inline "\00")
+  ;; CHECK-NEXT: (func $test-func-a (type $0)
+  ;; CHECK-NEXT: )
   (@metadata.code.inline "\00")
   (func $test-func-a
     ;; This VM annotation is kept.
   )
 
+  ;; CHECK:      (func $test-func-b (type $0)
+  ;; CHECK-NEXT: )
   (@binaryen.removable.if.unused)
   (func $test-func-b
     ;; Toolchain one is removed.
   )
 
+  ;; CHECK:      (@metadata.code.inline "\00")
+  ;; CHECK-NEXT: (func $test-func-c (type $0)
+  ;; CHECK-NEXT: )
   (@metadata.code.inline "\00")
   (@binaryen.removable.if.unused)
   (func $test-func-c
     ;; Toolchain one is removed, VM one is kept.
   )
 
+  ;; CHECK:      (@metadata.code.inline "\00")
+  ;; CHECK-NEXT: (func $test-func-d (type $0)
+  ;; CHECK-NEXT: )
   (@binaryen.removable.if.unused)
   (@metadata.code.inline "\00")
   (func $test-func-d
     ;; Reverse order of above.
   )
 )
 
-;; CHECK: (module
-;; CHECK-NEXT:  (type $0 (func))
-;; CHECK-NEXT:  (@metadata.code.inline "\00")
-;; CHECK-NEXT:  (func $test-func-a (type $0)
-;; CHECK-NEXT:  )
-;; CHECK-NEXT:  (func $test-func-b (type $0)
-;; CHECK-NEXT:  )
-;; CHECK-NEXT:  (@metadata.code.inline "\00")
-;; CHECK-NEXT:  (func $test-func-c (type $0)
-;; CHECK-NEXT:  )
-;; CHECK-NEXT:  (@metadata.code.inline "\00")
-;; CHECK-NEXT:  (func $test-func-d (type $0)
-;; CHECK-NEXT:  )
-;; CHECK-NEXT: )
 

test/lit/passes/vacuum-removable-if-unused-func.wast

@@ -1,7 +1,17 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
 ;; RUN: wasm-opt -all --vacuum %s -S -o - | filecheck %s
 
 ;; Test the function-level annotation of removable.if.unused.
 (module
+ ;; CHECK:      (@binaryen.removable.if.unused)
+ ;; CHECK-NEXT: (func $calls-marked (type $0) (param $x i32) (result i32)
+ ;; CHECK-NEXT:  (local.set $x
+ ;; CHECK-NEXT:   (call $calls-marked
+ ;; CHECK-NEXT:    (i32.const 1)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (i32.const 2)
+ ;; CHECK-NEXT: )
  (@binaryen.removable.if.unused)
  (func $calls-marked (param $x i32) (result i32)
   ;; The function is marked as removable if unused, and this is dropped, so optimize.
@@ -19,6 +29,19 @@
   (i32.const 2)
  )
 
+ ;; CHECK:      (func $calls-unmarked (type $0) (param $x i32) (result i32)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (call $calls-unmarked
+ ;; CHECK-NEXT:    (i32.const 0)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (local.set $x
+ ;; CHECK-NEXT:   (call $calls-unmarked
+ ;; CHECK-NEXT:    (i32.const 1)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (i32.const 2)
+ ;; CHECK-NEXT: )
  (func $calls-unmarked (param $x i32) (result i32)
   ;; As above, but unmarked with the hint. We change nothing here.
   (drop
@@ -34,6 +57,14 @@
   (i32.const 2)
  )
 
+ ;; CHECK:      (func $calls-other (type $0) (param $x i32) (result i32)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (call $calls-unmarked
+ ;; CHECK-NEXT:    (i32.const 1)
+ ;; CHECK-NEXT:   )
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT:  (i32.const 2)
+ ;; CHECK-NEXT: )
  (func $calls-other (param $x i32) (result i32)
   ;; As above, but calling another function, to check we look for annotations in
   ;; the right place. Both calls are dropped, and only the one to the marked
@@ -52,37 +83,4 @@
  )
 )
 
-;; CHECK:      (module
-;; CHECK-NEXT:  (type $0 (func (param i32) (result i32)))
-;; CHECK-NEXT:  (@binaryen.removable.if.unused)
-;; CHECK-NEXT:  (func $calls-marked (type $0) (param $x i32) (result i32)
-;; CHECK-NEXT:   (local.set $x
-;; CHECK-NEXT:    (call $calls-marked
-;; CHECK-NEXT:     (i32.const 1)
-;; CHECK-NEXT:    )
-;; CHECK-NEXT:   )
-;; CHECK-NEXT:   (i32.const 2)
-;; CHECK-NEXT:  )
-;; CHECK-NEXT:  (func $calls-unmarked (type $0) (param $x i32) (result i32)
-;; CHECK-NEXT:   (drop
-;; CHECK-NEXT:    (call $calls-unmarked
-;; CHECK-NEXT:     (i32.const 0)
-;; CHECK-NEXT:    )
-;; CHECK-NEXT:   )
-;; CHECK-NEXT:   (local.set $x
-;; CHECK-NEXT:    (call $calls-unmarked
-;; CHECK-NEXT:     (i32.const 1)
-;; CHECK-NEXT:    )
-;; CHECK-NEXT:   )
-;; CHECK-NEXT:   (i32.const 2)
-;; CHECK-NEXT:  )
-;; CHECK-NEXT:  (func $calls-other (type $0) (param $x i32) (result i32)
-;; CHECK-NEXT:   (drop
-;; CHECK-NEXT:    (call $calls-unmarked
-;; CHECK-NEXT:     (i32.const 1)
-;; CHECK-NEXT:    )
-;; CHECK-NEXT:   )
-;; CHECK-NEXT:   (i32.const 2)
-;; CHECK-NEXT:  )
-;; CHECK-NEXT: )
 

test/lit/removable-if-unused-func.wast

@@ -1,7 +1,14 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py and should not be edited.
 ;; RUN: wasm-opt -all             %s -S -o - | filecheck %s
 ;; RUN: wasm-opt -all --roundtrip %s -S -o - | filecheck %s
 
 (module
+  ;; CHECK:      (@binaryen.removable.if.unused)
+  ;; CHECK-NEXT: (func $func-annotation (type $0)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (i32.const 0)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
   (@binaryen.removable.if.unused)
   (func $func-annotation
     ;; The annotation here is on the function.
@@ -11,13 +18,4 @@
   )
 )
 
-;; CHECK: (module
-;; CHECK-NEXT:  (type $0 (func))
-;; CHECK-NEXT:  (@binaryen.removable.if.unused)
-;; CHECK-NEXT:  (func $func-annotation
-;; CHECK-NEXT:   (drop
-;; CHECK-NEXT:    (i32.const 0)
-;; CHECK-NEXT:   )
-;; CHECK-NEXT:  )
-;; CHECK-NEXT: )