Support semantic token deltas

lib/ruby_lsp/document.rb

@@ -16,6 +16,7 @@ class LanguageId < T::Enum
     extend T::Generic
 
     ParseResultType = type_member
+    EMPTY_CACHE = T.let(Object.new.freeze, Object)
 
     abstract!
 
@@ -34,9 +35,13 @@ class LanguageId < T::Enum
     sig { returns(Encoding) }
     attr_reader :encoding
 
+    sig { returns(T.any(Interface::SemanticTokens, Object)) }
+    attr_accessor :semantic_tokens
+
     sig { params(source: String, version: Integer, uri: URI::Generic, encoding: Encoding).void }
     def initialize(source:, version:, uri:, encoding: Encoding::UTF_8)
-      @cache = T.let({}, T::Hash[String, T.untyped])
+      @cache = T.let(Hash.new(EMPTY_CACHE), T::Hash[String, T.untyped])
+      @semantic_tokens = T.let(EMPTY_CACHE, T.any(Interface::SemanticTokens, Object))
       @encoding = T.let(encoding, Encoding)
       @source = T.let(source, String)
       @version = T.let(version, Integer)
@@ -63,7 +68,7 @@ def language_id; end
     end
     def cache_fetch(request_name, &block)
       cached = @cache[request_name]
-      return cached if cached
+      return cached if cached != EMPTY_CACHE
 
       result = block.call(self)
       @cache[request_name] = result

lib/ruby_lsp/listeners/inlay_hints.rb

@@ -69,6 +69,17 @@ def on_implicit_node_enter(node)
           tooltip: tooltip,
         )
       end
+
+      private
+
+      sig { params(node: T.nilable(Prism::Node), range: T.nilable(T::Range[Integer])).returns(T::Boolean) }
+      def visible?(node, range)
+        return true if range.nil?
+        return false if node.nil?
+
+        loc = node.location
+        range.cover?(loc.start_line - 1) && range.cover?(loc.end_line - 1)
+      end
     end
   end
 end

lib/ruby_lsp/listeners/semantic_highlighting.rb

@@ -22,12 +22,10 @@ class SemanticHighlighting
         params(
           dispatcher: Prism::Dispatcher,
           response_builder: ResponseBuilders::SemanticHighlighting,
-          range: T.nilable(T::Range[Integer]),
         ).void
       end
-      def initialize(dispatcher, response_builder, range: nil)
+      def initialize(dispatcher, response_builder)
         @response_builder = response_builder
-        @range = range
         @special_methods = T.let(nil, T.nilable(T::Array[String]))
         @current_scope = T.let(ParameterScope.new, ParameterScope)
         @inside_regex_capture = T.let(false, T::Boolean)
@@ -74,7 +72,6 @@ def initialize(dispatcher, response_builder, range: nil)
       sig { params(node: Prism::CallNode).void }
       def on_call_node_enter(node)
         return if @inside_implicit_node
-        return unless visible?(node, @range)
 
         message = node.message
         return unless message
@@ -107,51 +104,38 @@ def on_match_write_node_leave(node)
       sig { params(node: Prism::ConstantReadNode).void }
       def on_constant_read_node_enter(node)
         return if @inside_implicit_node
-        return unless visible?(node, @range)
 
         @response_builder.add_token(node.location, :namespace)
       end
 
       sig { params(node: Prism::ConstantWriteNode).void }
       def on_constant_write_node_enter(node)
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.name_loc, :namespace)
       end
 
       sig { params(node: Prism::ConstantAndWriteNode).void }
       def on_constant_and_write_node_enter(node)
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.name_loc, :namespace)
       end
 
       sig { params(node: Prism::ConstantOperatorWriteNode).void }
       def on_constant_operator_write_node_enter(node)
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.name_loc, :namespace)
       end
 
       sig { params(node: Prism::ConstantOrWriteNode).void }
       def on_constant_or_write_node_enter(node)
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.name_loc, :namespace)
       end
 
       sig { params(node: Prism::ConstantTargetNode).void }
       def on_constant_target_node_enter(node)
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.location, :namespace)
       end
 
       sig { params(node: Prism::DefNode).void }
       def on_def_node_enter(node)
         @current_scope = ParameterScope.new(@current_scope)
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.name_loc, :method, [:declaration])
       end
 
@@ -184,7 +168,6 @@ def on_block_parameter_node_enter(node)
       sig { params(node: Prism::RequiredKeywordParameterNode).void }
       def on_required_keyword_parameter_node_enter(node)
         @current_scope << node.name
-        return unless visible?(node, @range)
 
         location = node.name_loc
         @response_builder.add_token(location.copy(length: location.length - 1), :parameter)
