Handle constant aliases and ability to follow them

lib/ruby_indexer/lib/ruby_indexer/index.rb

@@ -5,6 +5,8 @@ module RubyIndexer
   class Index
     extend T::Sig
 
+    class UnresolvableAliasError < StandardError; end
+
     # The minimum Jaro-Winkler similarity score for an entry to be considered a match for a given fuzzy search query
     ENTRY_SIMILARITY_THRESHOLD = 0.7
 
@@ -125,13 +127,23 @@ def fuzzy_search(query)
     # 3. Baz
     sig { params(name: String, nesting: T::Array[String]).returns(T.nilable(T::Array[Entry])) }
     def resolve(name, nesting)
-      (nesting.length + 1).downto(0).each do |i|
-        prefix = T.must(nesting[0...i]).join("::")
-        full_name = prefix.empty? ? name : "#{prefix}::#{name}"
-        entries = @entries[full_name]
-        return entries if entries
+      nesting.length.downto(0).each do |i|
+        namespace = T.must(nesting[0...i]).join("::")
+        full_name = namespace.empty? ? name : "#{namespace}::#{name}"
+
+        # If we find an entry with `full_name` directly, then we can already return it, even if it contains aliases -
+        # because the user might be trying to jump to the alias definition.
+        #
+        # However, if we don't find it, then we need to search for possible aliases in the namespace. For example, in
+        # the LSP itself we alias `RubyLsp::Interface` to `LanguageServer::Protocol::Interface`, which means doing
+        # `RubyLsp::Interface::Location` is allowed. For these cases, we need some way to realize that the
+        # `RubyLsp::Interface` part is an alias, that has to be resolved
+        entries = @entries[full_name] || @entries[follow_aliased_namespace(full_name)]
+        return entries.map { |e| e.is_a?(Entry::UnresolvedAlias) ? resolve_alias(e) : e } if entries
       end
 
+      nil
+    rescue UnresolvableAliasError
       nil
     end
 
@@ -152,6 +164,68 @@ def index_single(indexable_path, source = nil)
       # If `path` is a directory, just ignore it and continue indexing
     end
 
+    private
+
+    # Follows aliases in a namespace. The algorithm keeps checking if the name is an alias and then recursively follows
+    # it. The idea is that we test the name in parts starting from the complete name to the first namespace. For
+    # `Foo::Bar::Baz`, we would test:
+    # 1. Is `Foo::Bar::Baz` an alias? Get the target and recursively follow its target
+    # 2. Is `Foo::Bar` an alias? Get the target and recursively follow its target
+    # 3. Is `Foo` an alias? Get the target and recursively follow its target
+    #
+    # If we find an alias, then we want to follow its target. In the same example, if `Foo::Bar` is an alias to
+    # `Something::Else`, then we first discover `Something::Else::Baz`. But `Something::Else::Baz` might contain other
+    # aliases, so we have to invoke `follow_aliased_namespace` again to check until we only return a real name
+    sig { params(name: String).returns(String) }
+    def follow_aliased_namespace(name)
+      parts = name.split("::")
+      real_parts = []
+
+      (parts.length - 1).downto(0).each do |i|
+        current_name = T.must(parts[0..i]).join("::")
+        entry = @entries[current_name]&.first
+
+        case entry
+        when Entry::Alias
+          target = entry.target
+          return follow_aliased_namespace("#{target}::#{real_parts.join("::")}")
+        when Entry::UnresolvedAlias
+          resolved = resolve_alias(entry)
+
+          if resolved.is_a?(Entry::UnresolvedAlias)
+            raise UnresolvableAliasError, "The constant #{resolved.name} is an alias to a non existing constant"
+          end
+
+          target = resolved.target
+          return follow_aliased_namespace("#{target}::#{real_parts.join("::")}")
+        else
+          real_parts.unshift(T.must(parts[i]))
+        end
+      end
+
+      real_parts.join("::")
+    end
+
+    # Attempts to resolve an UnresolvedAlias into a resolved Alias. If the unresolved alias is pointing to a constant
+    # that doesn't exist, then we return the same UnresolvedAlias
+    sig { params(entry: Entry::UnresolvedAlias).returns(T.any(Entry::Alias, Entry::UnresolvedAlias)) }
+    def resolve_alias(entry)
+      target = resolve(entry.target, entry.nesting)
+      return entry unless target
+
+      target_name = T.must(target.first).name
+      resolved_alias = Entry::Alias.new(target_name, entry)
+
+      # Replace the UnresolvedAlias by a resolved one so that we don't have to do this again later
+      original_entries = T.must(@entries[entry.name])
+      original_entries.delete(entry)
+      original_entries << resolved_alias
+
+      @entries_tree.insert(entry.name, original_entries)
+
+      resolved_alias
+    end
+
     class Entry
       extend T::Sig
 
