[WebAssembly] Support assembly parsing for new EH (#108668)

aheejin · web-flow · commit defb8fb2c67f · 2024-09-17T11:35:32.000-07:00
This adds assembly parsing support for the new EH (exnref) proposal. `try_table` parsing is a little tricky because catch clause lists use `()` and the multivalue block return types also use `()`. This handles all combinations below: - No return type (void) + no catch list - No return type (void) + catch list - Single return type + no catch list - Single return type + catch list - Multivalue return type + no catch list - Multivalue return type + catch list This does not include AsmTypeCheck support yet. That's the reason why this adds a new test file and use `--no-type-check` in the command line. After the type checker is added as a follow-up, I plan to merge https://github.com/llvm/llvm-project/blob/main/llvm/test/MC/WebAssembly/eh-assembly-legacy.s with this file. (Turning on `-mattr=+exception-handling` adds support for all legacy and new EH instructions in the assembly. `-wasm-enable-exnref` in `llc` only controls which instructions to generate and it doesn't affect `llvm-mc` and assembly parsing.)
diff --git a/llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp b/llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp
@@ -69,12 +69,23 @@ struct WebAssemblyOperand : public MCParsedAsmOperand {
     std::vector<unsigned> List;
   };
 
+  struct CaLOpElem {
+    uint8_t Opcode;
+    const MCExpr *Tag;
+    unsigned Dest;
+  };
+
+  struct CaLOp {
+    std::vector<CaLOpElem> List;
+  };
+
   union {
     struct TokOp Tok;
     struct IntOp Int;
     struct FltOp Flt;
     struct SymOp Sym;
     struct BrLOp BrL;
+    struct CaLOp CaL;
   };
 
   WebAssemblyOperand(SMLoc Start, SMLoc End, TokOp T)
@@ -85,12 +96,16 @@ struct WebAssemblyOperand : public MCParsedAsmOperand {
       : Kind(Float), StartLoc(Start), EndLoc(End), Flt(F) {}
   WebAssemblyOperand(SMLoc Start, SMLoc End, SymOp S)
       : Kind(Symbol), StartLoc(Start), EndLoc(End), Sym(S) {}
-  WebAssemblyOperand(SMLoc Start, SMLoc End)
-      : Kind(BrList), StartLoc(Start), EndLoc(End), BrL() {}
+  WebAssemblyOperand(SMLoc Start, SMLoc End, BrLOp B)
+      : Kind(BrList), StartLoc(Start), EndLoc(End), BrL(B) {}
+  WebAssemblyOperand(SMLoc Start, SMLoc End, CaLOp C)
+      : Kind(CatchList), StartLoc(Start), EndLoc(End), CaL(C) {}
 
   ~WebAssemblyOperand() {
     if (isBrList())
       BrL.~BrLOp();
+    if (isCatchList())
+      CaL.~CaLOp();
   }
 
   bool isToken() const override { return Kind == Token; }
@@ -153,7 +168,15 @@ struct WebAssemblyOperand : public MCParsedAsmOperand {
   }
 
   void addCatchListOperands(MCInst &Inst, unsigned N) const {
-    // TODO
+    assert(N == 1 && isCatchList() && "Invalid CatchList!");
+    Inst.addOperand(MCOperand::createImm(CaL.List.size()));
+    for (auto Ca : CaL.List) {
+      Inst.addOperand(MCOperand::createImm(Ca.Opcode));
+      if (Ca.Opcode == wasm::WASM_OPCODE_CATCH ||
+          Ca.Opcode == wasm::WASM_OPCODE_CATCH_REF)
+        Inst.addOperand(MCOperand::createExpr(Ca.Tag));
+      Inst.addOperand(MCOperand::createImm(Ca.Dest));
+    }
   }
 
   void print(raw_ostream &OS) const override {
@@ -174,7 +197,7 @@ struct WebAssemblyOperand : public MCParsedAsmOperand {
       OS << "BrList:" << BrL.List.size();
       break;
     case CatchList:
-      // TODO
+      OS << "CaList:" << CaL.List.size();
       break;
     }
   }
@@ -228,6 +251,7 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
     Loop,
     Try,
     CatchAll,
+    TryTable,
     If,
     Else,
     Undefined,
@@ -304,6 +328,8 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
       return {"try", "end_try/delegate"};
     case CatchAll:
       return {"catch_all", "end_try"};
+    case TryTable:
+      return {"try_table", "end_try_table"};
     case If:
       return {"if", "end_if"};
     case Else:
@@ -571,6 +597,7 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
     // proper nesting.
     bool ExpectBlockType = false;
     bool ExpectFuncType = false;
+    bool ExpectCatchList = false;
     std::unique_ptr<WebAssemblyOperand> FunctionTable;
     if (Name == "block") {
       push(Block);
@@ -593,12 +620,19 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
     } else if (Name == "catch_all") {
       if (popAndPushWithSameSignature(Name, Try, CatchAll))
         return true;
+    } else if (Name == "try_table") {
+      push(TryTable);
+      ExpectBlockType = true;
+      ExpectCatchList = true;
     } else if (Name == "end_if") {
       if (pop(Name, If, Else))
         return true;
     } else if (Name == "end_try") {
       if (pop(Name, Try, CatchAll))
         return true;
+    } else if (Name == "end_try_table") {
+      if (pop(Name, TryTable))
+        return true;
     } else if (Name == "delegate") {
       if (pop(Name, Try))
         return true;
@@ -622,7 +656,18 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
       ExpectFuncType = true;
     }
 
-    if (ExpectFuncType || (ExpectBlockType && Lexer.is(AsmToken::LParen))) {
+    // Returns true if the next tokens are a catch clause
+    auto PeekCatchList = [&]() {
+      if (Lexer.isNot(AsmToken::LParen))
+        return false;
+      AsmToken NextTok = Lexer.peekTok();
+      return NextTok.getKind() == AsmToken::Identifier &&
+             NextTok.getIdentifier().starts_with("catch");
+    };
+
+    // Parse a multivalue block type
+    if (ExpectFuncType ||
+        (Lexer.is(AsmToken::LParen) && ExpectBlockType && !PeekCatchList())) {
       // This has a special TYPEINDEX operand which in text we
       // represent as a signature, such that we can re-build this signature,
       // attach it to an anonymous symbol, which is what WasmObjectWriter
@@ -648,6 +693,23 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
           Loc.getLoc(), Loc.getEndLoc(), WebAssemblyOperand::SymOp{Expr}));
     }
 
