[WebAssembly] Allow try_table to target loops in AsmTypeCheck (#111432)

4 files changed, 72 insertions, 6 deletions

diff --git a/llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmTypeCheck.cpp b/llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmTypeCheck.cpp
index bfe6996..cc4338f 100644
--- a/llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmTypeCheck.cpp
+++ b/llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmTypeCheck.cpp
@@ -371,11 +371,16 @@ bool WebAssemblyAsmTypeCheck::checkTryTable(SMLoc ErrorLoc,
     if (Level < BlockInfoStack.size()) {
       const auto &DestBlockInfo =
           BlockInfoStack[BlockInfoStack.size() - Level - 1];
-      if (compareTypes(SentTypes, DestBlockInfo.Sig.Returns)) {
+      ArrayRef<wasm::ValType> DestTypes;
+      if (DestBlockInfo.IsLoop)
+        DestTypes = DestBlockInfo.Sig.Params;
+      else
+        DestTypes = DestBlockInfo.Sig.Returns;
+      if (compareTypes(SentTypes, DestTypes)) {
         std::string ErrorMsg =
             ErrorMsgBase + "type mismatch, catch tag type is " +
             getTypesString(SentTypes) + ", but destination's type is " +
-            getTypesString(DestBlockInfo.Sig.Returns);
+            getTypesString(DestTypes);
         Error |= typeError(ErrorLoc, ErrorMsg);
       }
     } else {
diff --git a/llvm/test/MC/WebAssembly/annotations.s b/llvm/test/MC/WebAssembly/annotations.s
index 2a5bea2..f761ef3 100644
--- a/llvm/test/MC/WebAssembly/annotations.s
+++ b/llvm/test/MC/WebAssembly/annotations.s
@@ -7,7 +7,7 @@
   .section .text.test_annotation,"",@
   .type    test_annotation,@function
 test_annotation:
-  .functype   test_annotation () -> ()
+  .functype   test_annotation (exnref) -> ()
   .tagtype  __cpp_exception i32
   .tagtype  __c_longjmp i32
   try
@@ -54,8 +54,18 @@ test_annotation:
     return
   end_block
   drop
-  end_function
 
+  i32.const 0
+  loop (i32) -> ()
+    local.get 0
+    loop (exnref) -> ()
+      try_table (catch __cpp_exception 1) (catch_all_ref 0)
+      end_try_table
+      drop
+    end_loop
+    drop
+  end_loop
+  end_function
 
 # CHECK:      test_annotation:
 # CHECK:        try
@@ -105,5 +115,16 @@ test_annotation:
 # CHECK-NEXT:   return
 # CHECK-NEXT:   end_block                               # label7:
 # CHECK-NEXT:   drop
-# CHECK-NEXT:   end_function
 
+# CHECK:        i32.const       0
+# CHECK-NEXT:   loop            (i32) -> ()                     # label12:
+# CHECK-NEXT:   local.get       0
+# CHECK-NEXT:   loop            (exnref) -> ()                  # label13:
+# CHECK-NEXT:   try_table        (catch __cpp_exception 1) (catch_all_ref 0) # 1: up to label12
+# CHECK-NEXT:                                 # 0: up to label13
+# CHECK-NEXT:   end_try_table                           # label14:
+# CHECK-NEXT:   drop
+# CHECK-NEXT:   end_loop
+# CHECK-NEXT:   drop
+# CHECK-NEXT:   end_loop
+# CHECK-NEXT:   end_function
diff --git a/llvm/test/MC/WebAssembly/eh-assembly.s b/llvm/test/MC/WebAssembly/eh-assembly.s
index 38cda10..31dfce5 100644
--- a/llvm/test/MC/WebAssembly/eh-assembly.s
+++ b/llvm/test/MC/WebAssembly/eh-assembly.s
@@ -7,7 +7,7 @@
   .functype  foo () -> ()
 
 eh_test:
-  .functype  eh_test () -> ()
+  .functype  eh_test (exnref) -> ()
 
   # try_table with all four kinds of catch clauses
   block exnref
@@ -82,6 +82,18 @@ eh_test:
   end_try_table
   drop
   drop
+
+  # try_table targeting loops
+  i32.const 0
+  loop (i32) -> ()
+    local.get 0
+    loop (exnref) -> ()
+      try_table (catch __cpp_exception 1) (catch_all_ref 0)
+      end_try_table
+      drop
+    end_loop
+    drop
+  end_loop
   end_function
 
 eh_legacy_test:
@@ -203,6 +215,17 @@ eh_legacy_test:
 # CHECK-NEXT:    drop
 # CHECK-NEXT:    drop
 
+# CHECK:         i32.const       0
+# CHECK-NEXT:    loop            (i32) -> ()
+# CHECK-NEXT:    local.get       0
+# CHECK-NEXT:    loop            (exnref) -> ()
+# CHECK-NEXT:    try_table        (catch __cpp_exception 1) (catch_all_ref 0)
+# CHECK:         end_try_table
+# CHECK-NEXT:    drop
+# CHECK-NEXT:    end_loop
+# CHECK-NEXT:    drop
+# CHECK-NEXT:    end_loop
+
 # CHECK:       eh_legacy_test:
 # CHECK:         try
 # CHECK-NEXT:    i32.const       3
diff --git a/llvm/test/MC/WebAssembly/type-checker-errors.s b/llvm/test/MC/WebAssembly/type-checker-errors.s
index c1c8209..9aa6523 100644
--- a/llvm/test/MC/WebAssembly/type-checker-errors.s
+++ b/llvm/test/MC/WebAssembly/type-checker-errors.s
@@ -966,4 +966,21 @@ eh_test:
     end_block
   end_block
   drop
+
+  loop
+  i32.const 0
+    loop (i32) -> ()
+      loop (i32) -> ()
+        loop
+# CHECK: :[[@LINE+4]]:11: error: try_table: catch index 0: type mismatch, catch tag type is [i32], but destination's type is []
+# CHECK: :[[@LINE+3]]:11: error: try_table: catch index 1: type mismatch, catch tag type is [i32, exnref], but destination's type is [i32]
+# CHECK: :[[@LINE+2]]:11: error: try_table: catch index 2: type mismatch, catch tag type is [], but destination's type is [i32]
+# CHECK: :[[@LINE+1]]:11: error: try_table: catch index 3: type mismatch, catch tag type is [exnref], but destination's type is []
+          try_table (catch __cpp_exception 0) (catch_ref __cpp_exception 1) (catch_all 2) (catch_all_ref 3)
+          end_try_table
+        end_loop
+        drop
+      end_loop
+    end_loop
+  end_loop
   end_function