[WebAssembly] Support annotation for try_table

[aheejin]

This adds support for annotations (down to labelN) for try_table.

[llvmbot]

@llvm/pr-subscribers-mc

@llvm/pr-subscribers-backend-webassembly

Author: Heejin Ahn (aheejin)

Changes

This adds support for annotations (down to labelN) for try_table.

---

Full diff: https://github.com/llvm/llvm-project/pull/109029.diff

2 Files Affected:

• (modified) llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.cpp (+34-11)
• (modified) llvm/test/MC/WebAssembly/annotations.s (+35)

diff --git a/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.cpp b/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.cpp
index 215722204ba4b3..4c29b59b3302e4 100644
--- a/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.cpp
+++ b/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.cpp
@@ -110,6 +110,20 @@ void WebAssemblyInstPrinter::printInst(const MCInst *MI, uint64_t Address,
   // Print any added annotation.
   printAnnotation(OS, Annot);
 
+  auto PrintBranchAnnotation = [&](const MCOperand &Op,
+                                   SmallSet<uint64_t, 8> &Printed) {
+    uint64_t Depth = Op.getImm();
+    if (!Printed.insert(Depth).second)
+      return;
+    if (Depth >= ControlFlowStack.size()) {
+      printAnnotation(OS, "Invalid depth argument!");
+    } else {
+      const auto &Pair = ControlFlowStack.rbegin()[Depth];
+      printAnnotation(OS, utostr(Depth) + ": " + (Pair.second ? "up" : "down") +
+                              " to label" + utostr(Pair.first));
+    }
+  };
+
   if (CommentStream) {
     // Observe any effects on the control flow stack, for use in annotating
     // control flow label references.
@@ -136,6 +150,23 @@ void WebAssemblyInstPrinter::printInst(const MCInst *MI, uint64_t Address,
       EHInstStack.push_back(TRY);
       return;
 
+    case WebAssembly::TRY_TABLE:
+    case WebAssembly::TRY_TABLE_S: {
+      SmallSet<uint64_t, 8> Printed;
+      unsigned OpIdx = 1;
+      const MCOperand &Op = MI->getOperand(OpIdx++);
+      unsigned NumCatches = Op.getImm();
+      for (unsigned I = 0; I < NumCatches; I++) {
+        int64_t CatchOpcode = MI->getOperand(OpIdx++).getImm();
+        if (CatchOpcode == wasm::WASM_OPCODE_CATCH ||
+            CatchOpcode == wasm::WASM_OPCODE_CATCH_REF)
+          OpIdx++; // Skip tag
+        PrintBranchAnnotation(MI->getOperand(OpIdx++), Printed);
+      }
+      ControlFlowStack.push_back(std::make_pair(ControlFlowCounter++, false));
+      return;
+    }
+
     case WebAssembly::END_LOOP:
     case WebAssembly::END_LOOP_S:
       if (ControlFlowStack.empty()) {
@@ -147,6 +178,8 @@ void WebAssemblyInstPrinter::printInst(const MCInst *MI, uint64_t Address,
 
     case WebAssembly::END_BLOCK:
     case WebAssembly::END_BLOCK_S:
+    case WebAssembly::END_TRY_TABLE:
+    case WebAssembly::END_TRY_TABLE_S:
       if (ControlFlowStack.empty()) {
         printAnnotation(OS, "End marker mismatch!");
       } else {
@@ -251,17 +284,7 @@ void WebAssemblyInstPrinter::printInst(const MCInst *MI, uint64_t Address,
         if (!MI->getOperand(I).isImm())
           continue;
       }
-      uint64_t Depth = MI->getOperand(I).getImm();
-      if (!Printed.insert(Depth).second)
-        continue;
-      if (Depth >= ControlFlowStack.size()) {
-        printAnnotation(OS, "Invalid depth argument!");
-      } else {
-        const auto &Pair = ControlFlowStack.rbegin()[Depth];
-        printAnnotation(OS, utostr(Depth) + ": " +
-                                (Pair.second ? "up" : "down") + " to label" +
-                                utostr(Pair.first));
-      }
+      PrintBranchAnnotation(MI->getOperand(I), Printed);
     }
   }
 }
diff --git a/llvm/test/MC/WebAssembly/annotations.s b/llvm/test/MC/WebAssembly/annotations.s
index b1f97daccccd65..3e727591afa942 100644
--- a/llvm/test/MC/WebAssembly/annotations.s
+++ b/llvm/test/MC/WebAssembly/annotations.s
@@ -33,6 +33,22 @@ test_annotation:
       rethrow   0
     end_try
   end_try
+
+  block exnref
+    block
+      block () -> (i32, exnref)
+        block i32
+          try_table (catch __cpp_exception 0) (catch_ref __c_longjmp 1) (catch_all 2) (catch_all_ref 3)
+          end_try_table
+          return
+        end_block
+        return
+      end_block
+      return
+    end_block
+    return
+  end_block
+  drop
   end_function
 
 
@@ -61,5 +77,24 @@ test_annotation:
 # CHECK-NEXT:   rethrow   0               # to caller
 # CHECK-NEXT:   end_try                   # label3:
 # CHECK-NEXT:   end_try                   # label0:
+
+# CHECK:        block           exnref
+# CHECK-NEXT:   block
+# CHECK-NEXT:   block           () -> (i32, exnref)
+# CHECK-NEXT:   block           i32
+# CHECK-NEXT:   try_table        (catch __cpp_exception 0) (catch_ref __c_longjmp 1) (catch_all 2) (catch_all_ref 3) # 0: down to label10
+# CHECK-NEXT:                             # 1: down to label9
+# CHECK-NEXT:                             # 2: down to label8
+# CHECK-NEXT:                             # 3: down to label7
+# CHECK-NEXT:   end_try_table                           # label11:
+# CHECK-NEXT:   return
+# CHECK-NEXT:   end_block                               # label10:
+# CHECK-NEXT:   return
+# CHECK-NEXT:   end_block                               # label9:
+# CHECK-NEXT:   return
+# CHECK-NEXT:   end_block                               # label8:
+# CHECK-NEXT:   return
+# CHECK-NEXT:   end_block                               # label7:
+# CHECK-NEXT:   drop
 # CHECK-NEXT:   end_function