From 89ecbba4c8a687d449f2f9dc12817c256c4c9571 Mon Sep 17 00:00:00 2001
From: David Anderson <dvander@alliedmods.net>
Date: Fri, 20 Oct 2023 21:20:49 -0700
Subject: [PATCH] Clean up injected token state.

Make sure that a successful function parse actually terminates the
injected token stream. Also, drop the injected token stream in case
parsing failed.
---
 compiler/lexer.cpp  | 20 ++++++++++++++++++--
 compiler/lexer.h    | 11 +++++++++--
 compiler/parser.cpp | 12 ++++++++++--
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/compiler/lexer.cpp b/compiler/lexer.cpp
index 08f1e8df8..dd1c9a363 100644
--- a/compiler/lexer.cpp
+++ b/compiler/lexer.cpp
@@ -1319,12 +1319,21 @@ int Lexer::lex() {
         return current_token()->id;
     }
 
-    if (!injected_token_stream_.empty())
-        return LexInjectedToken();
+    if (using_injected_tokens_) {
+        if (!injected_token_stream_.empty())
+            return LexInjectedToken();
+        return 0;
+    }
 
     return LexNewToken();
 }
 
+bool Lexer::freading() const {
+    if (using_injected_tokens_)
+        return !injected_token_stream_.empty();
+    return freading_;
+}
+
 int Lexer::LexNewToken() {
     full_token_t* tok = advance_token_ptr();
     *tok = {};
@@ -2579,6 +2588,7 @@ void Lexer::AssertCleanState() {
     assert(!in_string_continuation_);
     assert(allow_tags_);
     assert(injected_token_stream_.empty());
+    assert(!using_injected_tokens_);
 }
 
 TokenCache* Lexer::LexFunctionBody() {
@@ -2620,10 +2630,16 @@ void Lexer::InjectCachedTokens(TokenCache* cache) {
     AssertCleanState();
 
     injected_token_stream_ = std::move(cache->tokens);
+    using_injected_tokens_ = true;
     token_caches_.remove(cache);
     delete cache;
 
     freading_ = true;
 }
 
+void Lexer::DiscardCachedTokens() {
+    using_injected_tokens_ = false;
+    injected_token_stream_.clear();
+}
+
 } // namespace sp
diff --git a/compiler/lexer.h b/compiler/lexer.h
index 193358a10..bf1787c62 100644
--- a/compiler/lexer.h
+++ b/compiler/lexer.h
@@ -331,6 +331,9 @@ class Lexer
     // be replayed by lex().
     void InjectCachedTokens(TokenCache* cache);
 
+    // Throw away tokens injected by InjectCachedTokens().
+    void DiscardCachedTokens();
+
     full_token_t lex_tok() {
         lex();
         return *current_token();
@@ -344,7 +347,7 @@ class Lexer
     bool& allow_tags() { return allow_tags_; }
     bool& require_newdecls() { return state_.require_newdecls; }
     bool& need_semicolon() { return state_.need_semicolon; }
-    bool freading() const { return freading_; }
+    bool freading() const;
     int fcurrent() const { return state_.inpf->sources_index(); }
     unsigned fline() const { return state_.fline; }
     SourceFile* inpf() const { return state_.inpf.get(); }
@@ -513,9 +516,13 @@ class Lexer
 
     LexerState state_;
     tr::vector<LexerState> prev_state_;
+
+    // Set if tokens are being lexed into a new token cache.
+    bool caching_tokens_ = false;
     ke::InlineList<TokenCache> token_caches_;
+
     std::deque<full_token_t> injected_token_stream_;
-    bool caching_tokens_ = false;
+    bool using_injected_tokens_ = false;
 };
 
 std::string StringizePath(const std::filesystem::path& in_path);
diff --git a/compiler/parser.cpp b/compiler/parser.cpp
index 5e6284ff7..250c8d81f 100644
--- a/compiler/parser.cpp
+++ b/compiler/parser.cpp
@@ -193,8 +193,16 @@ Parser::Parse()
                                              tokens->need_semicolon);
 
         lexer_->InjectCachedTokens(tokens);
-        auto body = parse_stmt(false);
-        fun->set_body(BlockStmt::WrapStmt(body));
+
+        AutoCountErrors errors;
+        if (auto body = parse_stmt(false)) {
+            fun->set_body(BlockStmt::WrapStmt(body));
+
+            // If there were no errors, we should have consumed every token.
+            assert(!errors.ok() || !lexer_->freading());
+        }
+
+        lexer_->DiscardCachedTokens();
     }
 
     auto list = new StmtList(token_pos_t{}, stmts);