Use code units cache in semantic highlighting

Shopify · Oct 10, 2024 · cd7e682 · cd7e682
1 parent 9e47368
commit cd7e682
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 9 deletions.
diff --git a/lib/ruby_lsp/requests/semantic_highlighting.rb b/lib/ruby_lsp/requests/semantic_highlighting.rb
@@ -101,7 +101,7 @@ def initialize(global_state, dispatcher, document, previous_result_id, range: ni
         @range = range
         @result_id = T.let(SemanticHighlighting.next_result_id.to_s, String)
         @response_builder = T.let(
-          ResponseBuilders::SemanticHighlighting.new(global_state.encoding),
+          ResponseBuilders::SemanticHighlighting.new(document.code_units_cache),
           ResponseBuilders::SemanticHighlighting,
         )
         Listeners::SemanticHighlighting.new(dispatcher, @response_builder)

diff --git a/lib/ruby_lsp/response_builders/semantic_highlighting.rb b/lib/ruby_lsp/response_builders/semantic_highlighting.rb
@@ -55,21 +55,27 @@ class UndefinedTokenType < StandardError; end
 
       ResponseType = type_member { { fixed: Interface::SemanticTokens } }
 
-      sig { params(encoding: Encoding).void }
-      def initialize(encoding)
+      sig do
+        params(code_units_cache: T.any(
+          T.proc.params(arg0: Integer).returns(Integer),
+          Prism::CodeUnitsCache,
+        )).void
+      end
+      def initialize(code_units_cache)
         super()
-        @encoding = encoding
+        @code_units_cache = code_units_cache
         @stack = T.let([], T::Array[SemanticToken])
       end
 
       sig { params(location: Prism::Location, type: Symbol, modifiers: T::Array[Symbol]).void }
       def add_token(location, type, modifiers = [])
-        length = location.end_code_units_offset(@encoding) - location.start_code_units_offset(@encoding)
+        end_code_unit = location.cached_end_code_units_offset(@code_units_cache)
+        length = end_code_unit - location.cached_start_code_units_offset(@code_units_cache)
         modifiers_indices = modifiers.filter_map { |modifier| TOKEN_MODIFIERS[modifier] }
         @stack.push(
           SemanticToken.new(
             start_line: location.start_line,
-            start_code_unit_column: location.start_code_units_column(@encoding),
+            start_code_unit_column: location.cached_start_code_units_column(@code_units_cache),
             length: length,
             type: T.must(TOKEN_TYPES[type]),
             modifier: modifiers_indices,
@@ -83,7 +89,7 @@ def last_token_matches?(location)
         return false unless token
 
         token.start_line == location.start_line &&
-          token.start_code_unit_column == location.start_code_units_column(@encoding)
+          token.start_code_unit_column == location.cached_start_code_units_column(@code_units_cache)
       end
 
       sig { returns(T.nilable(SemanticToken)) }

diff --git a/test/expectations/semantic_highlighting/multibyte_characters.exp.json b/test/expectations/semantic_highlighting/multibyte_characters.exp.json
@@ -16,15 +16,15 @@
     },
     {
       "delta_line": 3,
-      "delta_start_char": 11,
+      "delta_start_char": 10,
       "length": 2,
       "token_type": 8,
       "token_modifiers": 0
     },
     {
       "delta_line": 1,
       "delta_start_char": 11,
-      "length": 2,
+      "length": 1,
       "token_type": 8,
       "token_modifiers": 0
     }