1 files changed, 29 insertions, 10 deletions

diff --git a/lib/syntax_suggest/lex_all.rb b/lib/syntax_suggest/lex_all.rb
index 132cba9f5d..c16fbb52d3 100644
--- a/lib/syntax_suggest/lex_all.rb
+++ b/lib/syntax_suggest/lex_all.rb
@@ -3,34 +3,53 @@
 module SyntaxSuggest
   # Ripper.lex is not guaranteed to lex the entire source document
   #
-  # lex = LexAll.new(source: source)
-  # lex.each do |value|
-  #   puts value.line
-  # end
+  # This class guarantees the whole document is lex-ed by iteratively
+  # lexing the document where ripper stopped.
+  #
+  # Prism likely doesn't have the same problem. Once ripper support is removed
+  # we can likely reduce the complexity here if not remove the whole concept.
+  #
+  # Example usage:
+  #
+  #   lex = LexAll.new(source: source)
+  #   lex.each do |value|
+  #     puts value.line
+  #   end
   class LexAll
     include Enumerable
 
     def initialize(source:, source_lines: nil)
-      @lex = Ripper::Lexer.new(source, "-", 1).parse.sort_by(&:pos)
-      lineno = @lex.last.pos.first + 1
+      @lex = self.class.lex(source, 1)
+      lineno = @lex.last[0][0] + 1
       source_lines ||= source.lines
       last_lineno = source_lines.length
 
       until lineno >= last_lineno
-        lines = source_lines[lineno..-1]
+        lines = source_lines[lineno..]
 
         @lex.concat(
-          Ripper::Lexer.new(lines.join, "-", lineno + 1).parse.sort_by(&:pos)
+          self.class.lex(lines.join, lineno + 1)
         )
-        lineno = @lex.last.pos.first + 1
+
+        lineno = @lex.last[0].first + 1
       end
 
       last_lex = nil
       @lex.map! { |elem|
-        last_lex = LexValue.new(elem.pos.first, elem.event, elem.tok, elem.state, last_lex)
+        last_lex = LexValue.new(elem[0].first, elem[1], elem[2], elem[3], last_lex)
       }
     end
 
+    if SyntaxSuggest.use_prism_parser?
+      def self.lex(source, line_number)
+        Prism.lex_compat(source, line: line_number).value.sort_by { |values| values[0] }
+      end
+    else
+      def self.lex(source, line_number)
+        Ripper::Lexer.new(source, "-", line_number).parse.sort_by(&:pos)
+      end
+    end
+
     def to_a
       @lex
     end