@@ -193,7 +176,6 @@ def on_required_keyword_parameter_node_enter(node)
       sig { params(node: Prism::OptionalKeywordParameterNode).void }
       def on_optional_keyword_parameter_node_enter(node)
         @current_scope << node.name
-        return unless visible?(node, @range)
 
         location = node.name_loc
         @response_builder.add_token(location.copy(length: location.length - 1), :parameter)
@@ -205,24 +187,19 @@ def on_keyword_rest_parameter_node_enter(node)
 
         if name
           @current_scope << name.to_sym
-
-          @response_builder.add_token(T.must(node.name_loc), :parameter) if visible?(node, @range)
+          @response_builder.add_token(T.must(node.name_loc), :parameter)
         end
       end
 
       sig { params(node: Prism::OptionalParameterNode).void }
       def on_optional_parameter_node_enter(node)
         @current_scope << node.name
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.name_loc, :parameter)
       end
 
       sig { params(node: Prism::RequiredParameterNode).void }
       def on_required_parameter_node_enter(node)
         @current_scope << node.name
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.location, :parameter)
       end
 
@@ -232,29 +209,23 @@ def on_rest_parameter_node_enter(node)
 
         if name
           @current_scope << name.to_sym
-
-          @response_builder.add_token(T.must(node.name_loc), :parameter) if visible?(node, @range)
+          @response_builder.add_token(T.must(node.name_loc), :parameter)
         end
       end
 
       sig { params(node: Prism::SelfNode).void }
       def on_self_node_enter(node)
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.location, :variable, [:default_library])
       end
 
       sig { params(node: Prism::LocalVariableWriteNode).void }
       def on_local_variable_write_node_enter(node)
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.name_loc, @current_scope.type_for(node.name))
       end
 
       sig { params(node: Prism::LocalVariableReadNode).void }
       def on_local_variable_read_node_enter(node)
         return if @inside_implicit_node
-        return unless visible?(node, @range)
 
         # Numbered parameters
         if /_\d+/.match?(node.name)
@@ -267,22 +238,16 @@ def on_local_variable_read_node_enter(node)
 
       sig { params(node: Prism::LocalVariableAndWriteNode).void }
       def on_local_variable_and_write_node_enter(node)
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.name_loc, @current_scope.type_for(node.name))
       end
 
       sig { params(node: Prism::LocalVariableOperatorWriteNode).void }
       def on_local_variable_operator_write_node_enter(node)
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.name_loc, @current_scope.type_for(node.name))
       end
 
       sig { params(node: Prism::LocalVariableOrWriteNode).void }
       def on_local_variable_or_write_node_enter(node)
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.name_loc, @current_scope.type_for(node.name))
       end
 
@@ -294,15 +259,11 @@ def on_local_variable_target_node_enter(node)
         # prevent pushing local variable target tokens. See https://github.com/ruby/prism/issues/1912
         return if @inside_regex_capture
 
-        return unless visible?(node, @range)
-
         @response_builder.add_token(node.location, @current_scope.type_for(node.name))
       end
 
       sig { params(node: Prism::ClassNode).void }
       def on_class_node_enter(node)
-        return unless visible?(node, @range)
-
         constant_path = node.constant_path
 
         if constant_path.is_a?(Prism::ConstantReadNode)
@@ -342,8 +303,6 @@ def on_class_node_enter(node)
 
       sig { params(node: Prism::ModuleNode).void }
       def on_module_node_enter(node)
-        return unless visible?(node, @range)
-
         constant_path = node.constant_path
 
         if constant_path.is_a?(Prism::ConstantReadNode)
@@ -365,8 +324,6 @@ def on_module_node_enter(node)
 
       sig { params(node: Prism::ImplicitNode).void }
       def on_implicit_node_enter(node)
-        return unless visible?(node, @range)
-
         @inside_implicit_node = true
       end
 
@@ -378,7 +335,6 @@ def on_implicit_node_leave(node)
       sig { params(node: Prism::ConstantPathNode).void }
       def on_constant_path_node_enter(node)
         return if @inside_implicit_node
-        return unless visible?(node, @range)
 
         @response_builder.add_token(node.name_loc, :namespace)
       end

lib/ruby_lsp/requests/semantic_highlighting.rb

@@ -35,28 +35,125 @@ def provider
               token_modifiers: ResponseBuilders::SemanticHighlighting::TOKEN_MODIFIERS.keys,
             ),
             range: true,
-            full: { delta: false },
+            full: { delta: true },
           )
         end
