Merge pull request #267 from soutaro/index

Support `workspace/symbol` feature

lib/steep.rb

@@ -76,6 +76,10 @@
 require "steep/type_inference/method_call"
 require "steep/ast/types"
 
+require "steep/index/rbs_index"
+require "steep/index/signature_symbol_provider"
+require "steep/index/source_index"
+
 require "steep/server/utils"
 require "steep/server/base_worker"
 require "steep/server/code_worker"

lib/steep/index/rbs_index.rb

@@ -0,0 +1,334 @@
+module Steep
+  module Index
+    class RBSIndex
+      class TypeEntry
+        attr_reader :type_name
+        attr_reader :declarations
+        attr_reader :references
+
+        def initialize(type_name:)
+          @type_name = type_name
+          @declarations = Set[]
+          @references = Set[]
+        end
+
+        def add_declaration(decl)
+          case decl
+          when RBS::AST::Declarations::Class, RBS::AST::Declarations::Module
+            declarations << decl
+          when RBS::AST::Declarations::Interface
+            declarations << decl
+          when RBS::AST::Declarations::Alias
+            declarations << decl
+          else
+            raise "Unexpected type declaration: #{decl}"
+          end
+
+          self
+        end
+
+        def add_reference(ref)
+          case ref
+          when RBS::AST::Members::MethodDefinition
+            references << ref
+          when RBS::AST::Members::AttrAccessor, RBS::AST::Members::AttrReader, RBS::AST::Members::AttrWriter
+            references << ref
+          when RBS::AST::Members::InstanceVariable, RBS::AST::Members::ClassInstanceVariable, RBS::AST::Members::ClassVariable
+            references << ref
+          when RBS::AST::Members::Include, RBS::AST::Members::Extend
+            references << ref
+          when RBS::AST::Declarations::Module, RBS::AST::Declarations::Class
+            references << ref
+          when RBS::AST::Declarations::Constant, RBS::AST::Declarations::Global
+            references << ref
+          when RBS::AST::Declarations::Alias
+            references << ref
+          else
+            raise "Unexpected type reference: #{ref}"
+          end
+
+          self
+        end
+      end
+
+      class MethodEntry
+        attr_reader :method_name
+        attr_reader :declarations
+        attr_reader :references
+
+        def initialize(method_name:)
+          @method_name = method_name
+          @declarations = Set[]
+          @references = Set[]
+        end
+
+        def add_declaration(decl)
+          case decl
+          when RBS::AST::Members::MethodDefinition,
+            RBS::AST::Members::Alias,
+            RBS::AST::Members::AttrWriter,
+            RBS::AST::Members::AttrReader,
+            RBS::AST::Members::AttrAccessor
+            declarations << decl
+          else
+            raise "Unexpected method declaration: #{decl}"
+          end
+
+          self
+        end
+      end
+
+      class ConstantEntry
+        attr_reader :const_name
+        attr_reader :declarations
+
+        def initialize(const_name:)
+          @const_name = const_name
+          @declarations = Set[]
+        end
+
+        def add_declaration(decl)
+          case decl
+          when RBS::AST::Declarations::Constant
+            declarations << decl
+          else
+            raise
+          end
+
+          self
+        end
+      end
+
+      class GlobalEntry
+        attr_reader :global_name
+        attr_reader :declarations
+
+        def initialize(global_name:)
+          @global_name = global_name
+          @declarations = Set[]
+        end
+
+        def add_declaration(decl)
+          case decl
+          when RBS::AST::Declarations::Global
+            declarations << decl
+          else
+            raise
+          end
+
+          self
+        end
+      end
+
+      attr_reader :type_index
+      attr_reader :method_index
+      attr_reader :const_index
+      attr_reader :global_index
+
+      def initialize()
+        @type_index = {}
+        @method_index = {}
+        @const_index = {}
+        @global_index = {}
+      end
+
+      def entry(type_name: nil, method_name: nil, const_name: nil, global_name: nil)
+        case
+        when type_name
+          type_index[type_name] ||= TypeEntry.new(type_name: type_name)
+        when method_name
+          method_index[method_name] ||= MethodEntry.new(method_name: method_name)
+        when const_name
+          const_index[const_name] ||= ConstantEntry.new(const_name: const_name)
+        when global_name
+          global_index[global_name] ||= GlobalEntry.new(global_name: global_name)
+        else
+          raise
+        end
+      end
+
+      def each_entry(&block)
+        if block_given?
+          type_index.each_value(&block)
+          method_index.each_value(&block)
+          const_index.each_value(&block)
+          global_index.each_value(&block)
+        else
+          enum_for(:each_entry)
+        end
+      end
+
+      def add_type_declaration(type_name, declaration)
+        entry(type_name: type_name).add_declaration(declaration)
+      end
+
+      def add_method_declaration(method_name, member)
+        entry(method_name: method_name).add_declaration(member)
+      end
+
+      def add_constant_declaration(const_name, decl)
+        entry(const_name: const_name).add_declaration(decl)
+      end
+
+      def add_global_declaration(global_name, decl)
+        entry(global_name: global_name).add_declaration(decl)
+      end
+
+      def each_declaration(type_name: nil, method_name: nil, const_name: nil, global_name: nil, &block)
+        if block
+          entry = entry(type_name: type_name, method_name: method_name, const_name: const_name, global_name: global_name)
+          entry.declarations.each(&block)
+        else
+          enum_for(:each_declaration, type_name: type_name, method_name: method_name, const_name: const_name, global_name: global_name)
+        end
+      end
+
+      def add_type_reference(type_name, ref)
+        entry(type_name: type_name).add_reference(ref)
+      end
+
+      def each_reference(type_name: nil, &block)
+        if block
+          entry(type_name: type_name).references.each(&block)
+        else
+          enum_for(:each_reference, type_name: type_name)
+        end
+      end
+
+      class Builder
+        attr_reader :index
+
+        def initialize(index:)
+          @index = index
+        end
+
+        def member(type_name, member)
+          case member
+          when RBS::AST::Members::MethodDefinition
+            member.types.each do |method_type|
+              method_type.each_type do |type|
+                type_reference type, from: member
+              end
+            end
+
+            if member.instance?
+              method_name = InstanceMethodName.new(type_name: type_name, method_name: member.name)
+              index.add_method_declaration(method_name, member)
+            end
+
+            if member.singleton?
+              method_name = SingletonMethodName.new(type_name: type_name, method_name: member.name)
+              index.add_method_declaration(method_name, member)
+            end
+
+          when RBS::AST::Members::AttrAccessor, RBS::AST::Members::AttrReader, RBS::AST::Members::AttrWriter
+            type_reference member.type, from: member
+
+            if member.is_a?(RBS::AST::Members::AttrReader) || member.is_a?(RBS::AST::Members::AttrAccessor)
+              method_name = case member.kind
+                            when :instance
+                              InstanceMethodName.new(type_name: type_name, method_name: member.name)
+                            when :singleton
+                              SingletonMethodName.new(type_name: type_name, method_name: member.name)
+                            end
+              index.add_method_declaration(method_name, member)
+            end
+
+            if member.is_a?(RBS::AST::Members::AttrWriter) || member.is_a?(RBS::AST::Members::AttrAccessor)
+              method_name = case member.kind
+                            when :instance
+                              InstanceMethodName.new(type_name: type_name, method_name: "#{member.name}=".to_sym)
+                            when :singleton
+                              SingletonMethodName.new(type_name: type_name, method_name: "#{member.name}=".to_sym)
+                            end
+              index.add_method_declaration(method_name, member)
+            end
+
+          when RBS::AST::Members::InstanceVariable, RBS::AST::Members::ClassVariable, RBS::AST::Members::ClassInstanceVariable
+            type_reference member.type, from: member
+
+          when RBS::AST::Members::Include, RBS::AST::Members::Extend
+            index.add_type_reference member.name, member
+            member.args.each do |type|
+              type_reference type, from: member
+            end
+
+          when RBS::AST::Members::Alias
+            if member.instance?
+              new_name = InstanceMethodName.new(type_name: type_name, method_name: member.new_name)
+              index.add_method_declaration(new_name, member)
+            end
+
+            if member.singleton?
+              new_name = SingletonMethodName.new(type_name: type_name, method_name: member.new_name)
+              index.add_method_declaration(new_name, member)
+            end
+          end
+        end
+
+        def type_reference(type, from:)
+          case type
+          when RBS::Types::ClassInstance, RBS::Types::ClassSingleton, RBS::Types::Alias, RBS::Types::Interface
+            index.add_type_reference(type.name, from)
+          end
+
+          type.each_type do |ty|
+            type_reference ty, from: from
+          end
+        end
+
+        def env(env)
+          env.class_decls.each do |name, decl|
+            decl.decls.each do |d|
+              index.add_type_declaration(name, d.decl)
+
+              case d.decl
+              when RBS::AST::Declarations::Class
+                if super_class = d.decl.super_class
+                  index.add_type_reference(super_class.name, d.decl)
+                  super_class.args.each do |type|
+                    type_reference(type, from: d.decl)
+                  end
+                end
+              when RBS::AST::Declarations::Module
+                d.decl.self_types.each do |self_type|
+                  index.add_type_reference(self_type.name, d.decl)
+                  self_type.args.each do |type|
+                    type_reference(type, from: d.decl)
+                  end
+                end
+              end
+
+              d.decl.members.each do |member|
+                member(name, member)
+              end
+            end
+          end
+
+          env.interface_decls.each do |name, decl|
+            index.add_type_declaration(name, decl.decl)
+
+            decl.decl.members.each do |member|
+              member(name, member)
+            end
+          end
+
+          env.alias_decls.each do |name, decl|
+            index.add_type_declaration(name, decl.decl)
+            type_reference decl.decl.type, from: decl.decl
+          end
+
+          env.constant_decls.each do |name, decl|
+            index.add_constant_declaration(name, decl.decl)
+            type_reference decl.decl.type, from: decl.decl
+          end
+
+          env.global_decls.each do |name, decl|
+            index.add_global_declaration(name, decl.decl)
+            type_reference decl.decl.type, from: decl.decl
+          end
+        end
+      end
+    end
+  end
+end