[llvm] [WebAssembly] Support annotation for try_table (PR #109029)

[Heejin Ahn via llvm-commits]

https://github.com/aheejin created https://github.com/llvm/llvm-project/pull/109029

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

>From 0ac2a9a71e28300dbdd717f9bfd07a48f7cd319c Mon Sep 17 00:00:00 2001
From: Heejin Ahn <aheejin at gmail.com>
Date: Mon, 16 Sep 2024 03:48:00 +0000
Subject: [PATCH] [WebAssembly] Support annotation for try_table

This adds support for annotations (`down to labelN`) for `try_table`.
---
 .../MCTargetDesc/WebAssemblyInstPrinter.cpp   | 45 ++++++++++++++-----
 llvm/test/MC/WebAssembly/annotations.s        | 35 +++++++++++++++
 2 files changed, 69 insertions(+), 11 deletions(-)

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