Add find references support for constants (#2632)

jekyll/design-and-roadmap.markdown

@@ -195,6 +195,9 @@ Interested in contributing? Check out the issues tagged with [help-wanted] or [g
 - Explore using variable/method call names as a type hint
 - [Develop strategy to index declarations made in native extensions or C code. For example, Ruby’s own Core classes]
 - [Add find references support]
+    - [References method support](https://github.com/Shopify/ruby-lsp/issues/2640)
+    - [References instance variable support](https://github.com/Shopify/ruby-lsp/issues/2641)
+    - [References local variable support](https://github.com/Shopify/ruby-lsp/issues/2642)
 - [Add rename support]
 - [Add show type hierarchy support]
 - [Show index view on the VS Code extension allowing users to browse indexed gems]

jekyll/index.markdown

@@ -46,6 +46,7 @@ Want to discuss Ruby developer experience? Consider joining the public
     - [ERB support](#erb-support)
     - [Guessed types](#guessed-types)
     - [Rename symbol](#rename-symbol)
+    - [Find references](#find-references)
 - [VS Code only features](#vs-code-features)
     - [Dependencies view](#dependencies-view)
     - [Rails generator integrations](#rails-generator-integrations)
@@ -418,6 +419,13 @@ edits that will be applied by pressing CTRL/CMD + Enter after typing the desired
 
 ![Rename demo](images/rename.gif)
 
+### Find references
+
+The find references request allows users to both see a list of references or jump to reference locations. Note that
+only constants are currently supported, but support for methods, instance variables and local variables is planned.
+
+![References demo](images/references.gif)
+
 ## VS Code features
 
 The following features are all custom made for VS Code.

lib/ruby_indexer/lib/ruby_indexer/reference_finder.rb

@@ -14,26 +14,29 @@ class Reference
       sig { returns(Prism::Location) }
       attr_reader :location
 
-      sig { params(name: String, location: Prism::Location).void }
-      def initialize(name, location)
+      sig { returns(T::Boolean) }
+      attr_reader :declaration
+
+      sig { params(name: String, location: Prism::Location, declaration: T::Boolean).void }
+      def initialize(name, location, declaration:)
         @name = name
         @location = location
+        @declaration = declaration
       end
     end
 
-    sig { returns(T::Array[Reference]) }
-    attr_reader :references
-
     sig do
       params(
         fully_qualified_name: String,
         index: RubyIndexer::Index,
         dispatcher: Prism::Dispatcher,
+        include_declarations: T::Boolean,
       ).void
     end
-    def initialize(fully_qualified_name, index, dispatcher)
+    def initialize(fully_qualified_name, index, dispatcher, include_declarations: true)
       @fully_qualified_name = fully_qualified_name
       @index = index
+      @include_declarations = include_declarations
       @stack = T.let([], T::Array[String])
       @references = T.let([], T::Array[Reference])
 
@@ -62,14 +65,21 @@ def initialize(fully_qualified_name, index, dispatcher)
       )
     end
 
+    sig { returns(T::Array[Reference]) }
+    def references
+      return @references if @include_declarations
+
+      @references.reject(&:declaration)
+    end
+
     sig { params(node: Prism::ClassNode).void }
     def on_class_node_enter(node)
       constant_path = node.constant_path
       name = constant_path.slice
       nesting = actual_nesting(name)
 
       if nesting.join("::") == @fully_qualified_name
-        @references << Reference.new(name, constant_path.location)
+        @references << Reference.new(name, constant_path.location, declaration: true)
       end
 
       @stack << name
@@ -87,7 +97,7 @@ def on_module_node_enter(node)
       nesting = actual_nesting(name)
 
       if nesting.join("::") == @fully_qualified_name
-        @references << Reference.new(name, constant_path.location)
+        @references << Reference.new(name, constant_path.location, declaration: true)
       end
 
       @stack << name
@@ -236,10 +246,16 @@ def collect_constant_references(name, location)
       entries = @index.resolve(name, @stack)
       return unless entries
 
+      previous_reference = @references.last
+
       entries.each do |entry|
         next unless entry.name == @fully_qualified_name
 
-        @references << Reference.new(name, location)
+        # When processing a class/module declaration, we eagerly handle the constant reference. To avoid duplicates,
+        # when we find the constant node defining the namespace, then we have to check if it wasn't already added
+        next if previous_reference&.location == location
+
+        @references << Reference.new(name, location, declaration: false)
       end
     end
 

lib/ruby_indexer/test/reference_finder_test.rb

@@ -80,7 +80,7 @@ def find_references(fully_qualified_name, source)
       dispatcher = Prism::Dispatcher.new
       finder = ReferenceFinder.new(fully_qualified_name, index, dispatcher)
       dispatcher.visit(parse_result.value)
-      finder.references.uniq(&:location)
+      finder.references
     end
   end
 end

lib/ruby_lsp/internal.rb

@@ -75,10 +75,11 @@
 require "ruby_lsp/requests/inlay_hints"
 require "ruby_lsp/requests/on_type_formatting"
 require "ruby_lsp/requests/prepare_type_hierarchy"
+require "ruby_lsp/requests/references"
+require "ruby_lsp/requests/rename"
 require "ruby_lsp/requests/selection_ranges"
 require "ruby_lsp/requests/semantic_highlighting"
 require "ruby_lsp/requests/show_syntax_tree"
 require "ruby_lsp/requests/signature_help"
 require "ruby_lsp/requests/type_hierarchy_supertypes"
 require "ruby_lsp/requests/workspace_symbol"
-require "ruby_lsp/requests/rename"

lib/ruby_lsp/requests/references.rb

@@ -0,0 +1,110 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyLsp
+  module Requests
+    # The
+    # [references](https://microsoft.github.io/language-server-protocol/specification#textDocument_references)
+    # request finds all references for the selected symbol.
+    class References < Request
+      extend T::Sig
+      include Support::Common
+
+      sig do
+        params(
+          global_state: GlobalState,
+          store: Store,
+          document: T.any(RubyDocument, ERBDocument),
+          params: T::Hash[Symbol, T.untyped],
+        ).void
+      end
+      def initialize(global_state, store, document, params)
+        super()
+        @global_state = global_state
+        @store = store
+        @document = document
+        @params = params
+        @locations = T.let([], T::Array[Interface::Location])
+      end
+
+      sig { override.returns(T::Array[Interface::Location]) }
+      def perform
+        position = @params[:position]
+        char_position = @document.create_scanner.find_char_position(position)
+
+        node_context = RubyDocument.locate(
+          @document.parse_result.value,
+          char_position,
+          node_types: [Prism::ConstantReadNode, Prism::ConstantPathNode, Prism::ConstantPathTargetNode],
+        )
+        target = node_context.node
+        parent = node_context.parent
+        return @locations if !target || target.is_a?(Prism::ProgramNode)
+
+        if target.is_a?(Prism::ConstantReadNode) && parent.is_a?(Prism::ConstantPathNode)
+          target = determine_target(
+            target,
+            parent,
+            position,
+          )
+        end
+
+        target = T.cast(
+          target,
+          T.any(Prism::ConstantReadNode, Prism::ConstantPathNode, Prism::ConstantPathTargetNode),
+        )
+
+        name = constant_name(target)
+        return @locations unless name
+
+        entries = @global_state.index.resolve(name, node_context.nesting)
+        return @locations unless entries
+
+        fully_qualified_name = T.must(entries.first).name
+
+        Dir.glob(File.join(@global_state.workspace_path, "**/*.rb")).each do |path|
+          uri = URI::Generic.from_path(path: path)
+          # If the document is being managed by the client, then we should use whatever is present in the store instead
+          # of reading from disk
+          next if @store.key?(uri)
+
+          parse_result = Prism.parse_file(path)
+          collect_references(fully_qualified_name, parse_result, uri)
+        end
+
+        @store.each do |_uri, document|
+          collect_references(fully_qualified_name, document.parse_result, document.uri)
+        end
+
+        @locations
+      end
+
+      private
+
+      sig do
+        params(
+          fully_qualified_name: String,
+          parse_result: Prism::ParseResult,
+          uri: URI::Generic,
+        ).void
+      end
+      def collect_references(fully_qualified_name, parse_result, uri)
+        dispatcher = Prism::Dispatcher.new
+        finder = RubyIndexer::ReferenceFinder.new(
+          fully_qualified_name,
+          @global_state.index,
+          dispatcher,
+          include_declarations: @params.dig(:context, :includeDeclaration) || true,
+        )
+        dispatcher.visit(parse_result.value)
+
+        finder.references.each do |reference|
+          @locations << Interface::Location.new(
+            uri: uri.to_s,
+            range: range_from_location(reference.location),
+          )
+        end
+      end
+    end
+  end
+end

lib/ruby_lsp/requests/rename.rb

@@ -163,7 +163,7 @@ def collect_changes(fully_qualified_name, parse_result, name, uri)
         finder = RubyIndexer::ReferenceFinder.new(fully_qualified_name, @global_state.index, dispatcher)
         dispatcher.visit(parse_result.value)
 
-        finder.references.uniq(&:location).map do |reference|
+        finder.references.map do |reference|
           adjust_reference_for_edit(name, reference)
         end
       end

lib/ruby_lsp/server.rb

@@ -69,6 +69,8 @@ def process_message(message)
         text_document_prepare_type_hierarchy(message)
       when "textDocument/rename"
         text_document_rename(message)
+      when "textDocument/references"
+        text_document_references(message)
       when "typeHierarchy/supertypes"
         type_hierarchy_supertypes(message)
       when "typeHierarchy/subtypes"
@@ -230,6 +232,7 @@ def run_initialize(message)
           signature_help_provider: signature_help_provider,
           type_hierarchy_provider: type_hierarchy_provider,
           rename_provider: !@global_state.has_type_checker,
+          references_provider: !@global_state.has_type_checker,
           experimental: {
             addon_detection: true,
           },
@@ -636,6 +639,24 @@ def text_document_rename(message)
       send_message(Error.new(id: message[:id], code: Constant::ErrorCodes::REQUEST_FAILED, message: e.message))
     end
 
+    sig { params(message: T::Hash[Symbol, T.untyped]).void }
+    def text_document_references(message)
+      params = message[:params]
+      document = @store.get(params.dig(:textDocument, :uri))
+
+      unless document.is_a?(RubyDocument)
+        send_empty_response(message[:id])
+        return
+      end
+
+      send_message(
+        Result.new(
+          id: message[:id],
+          response: Requests::References.new(@global_state, @store, document, params).perform,
+        ),
+      )
+    end
+
     sig { params(document: Document[T.untyped]).returns(RubyDocument::SorbetLevel) }
     def sorbet_level(document)
       return RubyDocument::SorbetLevel::Ignore unless @global_state.has_type_checker

test/requests/references_test.rb

@@ -0,0 +1,33 @@
+# typed: true
+# frozen_string_literal: true
+
+require "test_helper"
+
+class ReferencesTest < Minitest::Test
+  def test_finds_constant_references
+    refs = find_references("test/fixtures/rename_me.rb", { line: 0, character: 6 }).map do |ref|
+      ref.range.start.line
+    end
+
+    assert_equal([0, 3], refs)
+  end
+
+  private
+
+  def find_references(fixture_path, position)
+    source = File.read(fixture_path)
+    path = File.expand_path(fixture_path)
+    global_state = RubyLsp::GlobalState.new
+    global_state.index.index_single(RubyIndexer::IndexablePath.new(nil, path), source)
+
+    store = RubyLsp::Store.new
+    document = RubyLsp::RubyDocument.new(source: source, version: 1, uri: URI::Generic.from_path(path: path))
+
+    RubyLsp::Requests::References.new(
+      global_state,
+      store,
+      document,
+      { position: position },
+    ).perform
+  end
+end