Fix completion for inherited constants (#2586)

* Add constant completion method in index

* Fix alias following

* Fix completion for inherited constants

lib/ruby_indexer/lib/ruby_indexer/index.rb

@@ -242,6 +242,64 @@ def method_completion_candidates(name, receiver_name)
       completion_items.values.map!(&:first)
     end
 
+    sig do
+      params(
+        name: String,
+        nesting: T::Array[String],
+      ).returns(T::Array[T::Array[T.any(
+        Entry::Constant,
+        Entry::ConstantAlias,
+        Entry::Namespace,
+        Entry::UnresolvedConstantAlias,
+      )]])
+    end
+    def constant_completion_candidates(name, nesting)
+      # If we have a top level reference, then we don't need to include completions inside the current nesting
+      if name.start_with?("::")
+        return T.cast(
+          @entries_tree.search(name.delete_prefix("::")),
+          T::Array[T::Array[T.any(
+            Entry::Constant,
+            Entry::ConstantAlias,
+            Entry::Namespace,
+            Entry::UnresolvedConstantAlias,
+          )]],
+        )
+      end
+
+      # Otherwise, we have to include every possible constant the user might be referring to. This is essentially the
+      # same algorithm as resolve, but instead of returning early we concatenate all unique results
+
+      # Direct constants inside this namespace
+      entries = @entries_tree.search(nesting.any? ? "#{nesting.join("::")}::#{name}" : name)
+
+      # Constants defined in enclosing scopes
+      nesting.length.downto(1) do |i|
+        namespace = T.must(nesting[0...i]).join("::")
+        entries.concat(@entries_tree.search("#{namespace}::#{name}"))
+      end
+
+      # Inherited constants
+      if name.end_with?("::")
+        entries.concat(inherited_constant_completion_candidates(nil, nesting + [name]))
+      else
+        entries.concat(inherited_constant_completion_candidates(name, nesting))
+      end
+
+      # Top level constants
+      entries.concat(@entries_tree.search(name))
+      entries.uniq!
+      T.cast(
+        entries,
+        T::Array[T::Array[T.any(
+          Entry::Constant,
+          Entry::ConstantAlias,
+          Entry::Namespace,
+          Entry::UnresolvedConstantAlias,
+        )]],
+      )
+    end
+
     # Resolve a constant to its declaration based on its name and the nesting where the reference was found. Parameter
     # documentation:
     #
@@ -365,12 +423,10 @@ def index_single(indexable_path, source = nil, collect_comments: true)
     # aliases, so we have to invoke `follow_aliased_namespace` again to check until we only return a real name
     sig { params(name: String, seen_names: T::Array[String]).returns(String) }
     def follow_aliased_namespace(name, seen_names = [])
-      return name if @entries[name]
-
       parts = name.split("::")
       real_parts = []
 
-      (parts.length - 1).downto(0).each do |i|
+      (parts.length - 1).downto(0) do |i|
         current_name = T.must(parts[0..i]).join("::")
         entry = @entries[current_name]&.first
 
@@ -824,7 +880,7 @@ def resolve_alias(entry, seen_names)
       )]))
     end
     def lookup_enclosing_scopes(name, nesting, seen_names)
-      nesting.length.downto(1).each do |i|
+      nesting.length.downto(1) do |i|
         namespace = T.must(nesting[0...i]).join("::")
 
         # If we find an entry with `full_name` directly, then we can already return it, even if it contains aliases -
@@ -871,6 +927,51 @@ def lookup_ancestor_chain(name, nesting, seen_names)
       nil
     end
 
+    sig do
+      params(
+        name: T.nilable(String),
+        nesting: T::Array[String],
+      ).returns(T::Array[T::Array[T.any(
+        Entry::Namespace,
+        Entry::ConstantAlias,
+        Entry::UnresolvedConstantAlias,
+        Entry::Constant,
+      )]])
+    end
+    def inherited_constant_completion_candidates(name, nesting)
+      namespace_entries = if name
+        *nesting_parts, constant_name = build_non_redundant_full_name(name, nesting).split("::")
+        return [] if nesting_parts.empty?
+
+        resolve(nesting_parts.join("::"), [])
+      else
+        resolve(nesting.join("::"), [])
+      end
+      return [] unless namespace_entries
+
+      ancestors = linearized_ancestors_of(T.must(namespace_entries.first).name)
+      candidates = ancestors.flat_map do |ancestor_name|
+        @entries_tree.search("#{ancestor_name}::#{constant_name}")
+      end
+
+      # For candidates with the same name, we must only show the first entry in the inheritance chain, since that's the
+      # one the user will be referring to in completion
+      completion_items = candidates.each_with_object({}) do |entries, hash|
+        *parts, short_name = T.must(entries.first).name.split("::")
+        namespace_name = parts.join("::")
+        ancestor_index = ancestors.index(namespace_name)
+        existing_entry, existing_entry_index = hash[short_name]
+
+        next unless ancestor_index && (!existing_entry || ancestor_index < existing_entry_index)
+
+        hash[short_name] = [entries, ancestor_index]
+      end
+
+      completion_items.values.map!(&:first)
+    rescue NonExistingNamespaceError
+      []
+    end
+
     # Removes redudancy from a constant reference's full name. For example, if we find a reference to `A::B::Foo` inside
     # of the ["A", "B"] nesting, then we should not concatenate the nesting with the name or else we'll end up with
     # `A::B::A::B::Foo`. This method will remove any redundant parts from the final name based on the reference and the

