Translate existing AST into RuboCop AST

Author: vinistock

.rubocop.yml

@@ -10,6 +10,7 @@ require:
   - ./lib/rubocop/cop/ruby_lsp/use_register_with_handler_method
 
 AllCops:
+  ParserEngine: parser_prism
   NewCops: disable
   SuggestExtensions: false
   Include:
@@ -42,6 +43,7 @@ Sorbet/TrueSigil:
     - "lib/ruby_indexer/test/**/*.rb"
     - "lib/ruby_indexer/lib/ruby_indexer/prefix_tree.rb"
     - "lib/ruby_lsp/load_sorbet.rb"
+    - "lib/ruby_lsp/requests/support/ast_translation.rb"
   Exclude:
     - "**/*.rake"
     - "lib/**/*.rb"
@@ -57,3 +59,4 @@ Sorbet/StrictSigil:
     - "lib/ruby-lsp.rb"
     - "lib/ruby_indexer/lib/ruby_indexer/prefix_tree.rb"
     - "lib/ruby_lsp/load_sorbet.rb"
+    - "lib/ruby_lsp/requests/support/ast_translation.rb"

lib/ruby_lsp/document.rb

@@ -8,7 +8,7 @@ class Document
 
     abstract!
 
-    sig { returns(Prism::ParseResult) }
+    sig { returns(Prism::ParseLexResult) }
     attr_reader :parse_result
 
     sig { returns(String) }
@@ -36,7 +36,7 @@ def initialize(source:, version:, uri:, encoding: Encoding::UTF_8)
 
     sig { returns(Prism::ProgramNode) }
     def tree
-      @parse_result.value
+      @parse_result.value.first
     end
 
     sig { returns(T::Array[Prism::Comment]) }
@@ -113,7 +113,7 @@ def create_scanner
       ).returns([T.nilable(Prism::Node), T.nilable(Prism::Node), T::Array[String]])
     end
     def locate_node(position, node_types: [])
-      locate(@parse_result.value, create_scanner.find_char_position(position), node_types: node_types)
+      locate(tree, create_scanner.find_char_position(position), node_types: node_types)
     end
 
     sig do

lib/ruby_lsp/requests/support/ast_translation.rb

@@ -0,0 +1,100 @@
+# typed: true
+# frozen_string_literal: true
+
+begin
+  gem("rubocop", ">= 1.63.0")
+rescue LoadError
+  $stderr.puts("AST translation turned off because RuboCop >= 1.63.0 is required")
+  return
+end
+
+require "prism/translation/parser/rubocop"
+
+# Processed Source patch so that we can pass the existing AST to RuboCop without having to re-parse files a second time
+module ProcessedSourcePatch
+  extend T::Sig
+
+  sig do
+    params(
+      source: String,
+      ruby_version: Float,
+      path: T.nilable(String),
+      parser_engine: Symbol,
+      prism_result: T.nilable(Prism::ParseLexResult),
+    ).void
+  end
+  def initialize(source, ruby_version, path = nil, parser_engine: :parser_whitequark, prism_result: nil)
+    @prism_result = prism_result
+
+    # Invoking super will end up invoking our patched version of tokenize, which avoids re-parsing the file
+    super(source, ruby_version, path, parser_engine: parser_engine)
+  end
+
+  sig { params(parser: T.untyped).returns(T::Array[T.untyped]) }
+  def tokenize(parser)
+    begin
+      # This is where we need to pass the existing result to prevent a re-parse
+      ast, comments, tokens = parser.tokenize(@buffer, parse_result: @prism_result)
+
+      ast ||= nil
+    rescue Parser::SyntaxError
+      comments = []
+      tokens = []
+    end
+
+    ast&.complete!
+    tokens.map! { |t| RuboCop::AST::Token.from_parser_token(t) }
+
+    [ast, comments, tokens]
+  end
+
+  RuboCop::AST::ProcessedSource.prepend(self)
+end
+
+# This patch allows Prism's translation parser to accept an existing AST in `tokenize`. This doesn't match the original
+# signature of RuboCop itself, but there's no other way to allow reusing the AST
+module TranslatorPatch
+  extend T::Sig
+  extend T::Helpers
+
+  requires_ancestor { Prism::Translation::Parser }
+
+  sig do
+    params(
+      source_buffer: ::Parser::Source::Buffer,
+      recover: T::Boolean,
+      parse_result: T.nilable(Prism::ParseLexResult),
+    ).returns(T::Array[T.untyped])
+  end
+  def tokenize(source_buffer, recover = false, parse_result: nil)
+    @source_buffer = source_buffer
+    source = source_buffer.source
+
+    offset_cache = build_offset_cache(source)
+    result = if @prism_result
+      @prism_result
+    else
+      begin
+        unwrap(
+          Prism.parse_lex(source, filepath: source_buffer.name, version: convert_for_prism(version)),
+          offset_cache,
+        )
+      rescue ::Parser::SyntaxError
+        raise unless recover
+      end
+    end
+
+    program, tokens = result.value
+    ast = build_ast(program, offset_cache) if result.success?
+
+    [
+      ast,
+      build_comments(result.comments, offset_cache),
+      build_tokens(tokens, offset_cache),
+    ]
+  ensure
+    @source_buffer = nil
+  end
+
+  Prism::Translation::Parser.prepend(self)
+end

