Use code units cache API (#2704)

* Upgrade Prism

* Fix type checking errors

* Use code units cache in indexing

* Add code_units_cache helper

* Use cached code units in range formatting

* Use code units cache in semantic highlighting

* Use code units cache for locating targets

Gemfile.lock

@@ -3,7 +3,7 @@ PATH
   specs:
     ruby-lsp (0.19.2)
       language_server-protocol (~> 3.17.0)
-      prism (>= 1.1, < 2.0)
+      prism (>= 1.2, < 2.0)
       rbs (>= 3, < 4)
       sorbet-runtime (>= 0.5.10782)
 
@@ -39,7 +39,7 @@ GEM
       ast (~> 2.4.1)
       racc
     prettier_print (1.2.1)
-    prism (1.1.0)
+    prism (1.2.0)
     psych (5.1.2)
       stringio
     racc (1.8.1)

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -33,7 +33,10 @@ def initialize(index, dispatcher, parse_result, file_path, collect_comments: fal
         T::Hash[Integer, Prism::Comment],
       )
       @inside_def = T.let(false, T::Boolean)
-      @encoding = T.let(@index.configuration.encoding, Encoding)
+      @code_units_cache = T.let(
+        parse_result.code_units_cache(@index.configuration.encoding),
+        T.any(T.proc.params(arg0: Integer).returns(Integer), Prism::CodeUnitsCache),
+      )
 
       # The nesting stack we're currently inside. Used to determine the fully qualified name of constants, but only
       # stored by unresolved aliases which need the original nesting to be lazily resolved
@@ -112,8 +115,8 @@ def on_class_node_enter(node)
       entry = Entry::Class.new(
         nesting,
         @file_path,
-        Location.from_prism_location(node.location, @encoding),
-        Location.from_prism_location(constant_path.location, @encoding),
+        Location.from_prism_location(node.location, @code_units_cache),
+        Location.from_prism_location(constant_path.location, @code_units_cache),
         comments,
         parent_class,
       )
@@ -141,8 +144,8 @@ def on_module_node_enter(node)
       entry = Entry::Module.new(
         actual_nesting(name),
         @file_path,
-        Location.from_prism_location(node.location, @encoding),
-        Location.from_prism_location(constant_path.location, @encoding),
+        Location.from_prism_location(node.location, @code_units_cache),
+        Location.from_prism_location(constant_path.location, @code_units_cache),
         comments,
       )
 
@@ -174,16 +177,16 @@ def on_singleton_class_node_enter(node)
         if existing_entries
           entry = T.must(existing_entries.first)
           entry.update_singleton_information(
-            Location.from_prism_location(node.location, @encoding),
-            Location.from_prism_location(expression.location, @encoding),
+            Location.from_prism_location(node.location, @code_units_cache),
+            Location.from_prism_location(expression.location, @code_units_cache),
             collect_comments(node),
           )
         else
           entry = Entry::SingletonClass.new(
             real_nesting,
             @file_path,
-            Location.from_prism_location(node.location, @encoding),
-            Location.from_prism_location(expression.location, @encoding),
+            Location.from_prism_location(node.location, @code_units_cache),
+            Location.from_prism_location(expression.location, @code_units_cache),
             collect_comments(node),
             nil,
           )
@@ -312,7 +315,7 @@ def on_call_node_enter(node)
       end
 
       @enhancements.each do |enhancement|
-        enhancement.on_call_node(@index, @owner_stack.last, node, @file_path)
+        enhancement.on_call_node(@index, @owner_stack.last, node, @file_path, @code_units_cache)
       rescue StandardError => e
         @indexing_errors << "Indexing error in #{@file_path} with '#{enhancement.class.name}' enhancement: #{e.message}"
       end
@@ -342,8 +345,8 @@ def on_def_node_enter(node)
         @index.add(Entry::Method.new(
           method_name,
           @file_path,
-          Location.from_prism_location(node.location, @encoding),
-          Location.from_prism_location(node.name_loc, @encoding),
+          Location.from_prism_location(node.location, @code_units_cache),
+          Location.from_prism_location(node.name_loc, @code_units_cache),
           comments,
           [Entry::Signature.new(list_params(node.parameters))],
           current_visibility,
@@ -358,8 +361,8 @@ def on_def_node_enter(node)
           @index.add(Entry::Method.new(
             method_name,
             @file_path,
-            Location.from_prism_location(node.location, @encoding),
-            Location.from_prism_location(node.name_loc, @encoding),
+            Location.from_prism_location(node.location, @code_units_cache),
+            Location.from_prism_location(node.name_loc, @code_units_cache),
             comments,
             [Entry::Signature.new(list_params(node.parameters))],
             current_visibility,
@@ -442,7 +445,7 @@ def on_alias_method_node_enter(node)
           node.old_name.slice,
           @owner_stack.last,
           @file_path,
-          Location.from_prism_location(node.new_name.location, @encoding),
+          Location.from_prism_location(node.new_name.location, @code_units_cache),
           comments,
         ),
       )
@@ -469,7 +472,7 @@ def handle_global_variable(node, loc)
       @index.add(Entry::GlobalVariable.new(
         name,
         @file_path,
-        Location.from_prism_location(loc, @encoding),
+        Location.from_prism_location(loc, @code_units_cache),
         comments,
       ))
     end
@@ -501,7 +504,7 @@ def handle_instance_variable(node, loc)
       @index.add(Entry::InstanceVariable.new(
         name,
         @file_path,
-        Location.from_prism_location(loc, @encoding),
+        Location.from_prism_location(loc, @code_units_cache),
         collect_comments(node),
         owner,
       ))
@@ -565,7 +568,7 @@ def handle_alias_method(node)
           old_name_value,
           @owner_stack.last,
           @file_path,
-          Location.from_prism_location(new_name.location, @encoding),
+          Location.from_prism_location(new_name.location, @code_units_cache),
           comments,
         ),
       )
@@ -601,7 +604,7 @@ def add_constant(node, name, value = nil)
             @stack.dup,
             name,
             @file_path,
-            Location.from_prism_location(node.location, @encoding),
+            Location.from_prism_location(node.location, @code_units_cache),
             comments,
           )
         when Prism::ConstantWriteNode, Prism::ConstantAndWriteNode, Prism::ConstantOrWriteNode,
