[Parser] Support loops (#5966)

Parse loops in the new wat parser and add support for them to the IRBuilder.

4 files changed, 109 insertions, 33 deletions

diff --git a/src/parser/contexts.h b/src/parser/contexts.h
index 2ecaac70e..f2aeb2489 100644
--- a/src/parser/contexts.h
+++ b/src/parser/contexts.h
@@ -309,8 +309,12 @@ struct NullInstrParserCtx {
   Result<> makeIf(Index, std::optional<Name>, BlockTypeT) {
     return Ok{};
   }
-  Result<> visitEnd(Index) { return Ok{}; }
-  Result<> visitElse(Index) { return Ok{}; }
+  Result<> visitElse() { return Ok{}; }
+  template<typename BlockTypeT>
+  Result<> makeLoop(Index, std::optional<Name>, BlockTypeT) {
+    return Ok{};
+  }
+  Result<> visitEnd() { return Ok{}; }
 
   Result<> makeUnreachable(Index) { return Ok{}; }
   Result<> makeNop(Index) { return Ok{}; }
@@ -989,9 +993,17 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
       irBuilder.makeIf(label ? *label : Name{}, type.getSignature().results));
   }
 
-  Result<> visitEnd(Index pos) { return withLoc(pos, irBuilder.visitEnd()); }
+  Result<> visitElse() { return withLoc(irBuilder.visitElse()); }
+
+  Result<> makeLoop(Index pos, std::optional<Name> label, HeapType type) {
+    // TODO: validate labels?
+    // TODO: Move error on input types to here?
+    return withLoc(
+      pos,
+      irBuilder.makeLoop(label ? *label : Name{}, type.getSignature().results));
+  }
 