lib/ruby_lsp/requests/support/rubocop_formatter.rb

@@ -22,7 +22,7 @@ def run_formatting(uri, document)
           filename = T.must(uri.to_standardized_path || uri.opaque)
 
           # Invoke RuboCop with just this file in `paths`
-          @format_runner.run(filename, document.source)
+          @format_runner.run(filename, document.source, document.parse_result)
           @format_runner.formatted_source
         end
 
@@ -35,7 +35,7 @@ def run_formatting(uri, document)
         def run_diagnostic(uri, document)
           filename = T.must(uri.to_standardized_path || uri.opaque)
           # Invoke RuboCop with just this file in `paths`
-          @diagnostic_runner.run(filename, document.source)
+          @diagnostic_runner.run(filename, document.source, document.parse_result)
 
           @diagnostic_runner.offenses.map do |offense|
             Support::RuboCopDiagnostic.new(document, offense, uri).to_lsp_diagnostic

lib/ruby_lsp/requests/support/rubocop_runner.rb

@@ -17,6 +17,8 @@
   RuboCop::LSP.enable
 end
 
+require "ruby_lsp/requests/support/ast_translation"
+
 module RubyLsp
   module Requests
     module Support
@@ -74,6 +76,7 @@ def initialize(*args)
           @offenses = T.let([], T::Array[RuboCop::Cop::Offense])
           @errors = T.let([], T::Array[String])
           @warnings = T.let([], T::Array[String])
+          @parse_result = T.let(nil, T.nilable(Prism::ParseLexResult))
 
           args += DEFAULT_ARGS
           rubocop_options = ::RuboCop::Options.new.parse(args).first
@@ -82,14 +85,15 @@ def initialize(*args)
           super(rubocop_options, config_store)
         end
 
-        sig { params(path: String, contents: String).void }
-        def run(path, contents)
+        sig { params(path: String, contents: String, parse_result: Prism::ParseLexResult).void }
+        def run(path, contents, parse_result)
           # Clear Runner state between runs since we get a single instance of this class
           # on every use site.
           @errors = []
           @warnings = []
           @offenses = []
           @options[:stdin] = contents
+          @parse_result = parse_result
 
           super([path])
 
@@ -109,6 +113,29 @@ def formatted_source
           @options[:stdin]
         end
 
+        sig { params(file: String).returns(RuboCop::ProcessedSource) }
+        def get_processed_source(file)
+          config = @config_store.for_file(file)
+          parser_engine = config.parser_engine
+          return super unless parser_engine == :parser_prism
+
+          processed_source = T.unsafe(::RuboCop::AST::ProcessedSource).new(
+            @options[:stdin],
+            3.3,
+            file,
+            parser_engine: parser_engine,
+            prism_result: @parse_result,
+          )
+          processed_source.config = config
+          processed_source.registry = mobilized_cop_classes(config)
+
+          # We have to reset the result to nil after returning the processed source the first time. This is needed for
+          # formatting because RuboCop will keep re-parsing the same file until no more auto-corrects can be applied. If
+          # we didn't reset it, we would end up operating in a stale AST
+          @parse_result = nil
+          processed_source
+        end
+
         class << self
           extend T::Sig
 

sorbet/rbi/shims/rubocop.rbi

@@ -0,0 +1,7 @@
+# typed: true
+
+class RuboCop::Runner
+  def initialize(options, config_store)
+    @config_store = T.let(T.unsafe(nil), RuboCop::ConfigStore)
+  end
+end

test/ruby_document_test.rb

@@ -514,14 +514,14 @@ def test_reparsing_without_new_edits_does_nothing
       version: 2,
     )
 
-    parse_result = Prism.parse(text)
+    parse_result = Prism.parse_lex(text)
 
     # When there's a new edit, we parse it the first `parse` invocation
-    Prism.expects(:parse).with(document.source).once.returns(parse_result)
+    Prism.expects(:parse_lex).with(document.source).once.returns(parse_result)
     document.parse
 
     # If there are no new edits, we don't do anything
-    Prism.expects(:parse).never
+    Prism.expects(:parse_lex).never
     document.parse
 
     document.push_edits(
@@ -530,7 +530,7 @@ def test_reparsing_without_new_edits_does_nothing
     )
 
     # If there's another edit, we parse it once again
-    Prism.expects(:parse).with(document.source).once.returns(parse_result)
+    Prism.expects(:parse_lex).with(document.source).once.returns(parse_result)
     document.parse
   end