[ruby/prism] Fix lexing for unterminated strings/heredocs etc.

When we hit EOF and still have lex modes left, it means some content was unterminated.
Heredocs specifically have logic that needs to happen when the body finished lexing.
If we don't reset the mode back to how it was before, it will not continue lexing at the correct place.

Followup to https://github.com/ruby/prism/pull/3918.
We can't call into `parser_lex` since it resets token locations.

https://github.com/ruby/prism/commit/27c24fdc0d

4 files changed, 92 insertions, 3 deletions

diff --git a/test/prism/errors/unterminated_heredoc_and_embexpr.txt b/test/prism/errors/unterminated_heredoc_and_embexpr.txt
new file mode 100644
index 0000000000..bed7fcd24e
--- /dev/null
+++ b/test/prism/errors/unterminated_heredoc_and_embexpr.txt
@@ -0,0 +1,11 @@
+<<A+B
+  ^ unterminated heredoc; can't find string "A" anywhere before EOF
+   ^ unexpected '+', ignoring it
+  ^ unterminated heredoc; can't find string "A" anywhere before EOF
+#{C
+   ^ unexpected heredoc ending; expected an argument
+   ^ unexpected heredoc ending, expecting end-of-input
+   ^ unexpected heredoc ending, ignoring it
+   ^ unexpected end-of-input, assuming it is closing the parent top level context
+^ expected a `}` to close the embedded expression
+
diff --git a/test/prism/errors/unterminated_heredoc_and_embexpr_2.txt b/test/prism/errors/unterminated_heredoc_and_embexpr_2.txt
new file mode 100644
index 0000000000..a03ff1d212
--- /dev/null
+++ b/test/prism/errors/unterminated_heredoc_and_embexpr_2.txt
@@ -0,0 +1,9 @@
+<<A+B
+  ^ unterminated heredoc; can't find string "A" anywhere before EOF
+#{C + "#{"}
+           ^ unterminated string meets end of file
+           ^ unexpected end-of-input, assuming it is closing the parent top level context
+           ^ expected a `}` to close the embedded expression
+      ^ unterminated string; expected a closing delimiter for the interpolated string
+           ^ expected a `}` to close the embedded expression
+
diff --git a/test/prism/errors_test.rb b/test/prism/errors_test.rb
index b30a0f304d..27610e89d3 100644
--- a/test/prism/errors_test.rb
+++ b/test/prism/errors_test.rb
@@ -87,6 +87,28 @@ module Prism
       assert_nil(statement.end_keyword)
     end
 
+    def test_unclosed_interpolation
+      statement = Prism.parse_statement("\"\#{")
+      assert_equal('"', statement.opening)
+      assert_nil(statement.closing)
+
+      assert_equal(1, statement.parts.count)
+      assert_equal('#{', statement.parts[0].opening)
+      assert_equal("", statement.parts[0].closing)
+      assert_nil(statement.parts[0].statements)
+    end
+
+    def test_unclosed_heredoc_and_interpolation
+      statement = Prism.parse_statement("<<D\n\#{")
+      assert_equal("<<D", statement.opening)
+      assert_nil(statement.closing)
+
+      assert_equal(1, statement.parts.count)
+      assert_equal('#{', statement.parts[0].opening)
+      assert_equal("", statement.parts[0].closing)
+      assert_nil(statement.parts[0].statements)
+    end
+
     private
 
     def assert_errors(filepath, version)
diff --git a/test/prism/lex_test.rb b/test/prism/lex_test.rb
index 9a9f203c28..8ea7ce7e9b 100644
--- a/test/prism/lex_test.rb
+++ b/test/prism/lex_test.rb
@@ -48,11 +48,58 @@ module Prism
     end
 
     if RUBY_VERSION >= "3.3"
-      def test_lex_compare
-        prism = Prism.lex_compat(File.read(__FILE__), version: "current").value
-        ripper = Ripper.lex(File.read(__FILE__))
+      def test_lex_compat
+        source = "foo bar"
+        prism = Prism.lex_compat(source, version: "current").value
+        ripper = Ripper.lex(source)
         assert_equal(ripper, prism)
       end
     end
+
+    def test_lex_interpolation_unterminated
+      assert_equal(
+        %i[STRING_BEGIN EMBEXPR_BEGIN EOF],
+        token_types('"#{')
+      )
+
+      assert_equal(
+        %i[STRING_BEGIN EMBEXPR_BEGIN IGNORED_NEWLINE EOF],
+        token_types('"#{' + "\n")
+      )
+    end
+
+    def test_lex_interpolation_unterminated_with_content
+      # FIXME: Emits EOL twice.
+      assert_equal(
+        %i[STRING_BEGIN EMBEXPR_BEGIN CONSTANT EOF EOF],
+        token_types('"#{C')
+      )
+
+      assert_equal(
+        %i[STRING_BEGIN EMBEXPR_BEGIN CONSTANT NEWLINE EOF],
+        token_types('"#{C' + "\n")
+      )
+    end
+
+    def test_lex_heredoc_unterminated
+      code = <<~'RUBY'.strip
+        <<A+B
+        #{C
+      RUBY
+
+      assert_equal(
+        %i[HEREDOC_START EMBEXPR_BEGIN CONSTANT HEREDOC_END PLUS CONSTANT NEWLINE EOF],
+        token_types(code)
+      )
+
+      assert_equal(
+        %i[HEREDOC_START EMBEXPR_BEGIN CONSTANT NEWLINE HEREDOC_END PLUS CONSTANT NEWLINE EOF],
+        token_types(code + "\n")
+      )
+    end
+
+    def token_types(code)
+      Prism.lex(code).value.map { |token, _state| token.type }
+    end
   end
 end