Merge pull request #784 from soutaro/goto-type-decl

Implement *go to type definition*

lib/steep/server/master.rb

@@ -555,7 +555,7 @@ def process_message_from_client(message)
                       definition_provider: true,
                       declaration_provider: false,
                       implementation_provider: true,
-                      type_definition_provider: false
+                      type_definition_provider: true
                     )
                   )
                 }
@@ -654,7 +654,7 @@ def process_message_from_client(message)
             end
           end
 
-        when "textDocument/definition", "textDocument/implementation"
+        when "textDocument/definition", "textDocument/implementation", "textDocument/typeDefinition"
           if path = pathname(message[:params][:textDocument][:uri])
             result_controller << group_request do |group|
               typecheck_workers.each do |worker|

lib/steep/server/type_check_worker.rb

@@ -30,13 +30,25 @@ def self.definition(id:, params:)
           )
         end
 
+        def self.type_definition(id:, params:)
+          new(
+            kind: :type_definition,
+            id: id,
+            params: params
+          )
+        end
+
         def implementation?
           kind == :implementation
         end
 
         def definition?
           kind == :definition
         end
+
+        def type_definition?
+          kind == :type_definition
+        end
       end
 
       include ChangeBuffer
@@ -75,6 +87,8 @@ def handle_request(request)
           queue << GotoJob.definition(id: request[:id], params: request[:params])
         when "textDocument/implementation"
           queue << GotoJob.implementation(id: request[:id], params: request[:params])
+        when "textDocument/typeDefinition"
+          queue << GotoJob.type_definition(id: request[:id], params: request[:params])
         end
       end
 
@@ -266,7 +280,7 @@ def stats_result
       end
 
       def goto(job)
-        path = Steep::PathHelper.to_pathname(job.params[:textDocument][:uri])
+        path = Steep::PathHelper.to_pathname(job.params[:textDocument][:uri]) or return []
         line = job.params[:position][:line] + 1
         column = job.params[:position][:character]
 
@@ -277,6 +291,8 @@ def goto(job)
             goto_service.definition(path: path, line: line, column: column)
           when job.implementation?
             goto_service.implementation(path: path, line: line, column: column)
+          when job.type_definition?
+            goto_service.type_definition(path: path, line: line, column: column)
           else
             raise
           end
@@ -293,7 +309,7 @@ def goto(job)
           path = project.absolute_path(path)
 
           {
-            uri: Steep::PathHelper.to_uri(path.to_s).to_s,
+            uri: Steep::PathHelper.to_uri(path).to_s,
             range: loc.as_lsp_range
           }
         end

lib/steep/services/goto_service.rb

@@ -13,13 +13,14 @@ def from_rbs?
         end
       end
 
-      ConstantQuery = Struct.new(:name, :from, keyword_init: true) do
+      class ConstantQuery < Struct.new(:name, :from, keyword_init: true)
         include SourceHelper
       end
-      MethodQuery = Struct.new(:name, :from, keyword_init: true) do
+      class MethodQuery < Struct.new(:name, :from, keyword_init: true)
         include SourceHelper
       end
-      TypeNameQuery = Struct.new(:name, keyword_init: true)
+      class TypeNameQuery < Struct.new(:name, keyword_init: true)
+      end
 
       attr_reader :type_check, :assignment
 
@@ -33,9 +34,7 @@ def project
       end
 
       def implementation(path:, line:, column:)
-        locations = []
-
-        # relative_path = project.relative_path(path)
+        locations = [] #: Array[loc]
 
         queries = query_at(path: path, line: line, column: column)
         queries.uniq!
@@ -55,7 +54,7 @@ def implementation(path:, line:, column:)
       end
 
       def definition(path:, line:, column:)
-        locations = []
+        locations = [] #: Array[loc]
 
         queries = query_at(path: path, line: line, column: column)
         queries.uniq!
@@ -90,6 +89,61 @@ def definition(path:, line:, column:)
         end
       end
 
