[Parser] Support prologue and epilogue sourcemap annotations (#6370)

and fix a bug with sourcemap annotations on folded `if` conditions. Update
IRBuilder to apply prologue and epilogue source locations when beginning and ending
a function scope. Add basic support in the parser for explicitly tracking
annotations on module fields, although only do anything with them in the case of
prologue source location annotations.

4 files changed, 58 insertions, 32 deletions

diff --git a/src/parser/context-decls.cpp b/src/parser/context-decls.cpp
index d9279a0b6..c78c47d60 100644
--- a/src/parser/context-decls.cpp
+++ b/src/parser/context-decls.cpp
@@ -68,12 +68,14 @@ Result<> ParseDeclsCtx::addFunc(Name name,
                                 ImportNames* import,
                                 TypeUseT type,
                                 std::optional<LocalsT>,
+                                std::vector<Annotation>&& annotations,
                                 Index pos) {
   CHECK_ERR(checkImport(pos, import));
   auto f = addFuncDecl(pos, name, import);
   CHECK_ERR(f);
   CHECK_ERR(addExports(in, wasm, *f, exports, ExternalKind::Function));
-  funcDefs.push_back({name, pos, Index(funcDefs.size())});
+  funcDefs.push_back(
+    {name, pos, Index(funcDefs.size()), std::move(annotations)});
   return Ok{};
 }
 
@@ -109,7 +111,8 @@ Result<> ParseDeclsCtx::addTable(Name name,
   auto t = addTableDecl(pos, name, import, limits);
   CHECK_ERR(t);
   CHECK_ERR(addExports(in, wasm, *t, exports, ExternalKind::Table));
-  tableDefs.push_back({name, pos, Index(tableDefs.size())});
+  // TODO: table annotations
+  tableDefs.push_back({name, pos, Index(tableDefs.size()), {}});
   return Ok{};
 }
 
@@ -164,7 +167,8 @@ Result<> ParseDeclsCtx::addMemory(Name name,
   auto m = addMemoryDecl(pos, name, import, type);
   CHECK_ERR(m);
   CHECK_ERR(addExports(in, wasm, *m, exports, ExternalKind::Memory));
-  memoryDefs.push_back({name, pos, Index(memoryDefs.size())});
+  // TODO: memory annotations
+  memoryDefs.push_back({name, pos, Index(memoryDefs.size()), {}});
   return Ok{};
 }
 
@@ -209,7 +213,8 @@ Result<> ParseDeclsCtx::addGlobal(Name name,
   auto g = addGlobalDecl(pos, name, import);
   CHECK_ERR(g);
   CHECK_ERR(addExports(in, wasm, *g, exports, ExternalKind::Global));
-  globalDefs.push_back({name, pos, Index(globalDefs.size())});
+  // TODO: global annotations
+  globalDefs.push_back({name, pos, Index(globalDefs.size()), {}});
   return Ok{};
 }
 
@@ -228,7 +233,8 @@ Result<> ParseDeclsCtx::addElem(
     name = Names::getValidElementSegmentName(wasm, name);
     e->name = name;
   }
-  elemDefs.push_back({name, pos, Index(wasm.elementSegments.size())});
+  // TODO: element segment annotations
+  elemDefs.push_back({name, pos, Index(wasm.elementSegments.size()), {}});
   wasm.addElementSegment(std::move(e));
   return Ok{};
 }
@@ -252,7 +258,8 @@ Result<> ParseDeclsCtx::addData(Name name,
     d->name = name;
   }
   d->data = std::move(data);
-  dataDefs.push_back({name, pos, Index(wasm.dataSegments.size())});
+  // TODO: data segment annotations
+  dataDefs.push_back({name, pos, Index(wasm.dataSegments.size()), {}});
   wasm.addDataSegment(std::move(d));
   return Ok{};
 }
@@ -285,7 +292,8 @@ Result<> ParseDeclsCtx::addTag(Name name,
   auto t = addTagDecl(pos, name, import);
   CHECK_ERR(t);
   CHECK_ERR(addExports(in, wasm, *t, exports, ExternalKind::Tag));
-  tagDefs.push_back({name, pos, Index(tagDefs.size())});
+  // TODO: tag annotations
+  tagDefs.push_back({name, pos, Index(tagDefs.size()), {}});
   return Ok{};
 }
 
