[ruby/prism] Fix comments after HEREDOCs again.

The problem was deeper than just looking back a single token.
You can push the heredoc_end token way back into the list.
We need to save the last location of a heredoc end to see if
it's the last token in the file.

Fixes https://github.com/ruby/prism/pull/1954

https://github.com/ruby/prism/commit/91dfd4eecd

2 files changed, 10 insertions, 14 deletions

diff --git a/lib/prism/lex_compat.rb b/lib/prism/lex_compat.rb
index 0336f48d6d..2c23933714 100644
--- a/lib/prism/lex_compat.rb
+++ b/lib/prism/lex_compat.rb
@@ -610,6 +610,7 @@ module Prism
       result = Prism.lex(source, **options)
       result_value = result.value
       previous_state = nil
+      last_heredoc_end = nil
 
       # In previous versions of Ruby, Ripper wouldn't flush the bom before the
       # first token, so we had to have a hack in place to account for that. This
@@ -664,6 +665,7 @@ module Prism
           when :on_heredoc_end
             # Heredoc end tokens can be emitted in an odd order, so we don't
             # want to bother comparing the state on them.
+            last_heredoc_end = token.location.end_offset
             IgnoreStateToken.new([[lineno, column], event, value, lex_state])
           when :on_ident
             if lex_state == Ripper::EXPR_END
@@ -730,20 +732,13 @@ module Prism
             # Ripper will append a on_nl token (even though there isn't
             # necessarily a newline). We mirror that here.
             if previous_token.type == :COMMENT
-              # If the token before the comment was a heredoc end, then
-              # the comment's end_offset is before the heredoc end token.
+              # If the comment is at the start of a heredoc: <<HEREDOC # comment
+              # then the comment's end_offset is up near the heredoc_beg.
               # This is not the correct offset to use for figuring out if
-              # there is trailing whitespace after the comment.
-              # Use the end_offset of the heredoc end instead.
-              before_comment = result_value[index - 2]
-              before_comment &&= before_comment[0]
-
-              if before_comment&.type == :HEREDOC_END
-                start_offset = before_comment.location.end_offset
-              else
-                start_offset = previous_token.location.end_offset
-              end
-
+              # there is trailing whitespace after the last token.
+              # Use the greater offset of the two to determine the start of
+              # the trailing whitespace.
+              start_offset = [previous_token.location.end_offset, last_heredoc_end].compact.max
               end_offset = token.location.start_offset
 
               if start_offset < end_offset
diff --git a/test/prism/fixtures/heredoc_with_comment.txt b/test/prism/fixtures/heredoc_with_comment.txt
index cf48c12051..8e7b7275bd 100644
--- a/test/prism/fixtures/heredoc_with_comment.txt
+++ b/test/prism/fixtures/heredoc_with_comment.txt
@@ -1,2 +1,3 @@
-<<-TARGET # comment

+<<-TARGET.chomp # the heredoc end token doesn't always precede the comment

+  content makes for an obvious error

 TARGET
\ No newline at end of file