+    // If we are expecting a catch clause list, try to parse it here.
+    //
+    // If there is a multivalue block return type before this catch list, it
+    // should have been parsed above. If there is no return type before
+    // encountering this catch list, this means the type is void.
+    // The case when there is a single block return value and then a catch list
+    // will be handled below in the 'while' loop.
+    if (ExpectCatchList && PeekCatchList()) {
+      if (ExpectBlockType) {
+        ExpectBlockType = false;
+        addBlockTypeOperand(Operands, NameLoc, WebAssembly::BlockType::Void);
+      }
+      if (parseCatchList(Operands))
+        return true;
+      ExpectCatchList = false;
+    }
+
     while (Lexer.isNot(AsmToken::EndOfStatement)) {
       auto &Tok = Lexer.getTok();
       switch (Tok.getKind()) {
@@ -661,7 +723,15 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
           if (BT == WebAssembly::BlockType::Invalid)
             return error("Unknown block type: ", Id);
           addBlockTypeOperand(Operands, NameLoc, BT);
+          ExpectBlockType = false;
           Parser.Lex();
+          // Now that we've parsed a single block return type, if we are
+          // expecting a catch clause list, try to parse it.
+          if (ExpectCatchList && PeekCatchList()) {
+            if (parseCatchList(Operands))
+              return true;
+            ExpectCatchList = false;
+          }
         } else {
           // Assume this identifier is a label.
           const MCExpr *Val;
@@ -703,8 +773,8 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
       }
       case AsmToken::LCurly: {
         Parser.Lex();
-        auto Op =
-            std::make_unique<WebAssemblyOperand>(Tok.getLoc(), Tok.getEndLoc());
+        auto Op = std::make_unique<WebAssemblyOperand>(
+            Tok.getLoc(), Tok.getEndLoc(), WebAssemblyOperand::BrLOp{});
         if (!Lexer.is(AsmToken::RCurly))
           for (;;) {
             Op->BrL.List.push_back(Lexer.getTok().getIntVal());
@@ -724,10 +794,18 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
           return true;
       }
     }
-    if (ExpectBlockType && Operands.size() == 1) {
-      // Support blocks with no operands as default to void.
+
+    // If we are still expecting to parse a block type or a catch list at this
+    // point, we set them to the default/empty state.
+
+    // Support blocks with no operands as default to void.
+    if (ExpectBlockType)
       addBlockTypeOperand(Operands, NameLoc, WebAssembly::BlockType::Void);
-    }
+    // If no catch list has been parsed, add an empty catch list operand.
+    if (ExpectCatchList)
+      Operands.push_back(std::make_unique<WebAssemblyOperand>(
+          NameLoc, NameLoc, WebAssemblyOperand::CaLOp{}));
+
     if (FunctionTable)
       Operands.push_back(std::move(FunctionTable));
     Parser.Lex();
@@ -752,6 +830,55 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
     return false;
   }
 
