Avoid collecting duplicates in reference finder

lib/ruby_indexer/lib/ruby_indexer/location.rb

@@ -44,5 +44,17 @@ def initialize(start_line, end_line, start_column, end_column)
       @start_column = start_column
       @end_column = end_column
     end
+
+    sig do
+      params(
+        another_location: T.any(Location, Prism::Location)
+      ).returns(T::Boolean)
+    end
+    def ==(another_location)
+      start_line == another_location.start_line &&
+        end_line == another_location.end_line &&
+        start_column == another_location.start_column &&
+        end_column == another_location.end_column
+    end
   end
 end

lib/ruby_indexer/lib/ruby_indexer/reference_finder.rb

@@ -75,12 +75,14 @@ def initialize(name, location, declaration:)
         target: Target,
         index: RubyIndexer::Index,
         dispatcher: Prism::Dispatcher,
+        uri: URI::Generic,
         include_declarations: T::Boolean,
       ).void
     end
-    def initialize(target, index, dispatcher, include_declarations: true)
+    def initialize(target, index, dispatcher, uri, include_declarations: true)
       @target = target
       @index = index
+      @uri = uri
       @include_declarations = include_declarations
       @stack = T.let([], T::Array[String])
       @references = T.let([], T::Array[Reference])
@@ -126,15 +128,7 @@ def references
 
     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 @target.is_a?(ConstTarget) && nesting.join("::") == @target.fully_qualified_name
-        @references << Reference.new(name, constant_path.location, declaration: true)
-      end
-
-      @stack << name
+      @stack << node.constant_path.slice
     end
 
     sig { params(node: Prism::ClassNode).void }
@@ -144,15 +138,7 @@ def on_class_node_leave(node)
 
     sig { params(node: Prism::ModuleNode).void }
     def on_module_node_enter(node)
-      constant_path = node.constant_path
-      name = constant_path.slice
-      nesting = actual_nesting(name)
-
-      if @target.is_a?(ConstTarget) && nesting.join("::") == @target.fully_qualified_name
-        @references << Reference.new(name, constant_path.location, declaration: true)
-      end
-
-      @stack << name
+      @stack << node.constant_path.slice
     end
 
     sig { params(node: Prism::ModuleNode).void }
@@ -341,17 +327,27 @@ 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 == @target.fully_qualified_name
+      # Filter down to all constant declarations that match the expected name and type
+      matching_entries = entries.select do |e|
+        [
+          Entry::Namespace,
+          Entry::Constant,
+          Entry::ConstantAlias,
+          Entry::UnresolvedConstantAlias,
+          Entry::Class
+        ].include?(e.class) &&
+          e.name == @target.fully_qualified_name
+      end
 
-        # 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
+      return if matching_entries.empty?
 
-        @references << Reference.new(name, location, declaration: false)
+      # If any of the matching entries have the same location as the constant and were
+      # defined in the same file, then it is that constant's declaration
+      declaration = matching_entries.any? do |e|
+        e.name_location == location && e.uri == @uri
       end
+
+      @references << Reference.new(name, location, declaration: declaration)
     end
 
     sig { params(name: String, location: Prism::Location, declaration: T::Boolean).void }

lib/ruby_indexer/test/reference_finder_test.rb

@@ -274,6 +274,30 @@ def name
       assert_equal(8, refs[1].location.start_line)
     end
 
+    def test_accounts_for_reopened_classes
+      refs = find_const_references("Foo", <<~RUBY)
+        class Foo
+        end
+        class Foo
+          class Bar
+          end
+        end
+
+        Foo.new
+      RUBY
+
+      assert_equal(3, refs.size)
+
+      assert_equal("Foo", refs[0].name)
+      assert_equal(1, refs[0].location.start_line)
+
+      assert_equal("Foo", refs[1].name)
+      assert_equal(3, refs[1].location.start_line)
+
+      assert_equal("Foo", refs[2].name)
+      assert_equal(8, refs[2].location.start_line)
+    end
+
     private
 
     def find_const_references(const_name, source)
@@ -293,11 +317,12 @@ def find_instance_variable_references(instance_variable_name, source)
 
     def find_references(target, source)
       file_path = "/fake.rb"
+      uri = URI::Generic.from_path(path: file_path)
       index = Index.new
-      index.index_single(URI::Generic.from_path(path: file_path), source)
+      index.index_single(uri, source)
       parse_result = Prism.parse(source)
       dispatcher = Prism::Dispatcher.new
-      finder = ReferenceFinder.new(target, index, dispatcher)
+      finder = ReferenceFinder.new(target, index, dispatcher, uri)
       dispatcher.visit(parse_result.value)
       finder.references
     end

lib/ruby_lsp/requests/references.rb

@@ -158,6 +158,7 @@ def collect_references(target, parse_result, uri)
           target,
           @global_state.index,
           dispatcher,
+          uri,
           include_declarations: @params.dig(:context, :includeDeclaration) || true,
         )
         dispatcher.visit(parse_result.value)

lib/ruby_lsp/requests/rename.rb

@@ -179,7 +179,7 @@ def collect_text_edits(target, name)
       end
       def collect_changes(target, parse_result, name, uri)
         dispatcher = Prism::Dispatcher.new
-        finder = RubyIndexer::ReferenceFinder.new(target, @global_state.index, dispatcher)
+        finder = RubyIndexer::ReferenceFinder.new(target, @global_state.index, dispatcher, uri)
         dispatcher.visit(parse_result.value)
 
         finder.references.map do |reference|