@@ -614,7 +617,7 @@ def add_constant(node, name, value = nil)
             @stack.dup,
             name,
             @file_path,
-            Location.from_prism_location(node.location, @encoding),
+            Location.from_prism_location(node.location, @code_units_cache),
             comments,
           )
         when Prism::ConstantPathWriteNode, Prism::ConstantPathOrWriteNode, Prism::ConstantPathOperatorWriteNode,
@@ -625,14 +628,14 @@ def add_constant(node, name, value = nil)
             @stack.dup,
             name,
             @file_path,
-            Location.from_prism_location(node.location, @encoding),
+            Location.from_prism_location(node.location, @code_units_cache),
             comments,
           )
         else
           Entry::Constant.new(
             name,
             @file_path,
-            Location.from_prism_location(node.location, @encoding),
+            Location.from_prism_location(node.location, @code_units_cache),
             comments,
           )
         end,
@@ -700,7 +703,7 @@ def handle_attribute(node, reader:, writer:)
           @index.add(Entry::Accessor.new(
             name,
             @file_path,
-            Location.from_prism_location(loc, @encoding),
+            Location.from_prism_location(loc, @code_units_cache),
             comments,
             current_visibility,
             @owner_stack.last,
@@ -712,7 +715,7 @@ def handle_attribute(node, reader:, writer:)
         @index.add(Entry::Accessor.new(
           "#{name}=",
           @file_path,
-          Location.from_prism_location(loc, @encoding),
+          Location.from_prism_location(loc, @code_units_cache),
           comments,
           current_visibility,
           @owner_stack.last,

lib/ruby_indexer/lib/ruby_indexer/enhancement.rb

@@ -19,8 +19,12 @@ module Enhancement
         owner: T.nilable(Entry::Namespace),
         node: Prism::CallNode,
         file_path: String,
+        code_units_cache: T.any(
+          T.proc.params(arg0: Integer).returns(Integer),
+          Prism::CodeUnitsCache,
+        ),
       ).void
     end
-    def on_call_node(index, owner, node, file_path); end
+    def on_call_node(index, owner, node, file_path, code_units_cache); end
   end
 end

lib/ruby_indexer/lib/ruby_indexer/location.rb

@@ -8,13 +8,21 @@ class Location
     class << self
       extend T::Sig
 
-      sig { params(prism_location: Prism::Location, encoding: Encoding).returns(T.attached_class) }
-      def from_prism_location(prism_location, encoding)
+      sig do
+        params(
+          prism_location: Prism::Location,
+          code_units_cache: T.any(
+            T.proc.params(arg0: Integer).returns(Integer),
+            Prism::CodeUnitsCache,
+          ),
+        ).returns(T.attached_class)
+      end
+      def from_prism_location(prism_location, code_units_cache)
         new(
           prism_location.start_line,
           prism_location.end_line,
-          prism_location.start_code_units_column(encoding),
-          prism_location.end_code_units_column(encoding),
+          prism_location.cached_start_code_units_column(code_units_cache),
+          prism_location.cached_end_code_units_column(code_units_cache),
         )
       end
     end

lib/ruby_indexer/test/enhancements_test.rb

@@ -9,14 +9,14 @@ def test_enhancing_indexing_included_hook
       enhancement_class = Class.new do
         include Enhancement
 
-        def on_call_node(index, owner, node, file_path)
+        def on_call_node(index, owner, node, file_path, code_units_cache)
           return unless owner
           return unless node.name == :extend
 
           arguments = node.arguments&.arguments
           return unless arguments
 
-          location = Location.from_prism_location(node.location, index.configuration.encoding)
+          location = Location.from_prism_location(node.location, code_units_cache)
 
           arguments.each do |node|
             next unless node.is_a?(Prism::ConstantReadNode) || node.is_a?(Prism::ConstantPathNode)
@@ -101,7 +101,7 @@ def test_enhancing_indexing_configuration_dsl
       enhancement_class = Class.new do
         include Enhancement
 
-        def on_call_node(index, owner, node, file_path)
+        def on_call_node(index, owner, node, file_path, code_units_cache)
           return unless owner
 
           name = node.name
@@ -113,7 +113,7 @@ def on_call_node(index, owner, node, file_path)
           association_name = arguments.first
           return unless association_name.is_a?(Prism::SymbolNode)
 
-          location = Location.from_prism_location(association_name.location, index.configuration.encoding)
+          location = Location.from_prism_location(association_name.location, code_units_cache)
 
           index.add(Entry::Method.new(
             T.must(association_name.value),
@@ -164,7 +164,7 @@ def test_error_handling_in_enhancement
       enhancement_class = Class.new do
         include Enhancement
 
-        def on_call_node(index, owner, node, file_path)
+        def on_call_node(index, owner, node, file_path, code_units_cache)
           raise "Error"
         end
 

lib/ruby_lsp/erb_document.rb

@@ -11,12 +11,24 @@ class ERBDocument < Document
     sig { returns(String) }
     attr_reader :host_language_source
 
+    sig do
+      returns(T.any(
+        T.proc.params(arg0: Integer).returns(Integer),
+        Prism::CodeUnitsCache,
+      ))
+    end
+    attr_reader :code_units_cache
+
     sig { params(source: String, version: Integer, uri: URI::Generic, encoding: Encoding).void }
     def initialize(source:, version:, uri:, encoding: Encoding::UTF_8)
       # This has to be initialized before calling super because we call `parse` in the parent constructor, which
       # overrides this with the proper virtual host language source
       @host_language_source = T.let("", String)
       super
+      @code_units_cache = T.let(@parse_result.code_units_cache(@encoding), T.any(
+        T.proc.params(arg0: Integer).returns(Integer),
+        Prism::CodeUnitsCache,
+      ))
     end
 
     sig { override.returns(T::Boolean) }
@@ -30,6 +42,7 @@ def parse!
       # Use partial script to avoid syntax errors in ERB files where keywords may be used without the full context in
       # which they will be evaluated
       @parse_result = Prism.parse(scanner.ruby, partial_script: true)
+      @code_units_cache = @parse_result.code_units_cache(@encoding)
       true
     end
 
@@ -53,8 +66,8 @@ def locate_node(position, node_types: [])
       RubyDocument.locate(
         @parse_result.value,
         create_scanner.find_char_position(position),
+        code_units_cache: @code_units_cache,
         node_types: node_types,
-        encoding: @encoding,
       )
     end
 

lib/ruby_lsp/listeners/folding_ranges.rb

@@ -239,10 +239,10 @@ def add_statements_range(node)
         statements = node.statements
         return unless statements
 
-        body = statements.body
-        return if body.empty?
+        statement = statements.body.last
+        return unless statement
 
-        add_lines_range(node.location.start_line, body.last.location.end_line)
+        add_lines_range(node.location.start_line, statement.location.end_line)
       end
 
       sig { params(node: Prism::Node).void }

lib/ruby_lsp/requests/code_action_resolve.rb

@@ -99,7 +99,13 @@ def refactor_variable
 
         # Find the closest statements node, so that we place the refactor in a valid position
         node_context = RubyDocument
-          .locate(@document.parse_result.value, start_index, node_types: [Prism::StatementsNode, Prism::BlockNode])
+          .locate(@document.parse_result.value,
+            start_index,
+            node_types: [
+              Prism::StatementsNode,
+              Prism::BlockNode,
+            ],
+            code_units_cache: @document.code_units_cache)
 
         closest_statements = node_context.node
         parent_statements = node_context.parent
@@ -196,7 +202,7 @@ def refactor_method
           @document.parse_result.value,
           start_index,
           node_types: [Prism::DefNode],
-          encoding: @global_state.encoding,
+          code_units_cache: @document.code_units_cache,
         )
         closest_node = node_context.node
         return Error::InvalidTargetRange unless closest_node

lib/ruby_lsp/requests/completion.rb

@@ -57,7 +57,7 @@ def initialize(document, global_state, params, sorbet_level, dispatcher)
             Prism::InstanceVariableTargetNode,
             Prism::InstanceVariableWriteNode,
           ],
-          encoding: global_state.encoding,
+          code_units_cache: document.code_units_cache,
         )
         @response_builder = T.let(
           ResponseBuilders::CollectionResponseBuilder[Interface::CompletionItem].new,

lib/ruby_lsp/requests/definition.rb

@@ -58,7 +58,7 @@ def initialize(document, global_state, position, dispatcher, sorbet_level)
             Prism::SuperNode,
             Prism::ForwardingSuperNode,
           ],
-          encoding: global_state.encoding,
+          code_units_cache: document.code_units_cache,
         )
 
         target = node_context.node

lib/ruby_lsp/requests/document_highlight.rb

@@ -28,7 +28,11 @@ def initialize(global_state, document, position, dispatcher)
         char_position = document.create_scanner.find_char_position(position)
         delegate_request_if_needed!(global_state, document, char_position)
 
-        node_context = RubyDocument.locate(document.parse_result.value, char_position, encoding: global_state.encoding)
+        node_context = RubyDocument.locate(
+          document.parse_result.value,
+          char_position,
+          code_units_cache: document.code_units_cache,
+        )
 
         @response_builder = T.let(
           ResponseBuilders::CollectionResponseBuilder[Interface::DocumentHighlight].new,

lib/ruby_lsp/requests/hover.rb

@@ -41,7 +41,7 @@ def initialize(document, global_state, position, dispatcher, sorbet_level)
           document.parse_result.value,
           char_position,
           node_types: Listeners::Hover::ALLOWED_TARGETS,
-          encoding: global_state.encoding,
+          code_units_cache: document.code_units_cache,
         )
         target = node_context.node
         parent = node_context.parent