+  bool parseCatchList(OperandVector &Operands) {
+    auto Op = std::make_unique<WebAssemblyOperand>(
+        Lexer.getTok().getLoc(), SMLoc(), WebAssemblyOperand::CaLOp{});
+    SMLoc EndLoc;
+
+    while (Lexer.is(AsmToken::LParen)) {
+      if (expect(AsmToken::LParen, "("))
+        return true;
+
+      auto CatchStr = expectIdent();
+      if (CatchStr.empty())
+        return true;
+      uint8_t CatchOpcode =
+          StringSwitch<uint8_t>(CatchStr)
+              .Case("catch", wasm::WASM_OPCODE_CATCH)
+              .Case("catch_ref", wasm::WASM_OPCODE_CATCH_REF)
+              .Case("catch_all", wasm::WASM_OPCODE_CATCH_ALL)
+              .Case("catch_all_ref", wasm::WASM_OPCODE_CATCH_ALL_REF)
+              .Default(0xff);
+      if (CatchOpcode == 0xff)
+        return error(
+            "Expected catch/catch_ref/catch_all/catch_all_ref, instead got: " +
+            CatchStr);
+
+      const MCExpr *Tag = nullptr;
+      if (CatchOpcode == wasm::WASM_OPCODE_CATCH ||
+          CatchOpcode == wasm::WASM_OPCODE_CATCH_REF) {
+        if (Parser.parseExpression(Tag))
+          return error("Cannot parse symbol: ", Lexer.getTok());
+      }
+
+      auto &DestTok = Lexer.getTok();
+      if (DestTok.isNot(AsmToken::Integer))
+        return error("Expected integer constant, instead got: ", DestTok);
+      unsigned Dest = DestTok.getIntVal();
+      Parser.Lex();
+
+      EndLoc = Lexer.getTok().getEndLoc();
+      if (expect(AsmToken::RParen, ")"))
+        return true;
+
+      Op->CaL.List.push_back({CatchOpcode, Tag, Dest});
+    }
+
+    Op->EndLoc = EndLoc;
+    Operands.push_back(std::move(Op));
+    return false;
+  }
+
   bool CheckDataSection() {
     if (CurrentState != DataSection) {
       auto WS = cast<MCSectionWasm>(getStreamer().getCurrentSectionOnly());
diff --git a/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.cpp b/llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.cpp
@@ -378,7 +378,7 @@ void WebAssemblyInstPrinter::printCatchList(const MCInst *MI, unsigned OpNo,
     const MCSymbolRefExpr *TagExpr = nullptr;
     const MCSymbolWasm *TagSym = nullptr;
     if (Op.isExpr()) {
-      TagExpr = dyn_cast<MCSymbolRefExpr>(Op.getExpr());
+      TagExpr = cast<MCSymbolRefExpr>(Op.getExpr());
       TagSym = cast<MCSymbolWasm>(&TagExpr->getSymbol());
       O << TagSym->getName() << " ";
     } else {
diff --git a/llvm/test/MC/WebAssembly/basic-assembly-errors.s b/llvm/test/MC/WebAssembly/basic-assembly-errors.s
@@ -26,11 +26,37 @@ test0:
     catch_all
 # CHECK: error: Block construct type mismatch, expected: end_try, instead got: catch_all
     end
-# CHECK: Block construct type mismatch, expected: end_try, instead got: end_function
+
+# CHECK: error: Expected integer constant, instead got: )
+    try_table (catch __cpp_exception)
+    end_try_table
+
+    block
+# CHECK: error: invalid operand for instruction
+    try_table (catch_all 0) i32
+    i32.const 0
+    end_try_table
+    drop
+    end_block
+
+    block
+# CHECK: error: Expected identifier, got: )
+    try_table (catch_all 0) () -> (i32, i32)
+    i32.const 0
+    i32.const 0
+    end_try_table
+    drop
+    drop
+    end_block
+
+# CHECK: error: unknown type: not_catch
+    try_table (not_catch 0)
+
+# CHECK: Block construct type mismatch, expected: end_try_table, instead got: end_function
+    end_function
+# CHECK: error: Unmatched block construct(s) at function end: try_table
 # CHECK: error: Unmatched block construct(s) at function end: catch_all
 # CHECK: error: Unmatched block construct(s) at function end: loop
 # CHECK: error: Unmatched block construct(s) at function end: try
 # CHECK: error: Unmatched block construct(s) at function end: block
 # CHECK: error: Unmatched block construct(s) at function end: function
-    end_function
-
diff --git a/llvm/test/MC/WebAssembly/eh-assembly.s b/llvm/test/MC/WebAssembly/eh-assembly.s