lib/ruby_indexer/test/index_test.rb

@@ -1863,5 +1863,73 @@ def self.my_singleton_def; end
     def test_entries_for_returns_nil_if_no_matches
       assert_nil(@index.entries_for("non_existing_file.rb", Entry::Namespace))
     end
+
+    def test_constant_completion_candidates_all_possible_constants
+      index(<<~RUBY)
+        XQRK = 3
+
+        module Bar
+          XQRK = 2
+        end
+
+        module Foo
+          XQRK = 1
+        end
+
+        module Namespace
+          XQRK = 0
+
+          class Baz
+            include Foo
+            include Bar
+          end
+        end
+      RUBY
+
+      result = @index.constant_completion_candidates("X", ["Namespace", "Baz"])
+
+      result.each do |entries|
+        name = entries.first.name
+        assert(entries.all? { |e| e.name == name })
+      end
+
+      assert_equal(["Namespace::XQRK", "Bar::XQRK", "XQRK"], result.map { |entries| entries.first.name })
+
+      result = @index.constant_completion_candidates("::X", ["Namespace", "Baz"])
+      assert_equal(["XQRK"], result.map { |entries| entries.first.name })
+    end
+
+    def test_constant_completion_candidates_for_empty_name
+      index(<<~RUBY)
+        module Foo
+          Bar = 1
+        end
+
+        class Baz
+          include Foo
+        end
+      RUBY
+
+      result = @index.constant_completion_candidates("Baz::", [])
+      assert_includes(result.map { |entries| entries.first.name }, "Foo::Bar")
+    end
+
+    def test_follow_alias_namespace
+      index(<<~RUBY)
+        module First
+          module Second
+            class Foo
+            end
+          end
+        end
+
+        module Namespace
+          Second = First::Second
+        end
+      RUBY
+
+      real_namespace = @index.follow_aliased_namespace("Namespace::Second")
+      assert_equal("First::Second", real_namespace)
+    end
   end
 end

lib/ruby_lsp/listeners/completion.rb

@@ -104,7 +104,7 @@ def on_constant_read_node_enter(node)
         name = constant_name(node)
         return if name.nil?
 
-        candidates = @index.prefix_search(name, @node_context.nesting)
+        candidates = @index.constant_completion_candidates(name, @node_context.nesting)
         candidates.each do |entries|
           complete_name = T.must(entries.first).name
           @response_builder << build_entry_completion(
@@ -124,7 +124,13 @@ def on_constant_path_node_enter(node)
         # no sigil, Sorbet will still provide completion for constants
         return if @sorbet_level != RubyDocument::SorbetLevel::Ignore
 
-        name = constant_name(node)
+        name = begin
+          node.full_name
+        rescue Prism::ConstantPathNode::MissingNodesInConstantPathError
+          node.slice
+        rescue Prism::ConstantPathNode::DynamicPartsInConstantPathError
+          nil
+        end
         return if name.nil?
 
         constant_path_completion(name, range_from_location(node.location))
@@ -230,7 +236,7 @@ def constant_path_completion(name, range)
 
         real_namespace = @index.follow_aliased_namespace(T.must(namespace_entries.first).name)
 
-        candidates = @index.prefix_search(
+        candidates = @index.constant_completion_candidates(
           "#{real_namespace}::#{incomplete_name}",
           top_level_reference ? [] : nesting,
         )
@@ -240,8 +246,16 @@ def constant_path_completion(name, range)
           first_entry = T.must(entries.first)
           next if first_entry.private? && !first_entry.name.start_with?("#{nesting}::")
 
-          constant_name = first_entry.name.delete_prefix("#{real_namespace}::")
-          full_name = aliased_namespace.empty? ? constant_name : "#{aliased_namespace}::#{constant_name}"
+          entry_name = first_entry.name
+          full_name = if aliased_namespace != real_namespace
+            constant_name = entry_name.delete_prefix("#{real_namespace}::")
+            aliased_namespace.empty? ? constant_name : "#{aliased_namespace}::#{constant_name}"
+          elsif !entry_name.start_with?(aliased_namespace)
+            *_, short_name = entry_name.split("::")
+            "#{aliased_namespace}::#{short_name}"
+          else
+            entry_name
+          end
 
           @response_builder << build_entry_completion(
             full_name,
@@ -545,7 +559,7 @@ def build_entry_completion(real_name, incomplete_name, range, entries, top_level
       sig { params(entry_name: String).returns(T::Boolean) }
       def top_level?(entry_name)
         nesting = @node_context.nesting
-        nesting.length.downto(0).each do |i|
+        nesting.length.downto(0) do |i|
           prefix = T.must(nesting[0...i]).join("::")
           full_name = prefix.empty? ? entry_name : "#{prefix}::#{entry_name}"
           next if full_name == entry_name