@@ -199,6 +273,58 @@ class Class < Namespace
 
       class Constant < Entry
       end
+
+      # An UnresolvedAlias points to a constant alias with a right hand side that has not yet been resolved. For
+      # example, if we find
+      #
+      # ```ruby
+      #   CONST = Foo
+      # ```
+      # Before we have discovered `Foo`, there's no way to eagerly resolve this alias to the correct target constant.
+      # All aliases are inserted as UnresolvedAlias in the index first and then we lazily resolve them to the correct
+      # target in [rdoc-ref:Index#resolve]. If the right hand side contains a constant that doesn't exist, then it's not
+      # possible to resolve the alias and it will remain an UnresolvedAlias until the right hand side constant exists
+      class UnresolvedAlias < Entry
+        extend T::Sig
+
+        sig { returns(String) }
+        attr_reader :target
+
+        sig { returns(T::Array[String]) }
+        attr_reader :nesting
+
+        sig do
+          params(
+            target: String,
+            nesting: T::Array[String],
+            name: String,
+            file_path: String,
+            location: YARP::Location,
+            comments: T::Array[String],
+          ).void
+        end
+        def initialize(target, nesting, name, file_path, location, comments) # rubocop:disable Metrics/ParameterLists
+          super(name, file_path, location, comments)
+
+          @target = target
+          @nesting = nesting
+        end
+      end
+
+      # Alias represents a resolved alias, which points to an existing constant target
+      class Alias < Entry
+        extend T::Sig
+
+        sig { returns(String) }
+        attr_reader :target
+
+        sig { params(target: String, unresolved_alias: UnresolvedAlias).void }
+        def initialize(target, unresolved_alias)
+          super(unresolved_alias.name, unresolved_alias.file_path, unresolved_alias.location, unresolved_alias.comments)
+
+          @target = target
+        end
+      end
     end
   end
 end

lib/ruby_indexer/lib/ruby_indexer/visitor.rb

@@ -36,9 +36,16 @@ def visit(node)
       when YARP::ModuleNode
         add_index_entry(node, Index::Entry::Module)
       when YARP::ConstantWriteNode, YARP::ConstantOrWriteNode
-        add_constant(node)
-      when YARP::ConstantPathWriteNode, YARP::ConstantPathOrWriteNode
-        add_constant_with_path(node)
+        name = fully_qualify_name(node.name.to_s)
+        add_constant(node, name)
+      when YARP::ConstantPathWriteNode, YARP::ConstantPathOrWriteNode, YARP::ConstantPathOperatorWriteNode,
+        YARP::ConstantPathAndWriteNode
+
+        # ignore variable constants like `var::FOO` or `self.class::FOO`
+        return unless node.target.parent.nil? || node.target.parent.is_a?(YARP::ConstantReadNode)
+
+        name = fully_qualify_name(node.target.location.slice)
+        add_constant(node, name)
       when YARP::CallNode
         message = node.message
         handle_private_constant(node) if message == "private_constant"