+
+        # The compute_delta method receives the current semantic tokens and the previous semantic tokens and then tries
+        # to compute the smallest possible semantic token edit that will turn previous into current
+        sig do
+          params(
+            current_tokens: T::Array[Integer],
+            previous_tokens: T::Array[Integer],
+            result_id: String,
+          ).returns(Interface::SemanticTokensDelta)
+        end
+        def compute_delta(current_tokens, previous_tokens, result_id)
+          # Find the index of the first token that is different between the two sets of tokens
+          first_different_position = current_tokens.zip(previous_tokens).find_index { |new, old| new != old }
+
+          # When deleting a token from the end, the first_different_position will be nil, but since we're removing at
+          # the end, then we have to initialize it to the length of the current tokens after the deletion
+          if !first_different_position && current_tokens.length < previous_tokens.length
+            first_different_position = current_tokens.length
+          end
+
+          unless first_different_position
+            return Interface::SemanticTokensDelta.new(result_id: result_id, edits: [])
+          end
+
+          # Filter the tokens based on the first different position. This must happen at this stage, before we try to
+          # find the next position from the end or else we risk confusing sets of token that may have different lengths,
+          # but end with the exact same token
+          old_tokens = T.must(previous_tokens[first_different_position...])
+          new_tokens = T.must(current_tokens[first_different_position...])
+
+          # Then search from the end to find the first token that doesn't match. Since the user is normally editing the
+          # middle of the file, this will minimize the number of edits since the end of the token array has not changed
+          first_different_token_from_end = new_tokens.reverse.zip(old_tokens.reverse).find_index do |new, old|
+            new != old
+          end || 0
+
+          # Filter the old and new tokens to only the section that will be replaced/inserted/deleted
+          old_tokens = T.must(old_tokens[...old_tokens.length - first_different_token_from_end])
+          new_tokens = T.must(new_tokens[...new_tokens.length - first_different_token_from_end])
+
+          # And we send back a single edit, replacing an entire section for the new tokens
+          Interface::SemanticTokensDelta.new(
+            result_id: result_id,
+            edits: [{ start: first_different_position, deleteCount: old_tokens.length, data: new_tokens }],
+          )
+        end
+
+        sig { returns(Integer) }
+        def next_result_id
+          @mutex.synchronize do
+            @result_id += 1
+          end
+        end
       end
 
-      sig { params(global_state: GlobalState, dispatcher: Prism::Dispatcher, range: T.nilable(T::Range[Integer])).void }
-      def initialize(global_state, dispatcher, range: nil)
+      @result_id = T.let(0, Integer)
+      @mutex = T.let(Mutex.new, Mutex)
+
+      sig do
+        params(
+          global_state: GlobalState,
+          dispatcher: Prism::Dispatcher,
+          document: T.any(RubyDocument, ERBDocument),
+          previous_result_id: T.nilable(String),
+          range: T.nilable(T::Range[Integer]),
+        ).void
+      end
+      def initialize(global_state, dispatcher, document, previous_result_id, range: nil)
         super()
+
+        @document = document
+        @previous_result_id = previous_result_id
+        @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,
         )
-        Listeners::SemanticHighlighting.new(dispatcher, @response_builder, range: range)
+        Listeners::SemanticHighlighting.new(dispatcher, @response_builder)
 
         Addon.addons.each do |addon|
           addon.create_semantic_highlighting_listener(@response_builder, dispatcher)
         end
       end
 
-      sig { override.returns(Interface::SemanticTokens) }
+      sig { override.returns(T.any(Interface::SemanticTokens, Interface::SemanticTokensDelta)) }
       def perform
-        @response_builder.response
+        previous_tokens = @document.semantic_tokens
+        tokens = @response_builder.response
+        encoded_tokens = ResponseBuilders::SemanticHighlighting::SemanticTokenEncoder.new.encode(tokens)
+        full_response = Interface::SemanticTokens.new(result_id: @result_id, data: encoded_tokens)
+        @document.semantic_tokens = full_response
+
+        if @range
+          tokens_within_range = tokens.select { |token| @range.cover?(token.start_line - 1) }
+
+          return Interface::SemanticTokens.new(
+            result_id: @result_id,
+            data: ResponseBuilders::SemanticHighlighting::SemanticTokenEncoder.new.encode(tokens_within_range),
+          )
+        end
+
+        # Semantic tokens full delta
+        if @previous_result_id
+          response = if previous_tokens.is_a?(Interface::SemanticTokens) &&
+              previous_tokens.result_id == @previous_result_id
+            Requests::SemanticHighlighting.compute_delta(encoded_tokens, previous_tokens.data, @result_id)
+          else
+            full_response
+          end
+
+          return response
+        end
+
+        # Semantic tokens full
+        full_response
       end
     end
   end