+      def type_definition(path:, line:, column:)
+        locations = [] #: Array[loc]
+
+        relative_path = project.relative_path(path)
+
+        target = type_check.source_file?(relative_path) or return []
+        source = type_check.source_files[relative_path]
+        typing, signature = type_check_path(target: target, path: relative_path, content: source.content, line: line, column: column)
+
+        typing or return []
+        signature or return []
+
+        node, *_parents = typing.source.find_nodes(line: line, column: column)
+        node or return []
+
+        type = typing.type_of(node: node)
+
+        subtyping = signature.current_subtyping or return []
+
+        each_type_name(type).uniq.each do |name|
+          type_name_locations(name, locations: locations)
+        end
+
+        locations.uniq.select do |loc|
+          case loc
+          when RBS::Location
+            assignment =~ loc.name
+          else
+            true
+          end
+        end
+      end
+
+      def each_type_name(type, &block)
+        if block
+          case type
+          when AST::Types::Name::Instance, AST::Types::Name::Alias, AST::Types::Name::Singleton, AST::Types::Name::Interface
+            yield type.name
+          when AST::Types::Literal
+            yield type.back_type.name
+          when AST::Types::Nil
+            yield RBS::TypeName.new(name: :NilClass, namespace: RBS::Namespace.root)
+          when AST::Types::Boolean
+            yield RBS::BuiltinNames::TrueClass.name
+            yield RBS::BuiltinNames::FalseClass.name
+          end
+
+          type.each_child do |child|
+            each_type_name(child, &block)
+          end
+        else
+          enum_for :each_type_name, type
+        end
+      end
+
       def test_ast_location(loc, line:, column:)
         return false if line < loc.line
         return false if line == loc.line && column < loc.column
@@ -99,7 +153,7 @@ def test_ast_location(loc, line:, column:)
       end
 
       def query_at(path:, line:, column:)
-        queries = []
+        queries = [] #: Array[query]
 
         relative_path = project.relative_path(path)
 
@@ -110,29 +164,44 @@ def query_at(path:, line:, column:)
           if typing
             node, *parents = typing.source.find_nodes(line: line, column: column)
 
-            if node
+            if node && parents
               case node.type
               when :const, :casgn
-                if test_ast_location(node.location.name, line: line, column: column)
+                named_location = (_ = node.location) #: Parser::AST::_NamedLocation
+                if test_ast_location(named_location.name, line: line, column: column)
                   if name = typing.source_index.reference(constant_node: node)
                     queries << ConstantQuery.new(name: name, from: :ruby)
                   end
                 end
               when :def, :defs
-                if test_ast_location(node.location.name, line: line, column: column)
+                named_location = (_ = node.location) #: Parser::AST::_NamedLocation
+                if test_ast_location(named_location.name, line: line, column: column)
                   if method_context = typing.context_at(line: line, column: column).method_context
-                    type_name = method_context.method.defined_in
-                    name =
-                      if method_context.method.defs.any? {|defn| defn.member.singleton? }
-                        SingletonMethodName.new(type_name: type_name, method_name: method_context.name)
-                      else
-                        InstanceMethodName.new(type_name: type_name, method_name: method_context.name)
+                    if method = method_context.method
+                      method.defs.each do |defn|
+                        singleton_method =
+                          case defn.member
+                          when RBS::AST::Members::MethodDefinition
+                            defn.member.singleton?
+                          when RBS::AST::Members::Attribute
+                            defn.member.kind == :singleton
+                          end
+
+                        name =
+                          if singleton_method
+                            SingletonMethodName.new(type_name: defn.defined_in, method_name: method_context.name)
+                          else
+                            InstanceMethodName.new(type_name: defn.defined_in, method_name: method_context.name)
+                          end
+
+                        queries << MethodQuery.new(name: name, from: :ruby)
                       end
-                    queries << MethodQuery.new(name: name, from: :ruby)
+                    end
                   end
                 end
               when :send
-                if test_ast_location(node.location.selector, line: line, column: column)
+                location = (_ = node.location) #: Parser::AST::_SelectorLocation
+                if test_ast_location(location.selector, line: line, column: column)
                   if (parent = parents[0]) && parent.type == :block && parent.children[0] == node
                     node = parents[0]
                   end
@@ -161,6 +230,9 @@ def query_at(path:, line:, column:)
 
             locator = RBS::Locator.new(buffer: buffer, dirs: dirs, decls: decls)
             last, nodes = locator.find2(line: line, column: column)
