[Parser] Support loops (#5966)

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

2 files changed, 69 insertions, 26 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) {