[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.

7 files changed, 91 insertions, 33 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a85bcacaf..fc03cf492 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,10 @@ full changeset diff at the end of each section.
 Current Trunk
 -------------
 
+ - (If new wat parser is enabled) Source map comments on `else` branches must
+   now be placed above the instruction inside the `else` branch rather than on
+   the `else` branch itself.
+
 v117
 ----
 
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 {
diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp
index 628a257c7..ffd0b8674 100644
--- a/src/wasm/wasm-ir-builder.cpp
+++ b/src/wasm/wasm-ir-builder.cpp
@@ -211,12 +211,17 @@ Result<Expression*> IRBuilder::build() {
 }
 
 void IRBuilder::setDebugLocation(const Function::DebugLocation& loc) {
+  DBG(std::cerr << "setting debugloc " << loc.fileIndex << ":" << loc.lineNumber
+                << ":" << loc.columnNumber << "\n";);
   debugLoc = loc;
 }
 
 void IRBuilder::applyDebugLoc(Expression* expr) {
   if (debugLoc) {
     if (func) {
+      DBG(std::cerr << "applying debugloc " << debugLoc->fileIndex << ":"
+                    << debugLoc->lineNumber << ":" << debugLoc->columnNumber
+                    << " to expression " << ShallowExpression{expr} << "\n");
       func->debugLocations[expr] = *debugLoc;
     }
     debugLoc.reset();
@@ -698,6 +703,10 @@ Result<> IRBuilder::visitFunctionStart(Function* func) {
   if (!scopeStack.empty()) {
     return Err{"unexpected start of function"};
   }
+  if (debugLoc) {
+    func->prologLocation.insert(*debugLoc);
+    debugLoc.reset();
+  }
   scopeStack.push_back(ScopeCtx::makeFunc(func));
   this->func = func;
   return Ok{};
@@ -740,6 +749,11 @@ Result<> IRBuilder::visitTryTableStart(TryTable* trytable, Name label) {
 }
 
 Result<Expression*> IRBuilder::finishScope(Block* block) {
+  if (debugLoc) {
+    DBG(std::cerr << "discarding debugloc " << debugLoc->fileIndex << ":"
+                  << debugLoc->lineNumber << ":" << debugLoc->columnNumber
+                  << "\n");
+  }
   debugLoc.reset();
 
   if (scopeStack.empty() || scopeStack.back().isNone()) {
@@ -966,6 +980,10 @@ Result<> IRBuilder::visitEnd() {
   if (scope.isNone()) {
     return Err{"unexpected end"};
   }
+  if (auto* func = scope.getFunction(); func && debugLoc) {
+    func->epilogLocation.insert(*debugLoc);
+    debugLoc.reset();
+  }
   auto expr = finishScope(scope.getBlock());
   CHECK_ERR(expr);
 
diff --git a/test/lit/source-map.wast b/test/lit/source-map.wast
index 8d22360ad..2b4d3f5c7 100644
--- a/test/lit/source-map.wast
+++ b/test/lit/source-map.wast
@@ -1,7 +1,10 @@
 ;; RUN: wasm-opt %s -o %t.wasm -osm %t.map -g -q
 ;; RUN: wasm-opt %t.wasm -ism %t.map -q -o - -S | filecheck %s
 
+;; RUN: wasm-opt %s --new-wat-parser -S -o - | filecheck %s
+
 (module
+  ;;@ src.cpp:0:1
   (func $foo (param $x i32) (param $y i32)
     ;;@ src.cpp:10:1
     (if
@@ -12,8 +15,11 @@
         ;;@ src.cpp:40:1
         (local.get $y)
       )
+      ;; For the legacy parser
       ;;@ src.cpp:50:1
       (then
+        ;; For the new parser
+        ;;@ src.cpp:50:1
         (return)
       )
     )
@@ -24,10 +30,12 @@
       ;;@ src.cpp:80:1
       (local.get $y)
     )
+    ;;@ src.cpp:90:1
   )
 )
 
-;; CHECK:  (func $foo (param $x i32) (param $y i32)
+;; CHECK:       ;;@ src.cpp:0:1
+;; CHECK-NEXT:  (func $foo (param $x i32) (param $y i32)
 ;; CHECK-NEXT:   ;;@ src.cpp:10:1
 ;; CHECK-NEXT:   (if
 ;; CHECK-NEXT:    ;;@ src.cpp:20:1
@@ -49,3 +57,5 @@
 ;; CHECK-NEXT:    ;;@ src.cpp:80:1
 ;; CHECK-NEXT:    (local.get $y)
 ;; CHECK-NEXT:   )
+;; CHECK-NEXT:   ;;@ src.cpp:90:1
+;; CHECK-NEXT:  )