[WebAssembly] support "return" and unreachable code in asm type checker

[aardappel]

Summary

To support return (it not being supported well was the ground cause for
https://github.com/WebAssembly/wasi-sdk/issues/200) we also have to have
at least a basic notion of unreachable, which in this case just means to stop
type checking until there is an end_block (an incoming control flow edge).
This is conservative (may miss on some type checking opportunities) but is
simple and an improvement over what we had before.

[sbc100]

lgtm % clang-format stuff.

lld/test/wasm/lto/thinlto-thin-archive-collision.ll

3 ↗

(On Diff #383866)

Seems unrelated?

[aardappel]

lld/test/wasm/lto/thinlto-thin-archive-collision.ll

3 ↗

(On Diff #383866)

It was needed to make tests pass on my machine, and since it is a tiny change it seemed ok

[sbc100]

lld/test/wasm/lto/thinlto-thin-archive-collision.ll

3 ↗

(On Diff #383866)

Hmm. That is strange. Since it completely unrelated I don't think you should land this as part of this change.

[dschuff]

Are there any type checker tests that are expected to fail?

llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmTypeCheck.cpp

77

[aardappel]

Removed unrelated fix and formatting

llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmParser.cpp

void doBeforeLabelEmit(MCSymbol *Symbol) override {
     getStreamer().SwitchSection(WS);
     // Also generate DWARF for this section if requested.
     if (getContext().getGenDwarfForAssembly())
       getContext().addGenDwarfSection(WS);
   }
 
   void onEndOfFunction(SMLoc ErrorLoc) {
     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());

llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmTypeCheck.h

class WebAssemblyAsmTypeCheck final {
   MCAsmParser &Parser;
   const MCInstrInfo &MII;
 
   SmallVector<wasm::ValType, 8> Stack;
   SmallVector<wasm::ValType, 16> LocalTypes;
   SmallVector<wasm::ValType, 4> ReturnTypes;
   wasm::WasmSignature LastSig;
   bool TypeErrorThisFunction = false;
+  bool Unreachable = false;
   bool is64;
 
-  void Clear() {
-    Stack.clear();
-    LocalTypes.clear();
-    ReturnTypes.clear();
-    TypeErrorThisFunction = false;
-  }
-
   void dumpTypeStack(Twine Msg);
   bool typeError(SMLoc ErrorLoc, const Twine &Msg);
   bool popType(SMLoc ErrorLoc, Optional<wasm::ValType> EVT);
   bool getLocal(SMLoc ErrorLoc, const MCInst &Inst, wasm::ValType &Type);
   bool checkEnd(SMLoc ErrorLoc);
   bool checkSig(SMLoc ErrorLoc, const wasm::WasmSignature &Sig);
   bool getSymRef(SMLoc ErrorLoc, const MCInst &Inst,
                  const MCSymbolRefExpr *&SymRef);
   bool getGlobal(SMLoc ErrorLoc, const MCInst &Inst, wasm::ValType &Type);
 
 public:
   WebAssemblyAsmTypeCheck(MCAsmParser &Parser, const MCInstrInfo &MII, bool is64);
 
   void funcDecl(const wasm::WasmSignature &Sig);
   void localDecl(const SmallVector<wasm::ValType, 4> &Locals);
   void setLastSig(const wasm::WasmSignature &Sig) { LastSig = Sig; }
-  void endOfFunction(SMLoc ErrorLoc);
+  bool endOfFunction(SMLoc ErrorLoc);
   bool typeCheck(SMLoc ErrorLoc, const MCInst &Inst);
+
+  void Clear() {
+    Stack.clear();
+    LocalTypes.clear();
+    ReturnTypes.clear();
+    TypeErrorThisFunction = false;
+    Unreachable = false;
+  }
 };
 
 } // end namespace llvm
 
 #endif  // LLVM_LIB_TARGET_WEBASSEMBLY_ASMPARSER_TYPECHECK_H

llvm/lib/Target/WebAssembly/AsmParser/WebAssemblyAsmTypeCheck.cpp

void WebAssemblyAsmTypeCheck::dumpTypeStack(Twine Msg) {
   });
 }
 
 bool WebAssemblyAsmTypeCheck::typeError(SMLoc ErrorLoc, const Twine &Msg) {
   // Once you get one type error in a function, it will likely trigger more
   // which are mostly not helpful.
   if (TypeErrorThisFunction)
     return true;
+  // If we're currently in unreachable code, we surpress errors as well.
+  if (Unreachable)
+    return true;
   TypeErrorThisFunction = true;
   dumpTypeStack("current stack: ");
   return Parser.Error(ErrorLoc, Msg);
 }
 
 bool WebAssemblyAsmTypeCheck::popType(SMLoc ErrorLoc,
                                       Optional<wasm::ValType> EVT) {
   if (Stack.empty()) {
case wasm::WASM_SYMBOL_TYPE_DATA:
     LLVM_FALLTHROUGH;
   default:
     return typeError(ErrorLoc, StringRef("symbol ") + WasmSym->getName() +
                                     " missing .globaltype");
   }
   return false;
 }
 
-void WebAssemblyAsmTypeCheck::endOfFunction(SMLoc ErrorLoc) {
+bool WebAssemblyAsmTypeCheck::endOfFunction(SMLoc ErrorLoc) {
   // Check the return types.
   for (auto RVT : llvm::reverse(ReturnTypes)) {
-    popType(ErrorLoc, RVT);
+    if (popType(ErrorLoc, RVT))
+      return true;
   }
   if (!Stack.empty()) {
-    typeError(ErrorLoc,
-              std::to_string(Stack.size()) + " superfluous return values");
+    return typeError(ErrorLoc, std::to_string(Stack.size()) +
+                                   " superfluous return values");
   }
-  // Reset the type checker state.
-  Clear();
+  Unreachable = true;
+  return false;
 }
 
 bool WebAssemblyAsmTypeCheck::typeCheck(SMLoc ErrorLoc, const MCInst &Inst) {
   auto Opc = Inst.getOpcode();
   auto Name = GetMnemonic(Opc);
   dumpTypeStack("typechecking " + Name + ": ");
   wasm::ValType Type;
   if (Name == "local.get") {
if (popType(ErrorLoc, Type))
       return true;
   } else if (Name == "drop") {
     if (popType(ErrorLoc, {}))
       return true;
   } else if (Name == "end_block" || Name == "end_loop" || Name == "end_if" ||
              Name == "else" || Name == "end_try") {
     if (checkEnd(ErrorLoc))
       return true;
+    if (Name == "end_block")
+      Unreachable = false;
+  } else if (Name == "return") {
+    if (endOfFunction(ErrorLoc))
+      return true;
   } else if (Name == "call_indirect" || Name == "return_call_indirect") {
     // Function value.
     if (popType(ErrorLoc, wasm::ValType::I32)) return true;
     if (checkSig(ErrorLoc, LastSig)) return true;
+    if (Name == "return_call_indirect" && endOfFunction(ErrorLoc))
+      return true;
   } else if (Name == "call" || Name == "return_call") {
     const MCSymbolRefExpr *SymRef;
     if (getSymRef(ErrorLoc, Inst, SymRef))
       return true;
     auto WasmSym = cast<MCSymbolWasm>(&SymRef->getSymbol());
     auto Sig = WasmSym->getSignature();
     if (!Sig || WasmSym->getType() != wasm::WASM_SYMBOL_TYPE_FUNCTION)
       return typeError(ErrorLoc, StringRef("symbol ") + WasmSym->getName() +
                                       " missing .functype");
     if (checkSig(ErrorLoc, *Sig)) return true;
+    if (Name == "return_call" && endOfFunction(ErrorLoc))
+      return true;
   } else if (Name == "catch") {
     const MCSymbolRefExpr *SymRef;
     if (getSymRef(ErrorLoc, Inst, SymRef))
       return true;
     const auto *WasmSym = cast<MCSymbolWasm>(&SymRef->getSymbol());
     const auto *Sig = WasmSym->getSignature();
     if (!Sig || WasmSym->getType() != wasm::WASM_SYMBOL_TYPE_TAG)
       return typeError(ErrorLoc, StringRef("symbol ") + WasmSym->getName() +
                                      " missing .tagtype");
     // catch instruction pushes values whose types are specified in the tag's
     // "params" part
     Stack.insert(Stack.end(), Sig->Params.begin(), Sig->Params.end());
   } else if (Name == "ref.null") {
     auto VT = static_cast<wasm::ValType>(Inst.getOperand(0).getImm());
     Stack.push_back(VT);
+  } else if (Name == "unreachable") {
+    Unreachable = true;
   } else {
     // The current instruction is a stack instruction which doesn't have
     // explicit operands that indicate push/pop types, so we get those from
     // the register version of the same instruction.
     auto RegOpc = WebAssembly::getRegisterOpcode(Opc);
     assert(RegOpc != -1 && "Failed to get register version of MC instruction");
     const auto &II = MII.get(RegOpc);
     // First pop all the uses off the stack and check them.

llvm/test/MC/WebAssembly/basic-assembly.s

-# RUN: llvm-mc -triple=wasm32-unknown-unknown -mattr=+reference-types,atomics,+simd128,+nontrapping-fptoint,+exception-handling < %s | FileCheck %s
+# RUN: llvm-mc -triple=wasm32-unknown-unknown -mattr=+tail-call,+reference-types,atomics,+simd128,+nontrapping-fptoint,+exception-handling < %s | FileCheck %s
 # Check that it converts to .o without errors, but don't check any output:
-# RUN: llvm-mc -triple=wasm32-unknown-unknown -filetype=obj -mattr=+reference-types,+atomics,+simd128,+nontrapping-fptoint,+exception-handling -o %t.o < %s
+# RUN: llvm-mc -triple=wasm32-unknown-unknown -filetype=obj -mattr=+tail-call,+reference-types,+atomics,+simd128,+nontrapping-fptoint,+exception-handling -o %t.o < %s
 
 .functype   something1 () -> ()
 .functype   something2 (i64) -> (i32, f64)
+.functype   something3 () -> (i32)
 .globaltype __stack_pointer, i32
 
 empty_func:
     .functype empty_func () -> ()
     end_function
 
 test0:
 # local labels can appear between label and its .functype.
.LBB0_2:
     i32.const   3
     end_block            # "switch" exit.
     if                   # void
     if          i32
     end_if
     else
     end_if
     drop
+    block       void
+    i32.const   2
+    return
+    end_block
+    block       void
+    return_call something3
+    end_block
+    block       void
+    i32.const   3
+    return_call_indirect () -> (i32)
+    end_block
     local.get   4
     local.get   5
     f32x4.add
     drop
     # Test correct parsing of instructions with / and : in them:
     # TODO: enable once instruction has been added.
     #i32x4.trunc_sat_f32x4_s
     f32.const   1.0

 # CHECK-NEXT:      i32.const   3
 # CHECK-NEXT:      end_block           # label2:
 # CHECK-NEXT:      if
 # CHECK-NEXT:      if          i32
 # CHECK-NEXT:      end_if
 # CHECK-NEXT:      else
 # CHECK-NEXT:      end_if
 # CHECK-NEXT:      drop
+# CHECK-NEXT:      block
+# CHECK-NEXT:      i32.const   2
+# CHECK-NEXT:      return
+# CHECK-NEXT:      end_block
+# CHECK-NEXT:      block
+# CHECK-NEXT:      return_call something3
+# CHECK-NEXT:      end_block
+# CHECK-NEXT:      block
+# CHECK-NEXT:      i32.const   3
+# CHECK-NEXT:      return_call_indirect __indirect_function_table, () -> (i32)
+# CHECK-NEXT:      end_block
 # CHECK-NEXT:      local.get   4
 # CHECK-NEXT:      local.get   5
 # CHECK-NEXT:      f32x4.add
 # CHECK-NEXT:      drop
 # CHECK-NEXT:      f32.const   0x1p0
 # CHECK-NEXT:      i32.trunc_f32_s
 # CHECK-NEXT:      try
 # CHECK-NEXT:      i32.atomic.load 0