Find references from constant definitions

lib/ruby_indexer/lib/ruby_indexer/reference_finder.rb

@@ -85,7 +85,6 @@ def initialize(target, index, dispatcher, uri, include_declarations: true)
         :on_constant_path_or_write_node_enter,
         :on_constant_path_operator_write_node_enter,
         :on_constant_path_and_write_node_enter,
-        :on_constant_or_write_node_enter,
         :on_constant_path_node_enter,
         :on_constant_read_node_enter,
         :on_constant_write_node_enter,

lib/ruby_indexer/test/reference_finder_test.rb

@@ -70,6 +70,18 @@ class << self
       assert_equal(8, refs[2].location.start_line)
     end
 
+    def test_finds_constant_or_write_references_once
+      refs = find_const_references("Foo::BAR", <<~RUBY)
+        module Foo
+          BAR ||= 1
+        end
+      RUBY
+
+      assert_equal(1, refs.length)
+      assert_equal("BAR", refs[0].name)
+      assert_equal(2, refs[0].location.start_line)
+    end
+
     def test_finds_method_references
       refs = find_method_references("foo", <<~RUBY)
         class Bar

lib/ruby_lsp/requests/references.rb

@@ -32,6 +32,11 @@ def perform
             Prism::ConstantReadNode,
             Prism::ConstantPathNode,
             Prism::ConstantPathTargetNode,
+            Prism::ConstantAndWriteNode,
+            Prism::ConstantOperatorWriteNode,
+            Prism::ConstantOrWriteNode,
+            Prism::ConstantTargetNode,
+            Prism::ConstantWriteNode,
             Prism::InstanceVariableAndWriteNode,
             Prism::InstanceVariableOperatorWriteNode,
             Prism::InstanceVariableOrWriteNode,
@@ -55,7 +60,7 @@ def perform
           )
         end
 
-        target = target #: as Prism::ConstantReadNode | Prism::ConstantPathNode | Prism::ConstantPathTargetNode | Prism::InstanceVariableAndWriteNode | Prism::InstanceVariableOperatorWriteNode | Prism::InstanceVariableOrWriteNode | Prism::InstanceVariableReadNode | Prism::InstanceVariableTargetNode | Prism::InstanceVariableWriteNode | Prism::CallNode | Prism::DefNode,
+        target = target #: as Prism::ConstantReadNode | Prism::ConstantPathNode | Prism::ConstantPathTargetNode | Prism::ConstantAndWriteNode | Prism::ConstantOperatorWriteNode | Prism::ConstantOrWriteNode | Prism::ConstantTargetNode | Prism::ConstantWriteNode | Prism::InstanceVariableAndWriteNode | Prism::InstanceVariableOperatorWriteNode | Prism::InstanceVariableOrWriteNode | Prism::InstanceVariableReadNode | Prism::InstanceVariableTargetNode | Prism::InstanceVariableWriteNode | Prism::CallNode | Prism::DefNode,
 
         reference_target = create_reference_target(target, node_context)
         return @locations unless reference_target
@@ -81,19 +86,21 @@ def perform
 
       private
 
-      #: ((Prism::ConstantReadNode | Prism::ConstantPathNode | Prism::ConstantPathTargetNode | Prism::InstanceVariableAndWriteNode | Prism::InstanceVariableOperatorWriteNode | Prism::InstanceVariableOrWriteNode | Prism::InstanceVariableReadNode | Prism::InstanceVariableTargetNode | Prism::InstanceVariableWriteNode | Prism::CallNode | Prism::DefNode) target_node, NodeContext node_context) -> RubyIndexer::ReferenceFinder::Target?
+      #: ((Prism::ConstantReadNode | Prism::ConstantPathNode | Prism::ConstantPathTargetNode | Prism::ConstantAndWriteNode | Prism::ConstantOperatorWriteNode | Prism::ConstantOrWriteNode | Prism::ConstantTargetNode | Prism::ConstantWriteNode | Prism::InstanceVariableAndWriteNode | Prism::InstanceVariableOperatorWriteNode | Prism::InstanceVariableOrWriteNode | Prism::InstanceVariableReadNode | Prism::InstanceVariableTargetNode | Prism::InstanceVariableWriteNode | Prism::CallNode | Prism::DefNode) target_node, NodeContext node_context) -> RubyIndexer::ReferenceFinder::Target?
       def create_reference_target(target_node, node_context)
         case target_node
         when Prism::ConstantReadNode, Prism::ConstantPathNode, Prism::ConstantPathTargetNode
           name = RubyIndexer::Index.constant_name(target_node)
           return unless name
 
