Implement basic semantic highlighting (method calls, local vars)

lib/ruby_lsp/cli.rb

@@ -46,6 +46,10 @@ def self.start(_argv)
           respond_with_folding_ranges(request.dig(:params, :textDocument, :uri))
         end
 
+        on("textDocument/semanticTokens/full") do |request|
+          respond_with_semantic_highlighting(request.dig(:params, :textDocument, :uri))
+        end
+
         on("shutdown") { shutdown }
       end
 

lib/ruby_lsp/handler.rb

@@ -59,7 +59,18 @@ def respond_with_capabilities
           ),
           folding_range_provider: Interface::FoldingRangeClientCapabilities.new(
             line_folding_only: true
-          )
+          ),
+          semantic_tokens_provider: Interface::SemanticTokensRegistrationOptions.new(
+            document_selector: { scheme: "file", language: "ruby" },
+            legend: Interface::SemanticTokensLegend.new(
+              token_types: Requests::SemanticHighlighting::TOKEN_TYPES,
+              token_modifiers: Requests::SemanticHighlighting::TOKEN_MODIFIERS
+            ),
+            range: false,
+            full: {
+              delta: true,
+            }
+          ),
         )
       )
     end
@@ -71,5 +82,9 @@ def respond_with_document_symbol(uri)
     def respond_with_folding_ranges(uri)
       Requests::FoldingRanges.run(store[uri])
     end
+
+    def respond_with_semantic_highlighting(uri)
+      Requests::SemanticHighlighting.run(store[uri])
+    end
   end
 end

lib/ruby_lsp/requests.rb

@@ -4,5 +4,6 @@ module RubyLsp
   module Requests
     autoload :DocumentSymbol, "ruby_lsp/requests/document_symbol"
     autoload :FoldingRanges, "ruby_lsp/requests/folding_ranges"
+    autoload :SemanticHighlighting, "ruby_lsp/requests/semantic_highlighting"
   end
 end

lib/ruby_lsp/requests/semantic_highlighting.rb

@@ -0,0 +1,114 @@
+# frozen_string_literal: true
+
+module RubyLsp
+  module Requests
+    class SemanticHighlighting < Visitor
+      TOKEN_TYPES = [
+        :local_variable,
+        :method_call,
+      ].freeze
+      TOKEN_MODIFIERS = [].freeze
+
+      def self.run(parsed_tree)
+        new(parsed_tree).run
+      end
+
+      def initialize(parsed_tree)
+        @tokens = []
+        @parser = parsed_tree.parser
+        @tree = parsed_tree.tree
+        @current_row = 0
+        @current_column = 0
+
+        super()
+      end
+
+      def run
+        visit(@tree)
+        LanguageServer::Protocol::Interface::SemanticTokens.new(data: @tokens)
+      end
+
+      private
+
+      def visit_assign(node)
+        super
+      end
+
+      def visit_m_assign(node)
+        node.target.parts.each do |var_ref|
+          add_token(var_ref.value.location, :local_variable)
+        end
+      end
+
+      def visit_var_field(node)
+        add_token(node.value.location, :local_variable)
+      end
+
+      def visit_var_ref(node)
+        add_token(node.value.location, :local_variable)
+      end
+
+      def visit_a_ref_field(node)
+        add_token(node.target.collection.value.location, :local_variable)
+      end
+
+      def visit_call(node)
+        super
+        add_token(node.message.location, :method_call)
+      end
+
+      def visit_command(node)
+        add_token(node.message.location, :method_call)
+      end
+
+      def visit_f_call(node)
+        add_token(node.value.location, :method_call)
+      end
+
+      def visit_v_call(node)
+        add_token(node.value.location, :method_call)
+      end
+
+      def add_token(location, classification)
+        length = location.end_char - location.start_char
+
+        compute_delta(location) do |delta_line, delta_column|
+          @tokens.push(delta_line, delta_column, length, TOKEN_TYPES.index(classification), 0)
+        end
+      end
+
+      # The delta array is computed according to the LSP specification:
+      # > The protocol for the token format relative uses relative
+      # > positions, because most tokens remain stable relative to
+      # > each other when edits are made in a file. This simplifies
+      # > the computation of a delta if a server supports it. So each
+      # > token is represented using 5 integers.
+
+      # For more information on how each number is calculated, read:
+      # https://microsoft.github.io/language-server-protocol/specifications/specification-current/#textDocument_semanticTokens
+      def compute_delta(location)
+        row = location.start_line - 1
+
+        line = @parser.line_counts[location.start_line - 1]
+        column = location.start_char - line.start
+
+        if row < @current_row
+          raise InvalidTokenRowError, "Invalid token row detected: " \
+            "Ensure tokens are added in the expected order."
+        end
+
+        delta_line = row - @current_row
+
+        delta_column = column
+        delta_column -= @current_column if delta_line == 0
+
+        yield delta_line, delta_column
+
+        @current_row = row
+        @current_column = column
+      end
+
+      class InvalidTokenRowError < StandardError; end
+    end
+  end
+end