-  Result<> visitElse(Index pos) { return withLoc(pos, irBuilder.visitElse()); }
+  Result<> visitEnd() { return withLoc(irBuilder.visitEnd()); }
 
   Result<> makeUnreachable(Index pos) {
     return withLoc(pos, irBuilder.makeUnreachable());
diff --git a/src/parser/parsers.h b/src/parser/parsers.h
index 2dd0d9eb4..38005253f 100644
--- a/src/parser/parsers.h
+++ b/src/parser/parsers.h
@@ -55,6 +55,7 @@ template<typename Ctx> Result<typename Ctx::MemargT> memarg(Ctx&, uint32_t);
 template<typename Ctx> Result<typename Ctx::BlockTypeT> blocktype(Ctx&);
 template<typename Ctx> MaybeResult<> block(Ctx&, bool);
 template<typename Ctx> MaybeResult<> ifelse(Ctx&, bool);
+template<typename Ctx> MaybeResult<> loop(Ctx&, bool);
 template<typename Ctx> Result<> makeUnreachable(Ctx&, Index);
 template<typename Ctx> Result<> makeNop(Ctx&, Index);
 template<typename Ctx> Result<> makeBinary(Ctx&, Index, BinaryOp op);
@@ -553,6 +554,9 @@ template<typename Ctx> MaybeResult<> foldedBlockinstr(Ctx& ctx) {
   if (auto i = ifelse(ctx, true)) {
     return i;
   }
+  if (auto i = loop(ctx, true)) {
+    return i;
+  }
   // TODO: Other block instructions
   return {};
 }
@@ -564,6 +568,9 @@ template<typename Ctx> MaybeResult<> unfoldedBlockinstr(Ctx& ctx) {
   if (auto i = ifelse(ctx, false)) {
     return i;
   }
+  if (auto i = loop(ctx, false)) {
+    return i;
+  }
   // TODO: Other block instructions
   return {};
 }
@@ -741,14 +748,9 @@ template<typename Ctx> Result<typename Ctx::BlockTypeT> blocktype(Ctx& ctx) {
 template<typename Ctx> MaybeResult<> block(Ctx& ctx, bool folded) {
   auto pos = ctx.in.getPos();
 
-  if (folded) {
-    if (!ctx.in.takeSExprStart("block"sv)) {
-      return {};
-    }
-  } else {
-    if (!ctx.in.takeKeyword("block"sv)) {
-      return {};
-    }
+  if ((folded && !ctx.in.takeSExprStart("block"sv)) ||
+      (!folded && !ctx.in.takeKeyword("block"sv))) {
+    return {};
   }
 
   auto label = ctx.in.takeID();
@@ -774,7 +776,7 @@ template<typename Ctx> MaybeResult<> block(Ctx& ctx, bool folded) {
     }
   }
 
-  return ctx.visitEnd(pos);
+  return ctx.visitEnd();
 }
 
 // if ::= 'if' label blocktype instr1* ('else' id1? instr2*)? 'end' id2?
@@ -783,14 +785,9 @@ template<typename Ctx> MaybeResult<> block(Ctx& ctx, bool folded) {
 template<typename Ctx> MaybeResult<> ifelse(Ctx& ctx, bool folded) {
   auto pos = ctx.in.getPos();
 
-  if (folded) {
-    if (!ctx.in.takeSExprStart("if"sv)) {
-      return {};
-    }
-  } else {
-    if (!ctx.in.takeKeyword("if"sv)) {
-      return {};
-    }
+  if ((folded && !ctx.in.takeSExprStart("if"sv)) ||
+      (!folded && !ctx.in.takeKeyword("if"sv))) {
+    return {};
   }
 
   auto label = ctx.in.takeID();
@@ -821,7 +818,7 @@ template<typename Ctx> MaybeResult<> ifelse(Ctx& ctx, bool folded) {
       return ctx.in.err("else label does not match if label");
     }
 
-    ctx.visitElse(pos);
+    ctx.visitElse();
 
     CHECK_ERR(instrs(ctx));
 
@@ -838,14 +835,48 @@ template<typename Ctx> MaybeResult<> ifelse(Ctx& ctx, bool folded) {
     if (!ctx.in.takeKeyword("end"sv)) {
       return ctx.in.err("expected 'end' at end of if");
     }
+    auto id2 = ctx.in.takeID();
+    if (id2 && id2 != label) {
+      return ctx.in.err("end label does not match if label");
+    }
   }
 
-  auto id2 = ctx.in.takeID();
-  if (id2 && id2 != label) {
-    return ctx.in.err("end label does not match if label");
+  return ctx.visitEnd();
+}
+
+// loop ::= 'loop' label blocktype instr* end id?
+//        | '(' 'loop' label blocktype instr* ')'
+template<typename Ctx> MaybeResult<> loop(Ctx& ctx, bool folded) {
+  auto pos = ctx.in.getPos();
+
+  if ((folded && !ctx.in.takeSExprStart("loop"sv)) ||
+      (!folded && !ctx.in.takeKeyword("loop"sv))) {
+    return {};
   }
 
-  return ctx.visitEnd(pos);
+  auto label = ctx.in.takeID();
+
+  auto type = blocktype(ctx);
+  CHECK_ERR(type);
+
+  ctx.makeLoop(pos, label, *type);
+
+  CHECK_ERR(instrs(ctx));
+
+  if (folded) {
+    if (!ctx.in.takeRParen()) {
+      return ctx.in.err("expected ')' at end of loop");
+    }
+  } else {
+    if (!ctx.in.takeKeyword("end"sv)) {
+      return ctx.in.err("expected 'end' at end of loop");
+    }
+    auto id = ctx.in.takeID();
+    if (id && id != label) {
+      return ctx.in.err("end label does not match loop label");
+    }
+  }
+  return ctx.visitEnd();
 }
 
 template<typename Ctx> Result<> makeUnreachable(Ctx& ctx, Index pos) {
diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h
index c7185ebf0..a6afc63c4 100644
--- a/src/wasm-ir-builder.h
+++ b/src/wasm-ir-builder.h
@@ -54,6 +54,7 @@ public:
   [[nodiscard]] Result<> visitBlockStart(Block* block);
   [[nodiscard]] Result<> visitIfStart(If* iff, Name label = {});
   [[nodiscard]] Result<> visitElse();
+  [[nodiscard]] Result<> visitLoopStart(Loop* iff);
   [[nodiscard]] Result<> visitEnd();
 
   // Alternatively, call makeXYZ to have the IRBuilder allocate the nodes. This
@@ -62,7 +63,7 @@ public:
   [[nodiscard]] Result<> makeNop();
   [[nodiscard]] Result<> makeBlock(Name label, Type type);
   [[nodiscard]] Result<> makeIf(Name label, Type type);
-  // [[nodiscard]] Result<> makeLoop();
+  [[nodiscard]] Result<> makeLoop(Name label, Type type);
   // [[nodiscard]] Result<> makeBreak();
   // [[nodiscard]] Result<> makeSwitch();
   // [[nodiscard]] Result<> makeCall();
@@ -199,7 +200,11 @@ private:
       If* iff;
       Name label;
     };
-    using Scope = std::variant<NoScope, BlockScope, IfScope, ElseScope>;
+    struct LoopScope {
+      Loop* loop;
+    };
+    using Scope =
+      std::variant<NoScope, BlockScope, IfScope, ElseScope, LoopScope>;
 
     // The control flow structure we are building expressions for.
     Scope scope;
@@ -221,6 +226,7 @@ private:
     static ScopeCtx makeElse(If* iff, Name label = {}) {
       return ScopeCtx(ElseScope{iff, label});
     }
+    static ScopeCtx makeLoop(Loop* loop) { return ScopeCtx(LoopScope{loop}); }
 
     bool isNone() { return std::get_if<NoScope>(&scope); }
     Block* getBlock() {
@@ -241,6 +247,12 @@ private:
       }
       return nullptr;
     }
+    Loop* getLoop() {
+      if (auto* loopScope = std::get_if<LoopScope>(&scope)) {
+        return loopScope->loop;
+      }
+      return nullptr;
+    }
     Type getResultType() {
       if (auto* block = getBlock()) {
         return block->type;
@@ -251,6 +263,9 @@ private:
       if (auto* iff = getElse()) {
         return iff->type;
       }
+      if (auto* loop = getLoop()) {
+        return loop->type;
+      }
       WASM_UNREACHABLE("unexpected scope kind");
     }
     Name getLabel() {
@@ -263,6 +278,9 @@ private:
       if (auto* elseScope = std::get_if<ElseScope>(&scope)) {
         return elseScope->label;
       }
+      if (auto* loop = getLoop()) {
+        return loop->name;
+      }
       WASM_UNREACHABLE("unexpected scope kind");
     }
   };
diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index b8d742ae8..585f9be49 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -294,6 +294,11 @@ Result<> IRBuilder::visitIfStart(If* iff, Name label) {
   return Ok{};
 }
 
+Result<> IRBuilder::visitLoopStart(Loop* loop) {
+  scopeStack.push_back(ScopeCtx::makeLoop(loop));
+  return Ok{};
+}
+
 Result<Expression*> IRBuilder::finishScope(Block* block) {
   if (scopeStack.empty() || scopeStack.back().isNone()) {
     return Err{"unexpected end of scope"};
@@ -403,6 +408,13 @@ Result<> IRBuilder::visitEnd() {
     block->finalize(block->type);
     push(block);
     return Ok{};
+  } else if (auto* loop = scope.getLoop()) {
+    auto expr = finishScope();
+    CHECK_ERR(expr);
+    loop->body = *expr;
+    loop->finalize(loop->type);
+    push(loop);
+    return Ok{};
   }
   auto label = scope.getLabel();
   Expression* scopeExpr = nullptr;
@@ -411,13 +423,13 @@ Result<> IRBuilder::visitEnd() {
     CHECK_ERR(expr);
     iff->ifTrue = *expr;
     iff->ifFalse = nullptr;
-    iff->finalize();
+    iff->finalize(iff->type);
     scopeExpr = iff;
   } else if (auto* iff = scope.getElse()) {
     auto expr = finishScope();
     CHECK_ERR(expr);
     iff->ifFalse = *expr;
-    iff->finalize();
+    iff->finalize(iff->type);
     scopeExpr = iff;
   }
   assert(scopeExpr && "unexpected scope kind");
@@ -449,9 +461,12 @@ Result<> IRBuilder::makeIf(Name label, Type type) {
   return visitIfStart(iff, label);
 }
 
-// Result<> IRBuilder::makeIf() {}
-
-// Result<> IRBuilder::makeLoop() {}
+Result<> IRBuilder::makeLoop(Name label, Type type) {
+  auto* loop = wasm.allocator.alloc<Loop>();
+  loop->name = label;
+  loop->type = type;
+  return visitLoopStart(loop);
+}
 
 // Result<> IRBuilder::makeBreak() {}