4 files changed, 35 insertions, 16 deletions

diff --git a/src/parser/contexts.h b/src/parser/contexts.h
index 81537abaf..cead35f60 100644
--- a/src/parser/contexts.h
+++ b/src/parser/contexts.h
@@ -912,7 +912,7 @@ struct ParseDeclsCtx : NullTypeParserCtx, NullInstrParserCtx {
     return Ok{};
   }
 
-  ParseDeclsCtx(std::string_view in, Module& wasm) : in(in), wasm(wasm) {}
+  ParseDeclsCtx(Lexer& in, Module& wasm) : in(in), wasm(wasm) {}
 
   void addFuncType(SignatureT) {}
   void addContType(ContinuationT) {}
@@ -1049,9 +1049,7 @@ struct ParseTypeDefsCtx : TypeParserCtx<ParseTypeDefsCtx> {
   // The index of the subtype definition we are parsing.
   Index index = 0;
 
-  ParseTypeDefsCtx(std::string_view in,
-                   TypeBuilder& builder,
-                   const IndexMap& typeIndices)
+  ParseTypeDefsCtx(Lexer& in, TypeBuilder& builder, const IndexMap& typeIndices)
     : TypeParserCtx<ParseTypeDefsCtx>(typeIndices), in(in), builder(builder),
       names(builder.size()) {}
 
@@ -1121,7 +1119,7 @@ struct ParseImplicitTypeDefsCtx : TypeParserCtx<ParseImplicitTypeDefsCtx> {
   // Map signatures to the first defined heap type they match.
   std::unordered_map<Signature, HeapType> sigTypes;
 
-  ParseImplicitTypeDefsCtx(std::string_view in,
+  ParseImplicitTypeDefsCtx(Lexer& in,
                            std::vector<HeapType>& types,
                            std::unordered_map<Index, HeapType>& implicitTypes,
                            const IndexMap& typeIndices)
@@ -1192,7 +1190,7 @@ struct ParseModuleTypesCtx : TypeParserCtx<ParseModuleTypesCtx>,
   Index index = 0;
 
   ParseModuleTypesCtx(
-    std::string_view in,
+    Lexer& in,
     Module& wasm,
     const std::vector<HeapType>& types,
     const std::unordered_map<Index, HeapType>& implicitTypes,
@@ -1397,7 +1395,7 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
   }
 
   ParseDefsCtx(
-    std::string_view in,
+    Lexer& in,
     Module& wasm,
     const std::vector<HeapType>& types,
     const std::unordered_map<Index, HeapType>& implicitTypes,
diff --git a/src/parser/lexer.h b/src/parser/lexer.h
index 1a93d3e99..e601091db 100644
--- a/src/parser/lexer.h
+++ b/src/parser/lexer.h
@@ -157,12 +157,13 @@ extern Name srcAnnotationKind;
 
 struct Lexer {
 private:
-  std::string_view buffer;
   size_t index = 0;
   std::optional<Token> curr;
   std::vector<Annotation> annotations;
 
 public:
+  std::string_view buffer;
+
   Lexer(std::string_view buffer) : buffer(buffer) { setIndex(0); }
 
   size_t getIndex() const { return index; }
diff --git a/src/parser/wat-parser.cpp b/src/parser/wat-parser.cpp
index 7f6dd2975..cc7d87540 100644
--- a/src/parser/wat-parser.cpp
+++ b/src/parser/wat-parser.cpp
@@ -99,17 +99,11 @@ void propagateDebugLocations(Module& wasm) {
   runner.run();
 }
 
-// ================
-// Parser Functions
-// ================
-
-} // anonymous namespace
-
-Result<> parseModule(Module& wasm, std::string_view input) {
+Result<> doParseModule(Module& wasm, Lexer& input, bool allowExtra) {
   // Parse module-level declarations.
   ParseDeclsCtx decls(input, wasm);
   CHECK_ERR(module(decls));
-  if (!decls.in.empty()) {
+  if (!allowExtra && !decls.in.empty()) {
     return decls.in.err("Unexpected tokens after module");
   }
 
@@ -222,8 +216,27 @@ Result<> parseModule(Module& wasm, std::string_view input) {
   }
 
   propagateDebugLocations(wasm);
+  input = decls.in;
 
   return Ok{};
 }
 
+} // anonymous namespace
+
+Result<> parseModule(Module& wasm, std::string_view in) {
+  Lexer lexer(in);
+  return doParseModule(wasm, lexer, false);
+}
+
+Result<> parseModule(Module& wasm, Lexer& lexer) {
+  return doParseModule(wasm, lexer, true);
+}
+
+Result<Expression*> parseExpression(Module& wasm, Lexer& lexer) {
+  ParseDefsCtx ctx(lexer, wasm, {}, {}, {}, {}, {});
+  auto e = expr(ctx);
+  CHECK_ERR(e);
+  return *e;
+}
+
 } // namespace wasm::WATParser
diff --git a/src/parser/wat-parser.h b/src/parser/wat-parser.h
index b31523af9..3f7dd64c4 100644
--- a/src/parser/wat-parser.h
+++ b/src/parser/wat-parser.h
@@ -19,6 +19,7 @@
 
 #include <string_view>
 
+#include "parser/lexer.h"
 #include "support/result.h"
 #include "wasm.h"
 
@@ -27,6 +28,12 @@ namespace wasm::WATParser {
 // Parse a single WAT module.
 Result<> parseModule(Module& wasm, std::string_view in);
 
+// Parse a single WAT module that may have other things after it, as in a wast
+// file.
+Result<> parseModule(Module& wasm, Lexer& lexer);
+
+Result<Expression*> parseExpression(Module& wasm, Lexer& lexer);
+
 } // namespace wasm::WATParser
 
 #endif // parser_wat_parser_h