1 files changed, 45 insertions, 25 deletions

diff --git a/llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp b/llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp
index 7bafa53af2af..1cba0843f891 100644
--- a/llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp
+++ b/llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp
@@ -210,6 +210,7 @@ class WebAssemblyAsmParser final : public MCTargetAsmParser {
   // guarantee that correct order.
   enum ParserState {
     FileStart,
+    FunctionLabel,
     FunctionStart,
     FunctionLocals,
     Instructions,
@@ -273,11 +274,11 @@ public:
 #include "WebAssemblyGenAsmMatcher.inc"
 
   // TODO: This is required to be implemented, but appears unused.
-  bool ParseRegister(unsigned & /*RegNo*/, SMLoc & /*StartLoc*/,
+  bool parseRegister(MCRegister & /*RegNo*/, SMLoc & /*StartLoc*/,
                      SMLoc & /*EndLoc*/) override {
-    llvm_unreachable("ParseRegister is not implemented.");
+    llvm_unreachable("parseRegister is not implemented.");
   }
-  OperandMatchResultTy tryParseRegister(unsigned & /*RegNo*/,
+  OperandMatchResultTy tryParseRegister(MCRegister & /*RegNo*/,
                                         SMLoc & /*StartLoc*/,
                                         SMLoc & /*EndLoc*/) override {
     llvm_unreachable("tryParseRegister is not implemented.");
@@ -287,8 +288,8 @@ public:
     return Parser.Error(Tok.getLoc(), Msg + Tok.getString());
   }
 
-  bool error(const Twine &Msg) {
-    return Parser.Error(Lexer.getTok().getLoc(), Msg);
+  bool error(const Twine &Msg, SMLoc Loc = SMLoc()) {
+    return Parser.Error(Loc.isValid() ? Loc : Lexer.getTok().getLoc(), Msg);
   }
 
   void addSignature(std::unique_ptr<wasm::WasmSignature> &&Sig) {
@@ -336,11 +337,12 @@ public:
     return false;
   }
 
-  bool ensureEmptyNestingStack() {
+  bool ensureEmptyNestingStack(SMLoc Loc = SMLoc()) {
     auto Err = !NestingStack.empty();
     while (!NestingStack.empty()) {
       error(Twine("Unmatched block construct(s) at function end: ") +
-            nestingString(NestingStack.back().NT).first);
+                nestingString(NestingStack.back().NT).first,
+            Loc);
       NestingStack.pop_back();
     }
     return Err;
@@ -835,7 +837,8 @@ public:
       auto ElemTypeName = expectIdent();
       if (ElemTypeName.empty())
         return true;
-      Optional<wasm::ValType> ElemType = WebAssembly::parseType(ElemTypeName);
+      std::optional<wasm::ValType> ElemType =
+          WebAssembly::parseType(ElemTypeName);
       if (!ElemType)
         return error("Unknown type in .tabletype directive: ", ElemTypeTok);
 
@@ -864,12 +867,24 @@ public:
         return true;
       auto WasmSym = cast<MCSymbolWasm>(Ctx.getOrCreateSymbol(SymName));
       if (WasmSym->isDefined()) {
-        // This .functype indicates a start of a function.
-        if (ensureEmptyNestingStack())
-          return true;
+        // We push 'Function' either when a label is parsed or a .functype
+        // directive is parsed. The reason it is not easy to do this uniformly
+        // in a single place is,
+        // 1. We can't do this at label parsing time only because there are
+        //    cases we don't have .functype directive before a function label,
+        //    in which case we don't know if the label is a function at the time
+        //    of parsing.
+        // 2. We can't do this at .functype parsing time only because we want to
+        //    detect a function started with a label and not ended correctly
+        //    without encountering a .functype directive after the label.
+        if (CurrentState != FunctionLabel) {
+          // This .functype indicates a start of a function.
+          if (ensureEmptyNestingStack())
+            return true;
+          push(Function);
+        }
         CurrentState = FunctionStart;
         LastFunctionLabel = WasmSym;
-        push(Function);
       }
       auto Signature = std::make_unique<wasm::WasmSignature>();
       if (parseSignature(Signature.get()))
@@ -1057,7 +1072,7 @@ public:
     llvm_unreachable("Implement any new match types added!");
   }
 
-  void doBeforeLabelEmit(MCSymbol *Symbol) override {
+  void doBeforeLabelEmit(MCSymbol *Symbol, SMLoc IDLoc) override {
     // Code below only applies to labels in text sections.
     auto CWS = cast<MCSectionWasm>(getStreamer().getCurrentSection().first);
     if (!CWS || !CWS->getKind().isText())
@@ -1067,7 +1082,7 @@ public:
     // Unlike other targets, we don't allow data in text sections (labels
     // declared with .type @object).
     if (WasmSym->getType() == wasm::WASM_SYMBOL_TYPE_DATA) {
-      Parser.Error(Parser.getTok().getLoc(),
+      Parser.Error(IDLoc,
                    "Wasm doesn\'t support data symbols in text sections");
       return;
     }
@@ -1099,6 +1114,22 @@ public:
     // Also generate DWARF for this section if requested.
     if (getContext().getGenDwarfForAssembly())
       getContext().addGenDwarfSection(WS);
+
+    if (WasmSym->isFunction()) {
+      // We give the location of the label (IDLoc) here, because otherwise the
+      // lexer's next location will be used, which can be confusing. For
+      // example:
+      //
+      // test0: ; This function does not end properly
+      //   ...
+      //
+      // test1: ; We would like to point to this line for error
+      //   ...  . Not this line, which can contain any instruction
+      ensureEmptyNestingStack(IDLoc);
+      CurrentState = FunctionLabel;
+      LastFunctionLabel = Symbol;
+      push(Function);
+    }
   }
 
   void onEndOfFunction(SMLoc ErrorLoc) {
@@ -1106,17 +1137,6 @@ public:
       TC.endOfFunction(ErrorLoc);
     // Reset the type checker state.
     TC.Clear();
-
-    // Automatically output a .size directive, so it becomes optional for the
-    // user.
-    if (!LastFunctionLabel) return;
-    auto TempSym = getContext().createLinkerPrivateTempSymbol();
-    getStreamer().emitLabel(TempSym);
-    auto Start = MCSymbolRefExpr::create(LastFunctionLabel, getContext());
-    auto End = MCSymbolRefExpr::create(TempSym, getContext());
-    auto Expr =
-        MCBinaryExpr::create(MCBinaryExpr::Sub, End, Start, getContext());
-    getStreamer().emitELFSize(LastFunctionLabel, Expr);
   }
 
   void onEndOfFile() override { ensureEmptyNestingStack(); }