-          entries = @global_state.index.resolve(name, node_context.nesting)
-          return unless entries
-
-          fully_qualified_name = entries.first #: as !nil
-            .name
-          RubyIndexer::ReferenceFinder::ConstTarget.new(fully_qualified_name)
+          create_constant_reference_target(name, node_context)
+        when
+          Prism::ConstantAndWriteNode,
+          Prism::ConstantOperatorWriteNode,
+          Prism::ConstantOrWriteNode,
+          Prism::ConstantTargetNode,
+          Prism::ConstantWriteNode
+          create_constant_reference_target(target_node.name.to_s, node_context)
         when
           Prism::InstanceVariableAndWriteNode,
           Prism::InstanceVariableOperatorWriteNode,
@@ -111,6 +118,16 @@ def create_reference_target(target_node, node_context)
         end
       end
 
+      #: (String name, NodeContext node_context) -> RubyIndexer::ReferenceFinder::ConstTarget?
+      def create_constant_reference_target(name, node_context)
+        entries = @global_state.index.resolve(name, node_context.nesting)
+        return unless entries
+
+        fully_qualified_name = entries.first #: as !nil
+          .name
+        RubyIndexer::ReferenceFinder::ConstTarget.new(fully_qualified_name)
+      end
+
       #: (RubyIndexer::ReferenceFinder::Target target, Prism::LexResult parse_result, URI::Generic uri) -> void
       def collect_references(target, parse_result, uri)
         dispatcher = Prism::Dispatcher.new

test/requests/references_test.rb

@@ -12,24 +12,93 @@ def test_finds_constant_references
     assert_equal([0, 3], refs)
   end
 
+  def test_finds_constant_references_from_constant_write
+    source = <<~RUBY
+      module ReferencesTestNamespace
+        class Klass
+          CONSTANT = 1
+          CONSTANT
+        end
+      end
+
+      ReferencesTestNamespace::Klass::CONSTANT
+    RUBY
+
+    refs = find_references_in_source(source, { line: 2, character: 4 }).map do |ref|
+      ref.range.start.line
+    end
+
+    assert_equal([2, 3, 7], refs)
+  end
+
+  def test_finds_constant_references_from_constant_or_write
+    source = <<~RUBY
+      module ReferencesTestNamespace
+        CONSTANT ||= 1
+      end
+
+      ReferencesTestNamespace::CONSTANT
+    RUBY
+
+    refs = find_references_in_source(source, { line: 1, character: 2 }).map do |ref|
+      ref.range.start.line
+    end
+
+    assert_equal([1, 4], refs)
+  end
+
+  def test_finds_constant_references_from_constant_operator_write
+    source = <<~RUBY
+      module ReferencesTestNamespace
+        CONSTANT = 1
+        CONSTANT += 1
+      end
+    RUBY
+
+    refs = find_references_in_source(source, { line: 2, character: 2 }).map do |ref|
+      ref.range.start.line
+    end
+
+    assert_equal([1, 2], refs)
+  end
+
+  def test_finds_constant_references_from_constant_target
+    source = <<~RUBY
+      module ReferencesTestNamespace
+        FIRST, SECOND = 1, 2
+      end
+
+      ReferencesTestNamespace::SECOND
+    RUBY
+
+    refs = find_references_in_source(source, { line: 1, character: 9 }).map do |ref|
+      [ref.range.start.line, ref.range.start.character]
+    end
+
+    assert_equal([[1, 9], [4, 0]], refs)
+  end
+
   private
 
   def find_references(fixture_path, position)
-    source = File.read(fixture_path)
     path = File.expand_path(fixture_path)
+    find_references_in_source(File.read(fixture_path), position, URI::Generic.from_path(path: path))
+  end
+
+  def find_references_in_source(source, position, uri = URI::Generic.from_path(path: "/fake.rb"))
     global_state = RubyLsp::GlobalState.new
-    global_state.index.index_single(URI::Generic.from_path(path: path), source)
+    global_state.index.index_single(uri, source)
 
     store = RubyLsp::Store.new(global_state)
     document = RubyLsp::RubyDocument.new(
       source: source,
       version: 1,
-      uri: URI::Generic.from_path(path: path),
+      uri: uri,
       global_state: global_state,
     )
 
     # In addition to glob files from the workspace, we also want to test references collection from the store
-    store.set(uri: URI::Generic.from_path(path: path), source: source, version: 1, language_id: :ruby)
+    store.set(uri: uri, source: source, version: 1, language_id: :ruby)
 
     RubyLsp::Requests::References.new(
       global_state,