[Parser] Parse labels and br (#5970)

The parser previously parsed labels and could attach them to control flow
structures, but did not maintain the context necessary to correctly parse
branches. Support parsing labels as both names and indices in IRBuilder,
handling shadowing correctly, and use that support to implement parsing of br.

2 files changed, 87 insertions, 14 deletions

diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index 4312fefe3..4d1c3353d 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -181,6 +181,7 @@ Result<Expression*> IRBuilder::build() {
   assert(scopeStack.back().exprStack.size() == 1);
   auto* expr = scopeStack.back().exprStack.back();
   scopeStack.clear();
+  labelDepths.clear();
   return expr;
 }
 
@@ -206,6 +207,10 @@ Result<> IRBuilder::visitExpression(Expression* curr) {
   auto field = pop();                                                          \
   CHECK_ERR(field);                                                            \
   expr->field = *field;
+#define DELEGATE_FIELD_SCOPE_NAME_DEF(id, field)                               \
+  if (labelDepths.count(expr->field)) {                                        \
+    return Err{"repeated label"};                                              \
+  }
 #define DELEGATE_END(id)
 
 #define DELEGATE_FIELD_OPTIONAL_CHILD(id, field)                               \
@@ -214,15 +219,18 @@ Result<> IRBuilder::visitExpression(Expression* curr) {
 #define DELEGATE_FIELD_CHILD_VECTOR(id, field)                                 \
   WASM_UNREACHABLE("should have called visit" #id " because " #id              \
                    " has child vector " #field);
+#define DELEGATE_FIELD_SCOPE_NAME_USE(id, field)                               \
+  WASM_UNREACHABLE("should have called visit" #id " because " #id              \
+                   " has scope name use " #field);
+#define DELEGATE_FIELD_SCOPE_NAME_USE_VECTOR(id, field)                        \
+  WASM_UNREACHABLE("should have called visit" #id " because " #id              \
+                   " has scope name use vector " #field);
 
 #define DELEGATE_FIELD_INT(id, field)
 #define DELEGATE_FIELD_INT_ARRAY(id, field)
 #define DELEGATE_FIELD_LITERAL(id, field)
 #define DELEGATE_FIELD_NAME(id, field)
 #define DELEGATE_FIELD_NAME_VECTOR(id, field)
-#define DELEGATE_FIELD_SCOPE_NAME_DEF(id, field)
-#define DELEGATE_FIELD_SCOPE_NAME_USE(id, field)
-#define DELEGATE_FIELD_SCOPE_NAME_USE_VECTOR(id, field)
 #define DELEGATE_FIELD_TYPE(id, field)
 #define DELEGATE_FIELD_HEAPTYPE(id, field)
 #define DELEGATE_FIELD_ADDRESS(id, field)
@@ -281,6 +289,30 @@ Result<> IRBuilder::visitArrayNew(ArrayNew* curr) {
   return Ok{};
 }
 
+Result<> IRBuilder::visitBreak(Break* curr, std::optional<Index> label) {
+  if (!label) {
+    auto index = getLabelIndex(curr->name);
+    CHECK_ERR(index);
+    label = *index;
+  }
+  auto scope = getScope(*label);
+  CHECK_ERR(scope);
+  std::vector<Expression*> values((*scope)->getResultType().size());
+  for (size_t i = 0, size = values.size(); i < size; ++i) {
+    auto val = pop();
+    CHECK_ERR(val);
+    values[size - 1 - i] = *val;
+  }
+  if (values.size() == 0) {
+    curr->value = nullptr;
+  } else if (values.size() == 1) {
+    curr->value = values[0];
+  } else {
+    curr->value = builder.makeTupleMake(values);
+  }
+  return Ok{};
+}
+
 Result<> IRBuilder::visitFunctionStart(Function* func) {
   if (!scopeStack.empty()) {
     return Err{"unexpected start of function"};
@@ -291,7 +323,7 @@ Result<> IRBuilder::visitFunctionStart(Function* func) {
 }
 
 Result<> IRBuilder::visitBlockStart(Block* curr) {
-  scopeStack.push_back(ScopeCtx::makeBlock(curr));
+  pushScope(ScopeCtx::makeBlock(curr));
   return Ok{};
 }
 
@@ -299,12 +331,12 @@ Result<> IRBuilder::visitIfStart(If* iff, Name label) {
   auto cond = pop();
   CHECK_ERR(cond);
   iff->condition = *cond;
-  scopeStack.push_back(ScopeCtx::makeIf(iff, label));
+  pushScope(ScopeCtx::makeIf(iff, label));
   return Ok{};
 }
 
 Result<> IRBuilder::visitLoopStart(Loop* loop) {
-  scopeStack.push_back(ScopeCtx::makeLoop(loop));
+  pushScope(ScopeCtx::makeLoop(loop));
   return Ok{};
 }
 
@@ -356,6 +388,7 @@ Result<Expression*> IRBuilder::finishScope(Block* block) {
     auto hoisted = hoistLastValue();
     CHECK_ERR(hoisted);
   }
+
   Expression* ret = nullptr;
   if (scope.exprStack.size() == 0) {
     // No expressions for this scope, but we need something. If we were given a
@@ -385,6 +418,12 @@ Result<Expression*> IRBuilder::finishScope(Block* block) {
     }
     ret = block;
   }
+
+  // If this scope had a label, remove it from the context.
+  if (auto label = scope.getOriginalLabel()) {
+    labelDepths.at(label).pop_back();
+  }
+
   scopeStack.pop_back();
   return ret;
 }
@@ -395,11 +434,12 @@ Result<> IRBuilder::visitElse() {
   if (!iff) {
     return Err{"unexpected else"};
   }
-  auto label = scope.getLabel();
+  auto originalLabel = scope.getOriginalLabel();
+  auto label = scope.label;
   auto expr = finishScope();
   CHECK_ERR(expr);
   iff->ifTrue = *expr;
-  scopeStack.push_back(ScopeCtx::makeElse(iff, label));
+  pushScope(ScopeCtx::makeElse(iff, originalLabel, label));
   return Ok{};
 }
 
@@ -414,22 +454,25 @@ Result<> IRBuilder::visitEnd() {
   // If the scope expression cannot be directly labeled, we may need to wrap it
   // in a block.
   auto maybeWrapForLabel = [&](Expression* curr) -> Expression* {
-    if (auto label = scope.getLabel()) {
-      return builder.makeBlock(label, {curr}, curr->type);
+    if (scope.label) {
+      return builder.makeBlock(scope.label, {curr}, scope.getResultType());
     }
     return curr;
   };
 
   if (auto* func = scope.getFunction()) {
-    func->body = *expr;
+    func->body = maybeWrapForLabel(*expr);
+    labelDepths.clear();
   } else if (auto* block = scope.getBlock()) {
     assert(*expr == block);
+    block->name = scope.label;
     // TODO: Track branches so we can know whether this block is a target and
     // finalize more efficiently.
     block->finalize(block->type);
     push(block);
   } else if (auto* loop = scope.getLoop()) {
     loop->body = *expr;
+    loop->name = scope.label;
     loop->finalize(loop->type);
     push(loop);
   } else if (auto* iff = scope.getIf()) {
@@ -447,6 +490,25 @@ Result<> IRBuilder::visitEnd() {
   return Ok{};
 }
 
+Result<Index> IRBuilder::getLabelIndex(Name label) {
+  auto it = labelDepths.find(label);
+  if (it == labelDepths.end() || it->second.empty()) {
+    return Err{"unexpected label '"s + label.toString()};
+  }
+  return scopeStack.size() - it->second.back();
+}
+
+Result<Name> IRBuilder::getLabelName(Index label) {
+  auto scope = getScope(label);
+  CHECK_ERR(scope);
+  auto& scopeLabel = (*scope)->label;
+  if (!scopeLabel) {
+    // The scope does not already have a name, so we need to create one.
+    scopeLabel = makeFresh("label");
+  }
+  return scopeLabel;
+}
+
 Result<> IRBuilder::makeNop() {
   push(builder.makeNop());
   return Ok{};
@@ -472,7 +534,15 @@ Result<> IRBuilder::makeLoop(Name label, Type type) {
   return visitLoopStart(loop);
 }
 
-// Result<> IRBuilder::makeBreak() {}
+Result<> IRBuilder::makeBreak(Index label) {
+  auto name = getLabelName(label);
+  CHECK_ERR(name);
+  Break curr;
+  curr.name = *name;
+  CHECK_ERR(visitBreak(&curr, label));
+  push(builder.makeBreak(curr.name, curr.value));
+  return Ok{};
+}
 
 // Result<> IRBuilder::makeSwitch() {}
 
diff --git a/src/wasm/wasm-validator.cpp b/src/wasm/wasm-validator.cpp
index a8bb4a536..337065b91 100644
--- a/src/wasm/wasm-validator.cpp
+++ b/src/wasm/wasm-validator.cpp
@@ -651,8 +651,11 @@ void FunctionValidator::visitBlock(Block* curr) {
     auto iter = breakTypes.find(curr->name);
     assert(iter != breakTypes.end()); // we set it ourselves
     for (Type breakType : iter->second) {
-      // none or unreachable means a poison value that we should ignore - if
-      // consumed, it will error
+      if (breakType == Type::none && curr->type == Type::unreachable) {
+        // We allow empty breaks to unreachable blocks.
+        continue;
+      }
+
       shouldBeSubType(breakType,
                       curr->type,
                       curr,