+
+            nodes or raise
+
             case nodes[0]
             when RBS::AST::Declarations::Class, RBS::AST::Declarations::Module
               if last == :name
@@ -172,7 +244,8 @@ def query_at(path:, line:, column:)
               end
             when RBS::AST::Members::MethodDefinition
               if last == :name
-                type_name = nodes[1].name
+                parent_node = nodes[1] #: RBS::AST::Declarations::Class | RBS::AST::Declarations::Module | RBS::AST::Declarations::Interface
+                type_name = parent_node.name
                 method_name = nodes[0].name
                 if nodes[0].instance?
                   queries << MethodQuery.new(
@@ -226,13 +299,19 @@ def constant_definition_in_rbs(name, locations:)
 
           case entry = env.constant_entry(name)
           when RBS::Environment::ConstantEntry
-            locations << entry.decl.location&.[](:name)
+            if entry.decl.location
+              locations << entry.decl.location[:name]
+            end
           when RBS::Environment::ClassEntry, RBS::Environment::ModuleEntry
             entry.decls.each do |d|
-              locations << d.decl.location&.[](:name)
+              if d.decl.location
+                locations << d.decl.location[:name]
+              end
             end
           when RBS::Environment::ClassAliasEntry, RBS::Environment::ModuleAliasEntry
-            locations << entry.decl.location&.[](:new_name)
+            if entry.decl.location
+              locations << entry.decl.location[:new_name]
+            end
           end
         end
 
@@ -305,11 +384,17 @@ def method_locations(name, in_ruby:, in_rbs:, locations:)
             entry.declarations.each do |decl|
               case decl
               when RBS::AST::Members::MethodDefinition
-                locations << decl.location[:name]
+                if decl.location
+                  locations << decl.location[:name]
+                end
               when RBS::AST::Members::Alias
-                locations << decl.location[:new_name]
+                if decl.location
+                  locations << decl.location[:new_name]
+                end
               when RBS::AST::Members::AttrAccessor, RBS::AST::Members::AttrReader, RBS::AST::Members::AttrWriter
-                locations << decl.location[:name]
+                if decl.location
+                  locations << decl.location[:name]
+                end
               end
             end
           end
@@ -324,7 +409,18 @@ def type_name_locations(name, locations: [])
 
           entry = index.entry(type_name: name)
           entry.declarations.each do |decl|
-            locations << decl.location[:name]
+            case decl
+            when RBS::AST::Declarations::Class, RBS::AST::Declarations::Module, RBS::AST::Declarations::Interface, RBS::AST::Declarations::TypeAlias
+              if decl.location
+                locations << decl.location[:name]
+              end
+            when RBS::AST::Declarations::AliasDecl
+              if decl.location
+                locations << decl.location[:new_name]
+              end
+            else
+              raise
+            end
           end
         end
 

sig/shims/parser/nodes.rbs

@@ -33,5 +33,13 @@ module Parser
 
       %a{pure} def end: () -> Source::Range
     end
+
+    interface _NamedLocation
+      %a{pure} def name: () -> Source::Range
+    end
+
+    interface _SelectorLocation
+      %a{pure} def selector: () -> Source::Range
+    end
   end
 end

sig/steep/path_helper.rbs

@@ -10,6 +10,6 @@ module Steep
 
     # Receives a Pathname and returns a *file* URI
     #
-    def self?.to_uri: (Pathname path, ?dosish: bool) -> URI::File
+    def self?.to_uri: (Pathname | String path, ?dosish: bool) -> URI::File
   end
 end

sig/steep/server/type_check_worker.rbs

@@ -72,7 +72,7 @@ module Steep
       end
 
       class GotoJob
-        type kind = :implementation | :definition
+        type kind = :implementation | :definition | :type_definition
 
         attr_reader id: String
 
@@ -88,9 +88,13 @@ module Steep
 
         def self.definition: (id: String, params: params) -> GotoJob
 
+        def self.type_definition: (id: String, params: params) -> GotoJob
+
         def implementation?: () -> bool
 
         def definition?: () -> bool
+
+        def type_definition?: () -> bool
       end
 
       include ChangeBuffer