@@ -80,26 +87,39 @@ def handle_private_constant(node)
 
     sig do
       params(
-        node: T.any(YARP::ConstantWriteNode, YARP::ConstantOrWriteNode),
+        node: T.any(
+          YARP::ConstantWriteNode,
+          YARP::ConstantOrWriteNode,
+          YARP::ConstantPathWriteNode,
+          YARP::ConstantPathOrWriteNode,
+          YARP::ConstantPathOperatorWriteNode,
+          YARP::ConstantPathAndWriteNode,
+        ),
+        name: String,
       ).void
     end
-    def add_constant(node)
+    def add_constant(node, name)
+      value = node.value
       comments = collect_comments(node)
-      @index << Index::Entry::Constant.new(fully_qualify_name(node.name.to_s), @file_path, node.location, comments)
-    end
 
-    sig do
-      params(
-        node: T.any(YARP::ConstantPathWriteNode, YARP::ConstantPathOrWriteNode),
-      ).void
-    end
-    def add_constant_with_path(node)
-      # ignore variable constants like `var::FOO` or `self.class::FOO`
-      return unless node.target.parent.nil? || node.target.parent.is_a?(YARP::ConstantReadNode)
-
-      name = node.target.location.slice
-      comments = collect_comments(node)
-      @index << Index::Entry::Constant.new(fully_qualify_name(name), @file_path, node.location, comments)
+      @index << case value
+      when YARP::ConstantReadNode, YARP::ConstantPathNode
+        Index::Entry::UnresolvedAlias.new(value.slice, @stack.dup, name, @file_path, node.location, comments)
+      when YARP::ConstantWriteNode, YARP::ConstantAndWriteNode, YARP::ConstantOrWriteNode,
+        YARP::ConstantOperatorWriteNode
+
+        # If the right hand side is another constant assignment, we need to visit it because that constant has to be
+        # indexed too
+        visit(value)
+        Index::Entry::UnresolvedAlias.new(value.name.to_s, @stack.dup, name, @file_path, node.location, comments)
+      when YARP::ConstantPathWriteNode, YARP::ConstantPathOrWriteNode, YARP::ConstantPathOperatorWriteNode,
+        YARP::ConstantPathAndWriteNode
+
+        visit(value)
+        Index::Entry::UnresolvedAlias.new(value.target.slice, @stack.dup, name, @file_path, node.location, comments)
+      else
+        Index::Entry::Constant.new(name, @file_path, node.location, comments)
+      end
     end
 
     sig { params(node: T.any(YARP::ClassNode, YARP::ModuleNode), klass: T.class_of(Index::Entry)).void }

lib/ruby_indexer/test/constant_test.rb

@@ -181,5 +181,130 @@ module B
       b_const = @index["A::B::CONST_B"].first
       assert_equal(:private, b_const.visibility)
     end