test/requests/semantic_highlighting_test.rb

@@ -0,0 +1,125 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class SemanticHighlightingTest < Minitest::Test
+  def test_local_variables
+    tokens = [
+      { delta_line: 1, delta_start_char: 2, length: 3, token_type: 0, token_modifiers: 0 },
+      { delta_line: 1, delta_start_char: 2, length: 3, token_type: 0, token_modifiers: 0 },
+    ]
+
+    assert_tokens(tokens, <<~RUBY)
+      def my_method
+        var = 1
+        var
+      end
+    RUBY
+  end
+
+  def test_multi_assignment
+    tokens = [
+      { delta_line: 1, delta_start_char: 2, length: 1, token_type: 0, token_modifiers: 0 },
+      { delta_line: 0, delta_start_char: 3, length: 1, token_type: 0, token_modifiers: 0 },
+      { delta_line: 1, delta_start_char: 2, length: 1, token_type: 0, token_modifiers: 0 },
+      { delta_line: 1, delta_start_char: 2, length: 1, token_type: 0, token_modifiers: 0 },
+    ]
+
+    assert_tokens(tokens, <<~RUBY)
+      def my_method
+        a, b = [1, 2]
+        a
+        b
+      end
+    RUBY
+  end
+
+  def test_aref_variable
+    tokens = [
+      { delta_line: 1, delta_start_char: 2, length: 1, token_type: 0, token_modifiers: 0 },
+      { delta_line: 1, delta_start_char: 2, length: 1, token_type: 0, token_modifiers: 0 },
+    ]
+
+    assert_tokens(tokens, <<~RUBY)
+      def my_method
+        a = []
+        a[1]
+      end
+    RUBY
+  end
+
+  def test_command_invocation
+    tokens = [
+      { delta_line: 0, delta_start_char: 0, length: 4, token_type: 1, token_modifiers: 0 },
+    ]
+
+    assert_tokens(tokens, <<~RUBY)
+      puts "Hello"
+    RUBY
+  end
+
+  def test_call_invocation
+    tokens = [
+      { delta_line: 0, delta_start_char: 8, length: 6, token_type: 1, token_modifiers: 0 },
+    ]
+
+    assert_tokens(tokens, <<~RUBY)
+      "Hello".upcase
+    RUBY
+  end
+
+  def test_variable_receiver_in_call_invocation
+    tokens = [
+      { delta_line: 1, delta_start_char: 2, length: 3, token_type: 0, token_modifiers: 0 },
+      { delta_line: 1, delta_start_char: 2, length: 3, token_type: 0, token_modifiers: 0 },
+      { delta_line: 0, delta_start_char: 4, length: 6, token_type: 1, token_modifiers: 0 },
+    ]
+
+    assert_tokens(tokens, <<~RUBY)
+      def some_method
+        var = "Hello"
+        var.upcase
+      end
+    RUBY
+  end
+
+  def test_vcall_invocation
+    tokens = [
+      { delta_line: 1, delta_start_char: 2, length: 10, token_type: 1, token_modifiers: 0 },
+    ]
+
+    assert_tokens(tokens, <<~RUBY)
+      def some_method
+        invocation
+      end
+    RUBY
+  end
+
+  def test_fcall_invocation
+    tokens = [
+      { delta_line: 1, delta_start_char: 2, length: 10, token_type: 1, token_modifiers: 0 },
+    ]
+
+    assert_tokens(tokens, <<~RUBY)
+      def some_method
+        invocation(1, 2, 3)
+      end
+    RUBY
+  end
+
+  private
+
+  def assert_tokens(expected, source_code)
+    parsed_tree = RubyLsp::Store::ParsedTree.new(source_code)
+    assert_equal(
+      inline_tokens(expected),
+      RubyLsp::Requests::SemanticHighlighting.run(parsed_tree).data
+    )
+  end
+
+  def inline_tokens(tokens)
+    tokens.flat_map do |token|
+      [token[:delta_line], token[:delta_start_char], token[:length], token[:token_type], token[:token_modifiers]]
+    end
+  end
+end