diff --git a/src/parser/contexts.h b/src/parser/contexts.h
index 08cebe33b..e79330871 100644
--- a/src/parser/contexts.h
+++ b/src/parser/contexts.h
@@ -61,6 +61,7 @@ struct DefPos {
   Name name;
   Index pos;
   Index index;
+  std::vector<Annotation> annotations;
 };
 
 struct GlobalType {
@@ -406,7 +407,7 @@ struct NullInstrParserCtx {
   TagLabelListT makeTagLabelList() { return Ok{}; }
   void appendTagLabel(TagLabelListT&, TagIdxT, LabelIdxT) {}
 
-  void setSrcLoc(const Annotation&) {}
+  void setSrcLoc(const std::vector<Annotation>&) {}
 
   Result<> makeUnreachable(Index, const std::vector<Annotation>&) {
     return Ok{};
@@ -908,12 +909,14 @@ struct ParseDeclsCtx : NullTypeParserCtx, NullInstrParserCtx {
   void setOpen() {}
   Result<> addSubtype(Index) { return Ok{}; }
   void finishSubtype(Name name, Index pos) {
-    subtypeDefs.push_back({name, pos, Index(subtypeDefs.size())});
+    // TODO: type annotations
+    subtypeDefs.push_back({name, pos, Index(subtypeDefs.size()), {}});
   }
   size_t getRecGroupStartIndex() { return 0; }
   void addRecGroup(Index, size_t) {}
   void finishDeftype(Index pos) {
-    typeDefs.push_back({{}, pos, Index(typeDefs.size())});
+    // TODO: type annotations
+    typeDefs.push_back({{}, pos, Index(typeDefs.size()), {}});
   }
 
   Limits makeLimits(uint64_t n, std::optional<uint64_t> m) {
@@ -957,6 +960,7 @@ struct ParseDeclsCtx : NullTypeParserCtx, NullInstrParserCtx {
                    ImportNames* import,
                    TypeUseT type,
                    std::optional<LocalsT>,
+                   std::vector<Annotation>&&,
                    Index pos);
 
   Result<Table*>
@@ -991,7 +995,8 @@ struct ParseDeclsCtx : NullTypeParserCtx, NullInstrParserCtx {
     if (!startDefs.empty()) {
       return Err{"unexpected extra 'start' function"};
     }
-    startDefs.push_back({{}, pos, 0});
+    // TODO: start function annotations.
+    startDefs.push_back({{}, pos, 0, {}});
     return Ok{};
   }
 
@@ -1245,6 +1250,7 @@ struct ParseModuleTypesCtx : TypeParserCtx<ParseModuleTypesCtx>,
                    ImportNames*,
                    TypeUse type,
                    std::optional<LocalsT> locals,
+                   std::vector<Annotation>&&,
                    Index pos) {
     auto& f = wasm.functions[index];
     if (!type.type.isSignature()) {
@@ -1606,6 +1612,7 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
                    ImportNames*,
                    TypeUseT,
                    std::optional<LocalsT>,
+                   std::vector<Annotation>&&,
                    Index) {
     return Ok{};
   }
@@ -1692,9 +1699,17 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
     return wasm.memories[0]->name;
   }
 
-  void setSrcLoc(const Annotation& annotation) {
-    assert(annotation.kind == srcAnnotationKind);
-    Lexer lexer(annotation.contents);
+  void setSrcLoc(const std::vector<Annotation>& annotations) {
+    const Annotation* annotation = nullptr;
+    for (auto& a : annotations) {
+      if (a.kind == srcAnnotationKind) {
+        annotation = &a;
+      }
+    }
+    if (!annotation) {
+      return;
+    }
+    Lexer lexer(annotation->contents);
     auto contents = lexer.takeKeyword();
     if (!contents || !lexer.empty()) {
       return;
diff --git a/src/parser/parsers.h b/src/parser/parsers.h
index 0dde4b765..0150af811 100644
--- a/src/parser/parsers.h
+++ b/src/parser/parsers.h
@@ -844,20 +844,11 @@ template<typename Ctx> Result<uint32_t> tupleArity(Ctx& ctx) {
 // Instructions
 // ============
 
-template<typename Ctx>
-void setSrcLoc(Ctx& ctx, const std::vector<Annotation>& annotations) {
-  for (const auto& annotation : annotations) {
-    if (annotation.kind == srcAnnotationKind) {
-      ctx.setSrcLoc(annotation);
-    }
-  }
-}
-
 // blockinstr ::= block | loop | if-else | try-catch | try_table
 template<typename Ctx>
 MaybeResult<> foldedBlockinstr(Ctx& ctx,
                                const std::vector<Annotation>& annotations) {
-  setSrcLoc(ctx, annotations);
+  ctx.setSrcLoc(annotations);
   if (auto i = block(ctx, annotations, true)) {
     return i;
   }
@@ -879,7 +870,7 @@ MaybeResult<> foldedBlockinstr(Ctx& ctx,
 template<typename Ctx>
 MaybeResult<> unfoldedBlockinstr(Ctx& ctx,
                                  const std::vector<Annotation>& annotations) {
-  setSrcLoc(ctx, annotations);
+  ctx.setSrcLoc(annotations);
   if (auto i = block(ctx, annotations, false)) {
     return i;
   }
@@ -912,7 +903,7 @@ MaybeResult<> blockinstr(Ctx& ctx, const std::vector<Annotation>& annotations) {
 // plaininstr ::= ... all plain instructions ...
 template<typename Ctx>
 MaybeResult<> plaininstr(Ctx& ctx, const std::vector<Annotation>& annotations) {
-  setSrcLoc(ctx, annotations);
+  ctx.setSrcLoc(annotations);
   auto pos = ctx.in.getPos();
   auto keyword = ctx.in.takeKeyword();
   if (!keyword) {
@@ -1061,6 +1052,7 @@ Result<typename Ctx::MemargT> memarg(Ctx& ctx, uint32_t n) {
 // blocktype ::= (t:result)? => t? | x,I:typeuse => x if I = {}
 template<typename Ctx> Result<typename Ctx::BlockTypeT> blocktype(Ctx& ctx) {
   auto pos = ctx.in.getPos();
+  auto initialLexer = ctx.in;
 
   if (auto res = results(ctx)) {
     CHECK_ERR(res);
@@ -1071,7 +1063,7 @@ template<typename Ctx> Result<typename Ctx::BlockTypeT> blocktype(Ctx& ctx) {
 
   // We either had no results or multiple results. Reset and parse again as a
   // type use.
-  ctx.in.setIndex(pos);
+  ctx.in = initialLexer;
   auto use = typeuse(ctx);
   CHECK_ERR(use);
 
@@ -1138,6 +1130,7 @@ ifelse(Ctx& ctx, const std::vector<Annotation>& annotations, bool folded) {
 
   if (folded) {
     CHECK_ERR(foldedinstrs(ctx));
+    ctx.setSrcLoc(annotations);
   }
 
   ctx.makeIf(pos, annotations, label, *type);
@@ -2928,8 +2921,9 @@ template<typename Ctx> MaybeResult<> import_(Ctx& ctx) {
     auto name = ctx.in.takeID();
     auto type = typeuse(ctx);
     CHECK_ERR(type);
-    CHECK_ERR(
-      ctx.addFunc(name ? *name : Name{}, {}, &names, *type, std::nullopt, pos));
+    // TODO: function import annotations
+    CHECK_ERR(ctx.addFunc(
+      name ? *name : Name{}, {}, &names, *type, std::nullopt, {}, pos));
   } else if (ctx.in.takeSExprStart("table"sv)) {
     auto name = ctx.in.takeID();
     auto type = tabletype(ctx);
@@ -2971,6 +2965,8 @@ template<typename Ctx> MaybeResult<> import_(Ctx& ctx) {
 //              '(' 'import' mod:name nm:name ')' typeuse ')'
 template<typename Ctx> MaybeResult<> func(Ctx& ctx) {
   auto pos = ctx.in.getPos();
+  auto annotations = ctx.in.getAnnotations();
+
   if (!ctx.in.takeSExprStart("func"sv)) {
     return {};
   }
@@ -2996,14 +2992,20 @@ template<typename Ctx> MaybeResult<> func(Ctx& ctx) {
       localVars = *l;
     }
     CHECK_ERR(instrs(ctx));
+    ctx.setSrcLoc(ctx.in.takeAnnotations());
   }
 
   if (!ctx.in.takeRParen()) {
     return ctx.in.err("expected end of function");
   }
 
-  CHECK_ERR(
-    ctx.addFunc(name, *exports, import.getPtr(), *type, localVars, pos));
+  CHECK_ERR(ctx.addFunc(name,
+                        *exports,
+                        import.getPtr(),
+                        *type,
+                        localVars,
+                        std::move(annotations),
+                        pos));
   return Ok{};
 }
 
diff --git a/src/parser/wat-parser.cpp b/src/parser/wat-parser.cpp
index ff0dea768..8675f4e36 100644
--- a/src/parser/wat-parser.cpp
+++ b/src/parser/wat-parser.cpp
@@ -184,10 +184,11 @@ Result<> parseModule(Module& wasm, std::string_view input) {
     for (Index i = 0; i < decls.funcDefs.size(); ++i) {
       ctx.index = i;
       auto* f = wasm.functions[i].get();
+      WithPosition with(ctx, decls.funcDefs[i].pos);
+      ctx.setSrcLoc(decls.funcDefs[i].annotations);
       if (!f->imported()) {
         CHECK_ERR(ctx.visitFunctionStart(f));
       }
-      WithPosition with(ctx, decls.funcDefs[i].pos);
       if (auto parsed = func(ctx)) {
         CHECK_ERR(parsed);
       } else {