[WebAssembly] `AsmTypeCheck` support to br instr

[HerrCai0907]

Summary

This PR introduces the BrStack member to store the info about
loop, block, if and try. It can check whether br immediate number
out of range.

[aheejin]

Thanks for the simplified version!

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

144–146

I think it'd be better to start BrStack as the function return signature in the beginning than treating this as a special case. Currently we cannot check if the stack-remaining value types at the time of br matches the function return signature.

147

llvm/test/MC/WebAssembly/type-checker-errors.s

744

try~end_try without a catch/catch_all is malformed.

749–750

772

Can we add a few more tests on the depth?

1. When the depth is +1 than the max possible depth (including the function boundary) (incorrect)
2. When the depth targets the function boundary (correct)
3. When the depth targets the function boundary, but the signature doesn't match (incorrect)

[HerrCai0907]

push func body in BrStack

[HerrCai0907]

clean code

[HerrCai0907]

add more testcases

[aheejin]

llvm/test/MC/WebAssembly/type-checker-errors.s

744

Nit: I think tag_f32 (as in my previous suggestion) is better, because tag_i32 means introducing another type error within the catch. But we have another test for checking catch part below.

772

Ping. I don't think any of these were added?

[HerrCai0907]

add testcase

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

 
 namespace llvm {
 
 class WebAssemblyAsmTypeCheck final {
   MCAsmParser &Parser;
   const MCInstrInfo &MII;
 
   SmallVector<wasm::ValType, 8> Stack;
+  SmallVector<SmallVector<wasm::ValType, 4>, 8> BrStack;
   SmallVector<wasm::ValType, 16> LocalTypes;
   SmallVector<wasm::ValType, 4> ReturnTypes;
   wasm::WasmSignature LastSig;
   bool TypeErrorThisFunction = false;
   bool Unreachable = false;
   bool is64;
 
   void dumpTypeStack(Twine Msg);
   bool typeError(SMLoc ErrorLoc, const Twine &Msg);
   bool popType(SMLoc ErrorLoc, std::optional<wasm::ValType> EVT);
   bool popRefType(SMLoc ErrorLoc);
   bool getLocal(SMLoc ErrorLoc, const MCInst &Inst, wasm::ValType &Type);
   bool checkEnd(SMLoc ErrorLoc, bool PopVals = false);
+  bool checkBr(SMLoc ErrorLoc, size_t Level);
   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);
   bool getTable(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 SmallVectorImpl<wasm::ValType> &Locals);
   void setLastSig(const wasm::WasmSignature &Sig) { LastSig = Sig; }
   bool endOfFunction(SMLoc ErrorLoc);
   bool typeCheck(SMLoc ErrorLoc, const MCInst &Inst, OperandVector &Operands);
 
   void Clear() {
     Stack.clear();
+    BrStack.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

 WebAssemblyAsmTypeCheck::WebAssemblyAsmTypeCheck(MCAsmParser &Parser,
                                                  const MCInstrInfo &MII,
                                                  bool is64)
     : Parser(Parser), MII(MII), is64(is64) {}
 
 void WebAssemblyAsmTypeCheck::funcDecl(const wasm::WasmSignature &Sig) {
   LocalTypes.assign(Sig.Params.begin(), Sig.Params.end());
   ReturnTypes.assign(Sig.Returns.begin(), Sig.Returns.end());
+  BrStack.emplace_back(Sig.Returns.begin(), Sig.Returns.end());
 }
 
 void WebAssemblyAsmTypeCheck::localDecl(
     const SmallVectorImpl<wasm::ValType> &Locals) {
   LocalTypes.insert(LocalTypes.end(), Locals.begin(), Locals.end());
 }
 
 void WebAssemblyAsmTypeCheck::dumpTypeStack(Twine Msg) {
bool WebAssemblyAsmTypeCheck::getLocal(SMLoc ErrorLoc, const MCInst &Inst,
   auto Local = static_cast<size_t>(Inst.getOperand(0).getImm());
   if (Local >= LocalTypes.size())
     return typeError(ErrorLoc, StringRef("no local type specified for index ") +
                                    std::to_string(Local));
   Type = LocalTypes[Local];
   return false;
 }
 
+static std::optional<std::string>
+checkStackTop(const SmallVectorImpl<wasm::ValType> &ExpectedStackTop,
+              const SmallVectorImpl<wasm::ValType> &Got) {
+  for (size_t I = 0; I < ExpectedStackTop.size(); I++) {
+    auto EVT = ExpectedStackTop[I];
+    auto PVT = Got[Got.size() - ExpectedStackTop.size() + I];
+    if (PVT != EVT)
+      return std::string{"got "} + WebAssembly::typeToString(PVT) +
+             ", expected " + WebAssembly::typeToString(EVT);
+  }
+  return std::nullopt;
+}
+
+bool WebAssemblyAsmTypeCheck::checkBr(SMLoc ErrorLoc, size_t Level) {
+  if (Level >= BrStack.size())
+    return typeError(ErrorLoc,
+                     StringRef("br: invalid depth ") + std::to_string(Level));
+  const SmallVector<wasm::ValType, 4> &Expected =
+      BrStack[BrStack.size() - Level - 1];
+  if (Expected.size() > Stack.size())
+    return typeError(ErrorLoc, "br: insufficient values on the type stack");
+  auto IsStackTopInvalid = checkStackTop(Expected, Stack);
+  if (IsStackTopInvalid)
+    return typeError(ErrorLoc, "br " + IsStackTopInvalid.value());
+  return false;
+}
+
 bool WebAssemblyAsmTypeCheck::checkEnd(SMLoc ErrorLoc, bool PopVals) {
+  if (!PopVals)
+    BrStack.pop_back();
   if (LastSig.Returns.size() > Stack.size())
     return typeError(ErrorLoc, "end: insufficient values on the type stack");
 
   if (PopVals) {
     for (auto VT : llvm::reverse(LastSig.Returns)) {
       if (popType(ErrorLoc, VT))
         return true;
     }
     return false;
   }
 
-  for (size_t i = 0; i < LastSig.Returns.size(); i++) {
-    auto EVT = LastSig.Returns[i];
-    auto PVT = Stack[Stack.size() - LastSig.Returns.size() + i];
-    if (PVT != EVT)
-      return typeError(ErrorLoc,
-                       StringRef("end got ") + WebAssembly::typeToString(PVT) +
-                           ", expected " + WebAssembly::typeToString(EVT));
-  }
+  auto IsStackTopInvalid = checkStackTop(LastSig.Returns, Stack);
+  if (IsStackTopInvalid)
+    return typeError(ErrorLoc, "end " + IsStackTopInvalid.value());
   return false;
 }
 
 bool WebAssemblyAsmTypeCheck::checkSig(SMLoc ErrorLoc,
                                        const wasm::WasmSignature &Sig) {
   for (auto VT : llvm::reverse(Sig.Params))
     if (popType(ErrorLoc, VT))
       return true;
if (popType(ErrorLoc, wasm::ValType::I32))
       return true;
     if (popType(ErrorLoc, wasm::ValType::I32))
       return true;
     if (popType(ErrorLoc, Type))
       return true;
   } else if (Name == "drop") {
     if (popType(ErrorLoc, {}))
       return true;
+  } else if (Name == "try" || Name == "block" || Name == "loop" ||
+             Name == "if") {
+    if (Name == "if" && popType(ErrorLoc, wasm::ValType::I32))
+      return true;
+    if (Name == "loop")
+      BrStack.emplace_back(LastSig.Params.begin(), LastSig.Params.end());
+    else
+      BrStack.emplace_back(LastSig.Returns.begin(), LastSig.Returns.end());
   } else if (Name == "end_block" || Name == "end_loop" || Name == "end_if" ||
              Name == "else" || Name == "end_try" || Name == "catch" ||
              Name == "catch_all" || Name == "delegate") {
     if (checkEnd(ErrorLoc,
                  Name == "else" || Name == "catch" || Name == "catch_all"))
       return true;
     Unreachable = false;
     if (Name == "catch") {
       const MCSymbolRefExpr *SymRef;
       if (getSymRef(Operands[1]->getStartLoc(), 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(Operands[1]->getStartLoc(), 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 == "br") {
+    const MCOperand &Operand = Inst.getOperand(0);
+    if (!Operand.isImm())
+      return false;
+    if (checkBr(ErrorLoc, static_cast<size_t>(Operand.getImm())))
+      return true;
   } 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))

llvm/test/MC/WebAssembly/type-checker-br.s

+# RUN: llvm-mc -triple=wasm32 %s 2>&1
+
+br_block:
+  .functype br_block () -> ()
+  block f32
+    block i32
+      f32.const 1.0
+      br 1
+      drop
+      i32.const 1
+    end_block
+    drop
+    f32.const 1.0
+  end_block
+  drop
+  end_function
+
+br_func:
+  .functype br_block () -> ()
+  block i32
+    br 1
+    i32.const 1
+  end_block
+  drop
+  end_function

llvm/test/MC/WebAssembly/type-checker-errors.s

 check_after_unreachable_within_try_4:
   .functype check_after_unreachable_within_try_4 () -> ()
   try
   unreachable
   delegate 0
 # CHECK: :[[@LINE+1]]:3: error: empty stack while popping value
   drop
   end_function
+
+br_invalid_type_loop:
+  .functype br_invalid_type_loop () -> ()
+  i32.const 1
+  loop (i32) -> (f32)
+    drop
+    f32.const 1.0
+# CHECK: :[[@LINE+1]]:5: error: br got f32, expected i32
+    br 0
+  end_loop
+  drop
+  end_function
+
+br_invalid_type_block:
+  .functype br_invalid_type_block () -> ()
+  i32.const 1
+  block (i32) -> (f32)
+# CHECK: :[[@LINE+1]]:5: error: br got i32, expected f32
+    br 0
+    f32.const 1.0
+  end_block
+  drop
+  end_function
+
+br_invalid_type_if:
+  .functype br_invalid_type_if () -> ()
+  i32.const 1
+  if f32
+    f32.const 1.0
+  else
+    i32.const 1
+# CHECK: :[[@LINE+1]]:5: error: br got i32, expected f32
+    br 0
+  end_if
+  drop
+  end_function
+
+br_invalid_type_try:
+  .functype br_invalid_type_try () -> ()
+  try f32
+    i32.const 1
+# CHECK: :[[@LINE+1]]:5: error: br got i32, expected f32
+    br 0
+  catch tag_f32
+  end_try
+  drop
+  end_function
+
+br_invalid_type_catch:
+  .functype br_invalid_type_catch () -> ()
+  try f32
+    f32.const 1.0
+  catch tag_i32
+# CHECK: :[[@LINE+1]]:5: error: br got i32, expected f32
+    br 0
+  end_try
+  drop
+  end_function
+
+br_invalid_type_catch_all:
+  .functype br_invalid_type_catch_all () -> ()
+  try f32
+    f32.const 1.0
+  catch_all
+    i32.const 1
+# CHECK: :[[@LINE+1]]:5: error: br got i32, expected f32
+    br 0
+  end_try
+  drop
+  end_function
+
+br_invalid_depth_out_of_range:
+  .functype br_invalid_depth_out_of_range () -> ()
+  block
+  block
+  block
+# CHECK: :[[@LINE+1]]:5: error: br: invalid depth 4
+    br 4
+  end_block
+  end_block
+  end_block
+  end_function
+
+br_incorrect_signature:
+  .functype br_incorrect_signature () -> ()
+  block f32
+    block i32
+      i32.const 1
+# CHECK: :[[@LINE+1]]:7: error: br got i32, expected f32
+      br 1
+    end_block
+    drop
+    f32.const 1.0
+  end_block
+  drop
+  end_function
+
+br_incorrect_func_signature:
+  .functype br_incorrect_func_signature () -> (i32)
+  block f32
+    f32.const 1.0
+# CHECK: :[[@LINE+1]]:5: error: br got f32, expected i32
+    br 1
+  end_block
+  drop
+  i32.const 1
+  end_function