+
+    def test_indexing_constant_aliases
+      index(<<~RUBY)
+        module A
+          module B
+            module C
+            end
+          end
+
+          FIRST = B::C
+        end
+
+        SECOND = A::FIRST
+      RUBY
+
+      unresolve_entry = @index["A::FIRST"].first
+      assert_instance_of(Index::Entry::UnresolvedAlias, unresolve_entry)
+      assert_equal(["A"], unresolve_entry.nesting)
+      assert_equal("B::C", unresolve_entry.target)
+
+      resolved_entry = @index.resolve("A::FIRST", []).first
+      assert_instance_of(Index::Entry::Alias, resolved_entry)
+      assert_equal("A::B::C", resolved_entry.target)
+    end
+
+    def test_aliasing_namespaces
+      index(<<~RUBY)
+        module A
+          module B
+            module C
+            end
+          end
+
+          ALIAS = B
+        end
+
+        module Other
+          ONE_MORE = A::ALIAS
+        end
+      RUBY
+
+      unresolve_entry = @index["A::ALIAS"].first
+      assert_instance_of(Index::Entry::UnresolvedAlias, unresolve_entry)
+      assert_equal(["A"], unresolve_entry.nesting)
+      assert_equal("B", unresolve_entry.target)
+
+      resolved_entry = @index.resolve("ALIAS", ["A"]).first
+      assert_instance_of(Index::Entry::Alias, resolved_entry)
+      assert_equal("A::B", resolved_entry.target)
+
+      resolved_entry = @index.resolve("ALIAS::C", ["A"]).first
+      assert_instance_of(Index::Entry::Module, resolved_entry)
+      assert_equal("A::B::C", resolved_entry.name)
+
+      unresolve_entry = @index["Other::ONE_MORE"].first
+      assert_instance_of(Index::Entry::UnresolvedAlias, unresolve_entry)
+      assert_equal(["Other"], unresolve_entry.nesting)
+      assert_equal("A::ALIAS", unresolve_entry.target)
+
+      resolved_entry = @index.resolve("Other::ONE_MORE::C", []).first
+      assert_instance_of(Index::Entry::Module, resolved_entry)
+    end
+
+    def test_indexing_same_line_constant_aliases
+      index(<<~RUBY)
+        module A
+          B = C = 1
+          D = E ||= 1
+          F = G::H &&= 1
+          I::J = K::L = M = 1
+        end
+      RUBY
+
+      # B and C
+      unresolve_entry = @index["A::B"].first
+      assert_instance_of(Index::Entry::UnresolvedAlias, unresolve_entry)
+      assert_equal(["A"], unresolve_entry.nesting)
+      assert_equal("C", unresolve_entry.target)
+
+      resolved_entry = @index.resolve("A::B", []).first
+      assert_instance_of(Index::Entry::Alias, resolved_entry)
+      assert_equal("A::C", resolved_entry.target)
+
+      constant = @index["A::C"].first
+      assert_instance_of(Index::Entry::Constant, constant)
+
+      # D and E
+      unresolve_entry = @index["A::D"].first
+      assert_instance_of(Index::Entry::UnresolvedAlias, unresolve_entry)
+      assert_equal(["A"], unresolve_entry.nesting)
+      assert_equal("E", unresolve_entry.target)
+
+      resolved_entry = @index.resolve("A::D", []).first
+      assert_instance_of(Index::Entry::Alias, resolved_entry)
+      assert_equal("A::E", resolved_entry.target)
+
+      # F and G::H
+      unresolve_entry = @index["A::F"].first
+      assert_instance_of(Index::Entry::UnresolvedAlias, unresolve_entry)
+      assert_equal(["A"], unresolve_entry.nesting)
+      assert_equal("G::H", unresolve_entry.target)
+
+      resolved_entry = @index.resolve("A::F", []).first
+      assert_instance_of(Index::Entry::Alias, resolved_entry)
+      assert_equal("A::G::H", resolved_entry.target)
+
+      # I::J, K::L and M
+      unresolve_entry = @index["A::I::J"].first
+      assert_instance_of(Index::Entry::UnresolvedAlias, unresolve_entry)
+      assert_equal(["A"], unresolve_entry.nesting)
+      assert_equal("K::L", unresolve_entry.target)
+
+      resolved_entry = @index.resolve("A::I::J", []).first
+      assert_instance_of(Index::Entry::Alias, resolved_entry)
+      assert_equal("A::K::L", resolved_entry.target)
+
+      # When we are resolving A::I::J, we invoke `resolve("K::L", ["A"])`, which recursively resolves A::K::L too.
+      # Therefore, both A::I::J and A::K::L point to A::M by the end of the previous resolve invocation
+      resolved_entry = @index["A::K::L"].first
+      assert_instance_of(Index::Entry::Alias, resolved_entry)
+      assert_equal("A::M", resolved_entry.target)
+
+      constant = @index["A::M"].first
+      assert_instance_of(Index::Entry::Constant, constant)
+    end
   end
 end