type-guessr 0.0.2 → 0.0.3

data/lib/type_guessr/core/registry/signature_registry.rb

@@ -1,10 +1,9 @@
 # frozen_string_literal: true
 
-require "singleton"
 require "rbs"
 require_relative "../types"
 require_relative "../logger"
-require_relative "../converter/rbs_converter"
+require_relative "../converter"
 require_relative "../type_serializer"
 
 module TypeGuessr
@@ -13,7 +12,9 @@ module TypeGuessr
       # Preloads stdlib RBS signatures and provides O(1) hash lookup
       # Singleton to ensure RBS is loaded only once across all usages
       class SignatureRegistry
-        include Singleton
+        class << self
+          attr_accessor :instance
+        end
 
         # Represents a method signature entry from RBS
         # Handles overload resolution and type conversion
@@ -205,11 +206,14 @@ module TypeGuessr
         class GemMethodEntry
           attr_reader :params
 
-          def initialize(return_type, params = [])
+          def initialize(return_type, params = [], skip_stdlib_rbs: false)
             @return_type = return_type
             @params = params
+            @skip_stdlib_rbs = skip_stdlib_rbs
           end
 
+          def skip_stdlib_rbs? = @skip_stdlib_rbs
+
           def return_type(_arg_types = [])
             @return_type
           end
@@ -233,14 +237,16 @@ module TypeGuessr
           # Get formatted signature strings for display
           # @return [Array<String>] human-readable method signatures
           def signature_strings
-            param_str = @params.map(&:to_s).join(", ")
+            param_str = @params.join(", ")
             ["(#{param_str}) -> #{@return_type}"]
           end
         end
 
+        attr_accessor :code_index
         attr_writer :on_demand_inferrer # ->(class_name, method_name, kind) { ... }
 
-        def initialize
+        def initialize(code_index: nil)
+          @code_index = code_index
           @instance_methods = {} # { "String" => { "upcase" => MethodEntry } }
           @class_methods = {}    # { "File" => { "read" => MethodEntry } }
           @converter = Converter::RBSConverter.new
@@ -265,19 +271,42 @@ module TypeGuessr
         end
 
         # Look up instance method entry
+        # Falls back to ancestor chain traversal when code_index is available
         # @param class_name [String] the class name
         # @param method_name [String] the method name
         # @return [MethodEntry, nil] method entry or nil if not found
         def lookup(class_name, method_name)
-          @instance_methods.dig(class_name, method_name)
+          result = @instance_methods.dig(class_name, method_name)
+          return result if result
+
+          return nil unless @code_index
+
+          @code_index.ancestors_of(class_name).each do |ancestor|
+            next if ancestor == class_name
+
+            result = @instance_methods.dig(ancestor, method_name)
+            return result if result
+          end
+          nil
         end
 
         # Look up class method entry
+        # Falls back to ancestor chain traversal when code_index is available
         # @param class_name [String] the class name
         # @param method_name [String] the method name
         # @return [MethodEntry, nil] method entry or nil if not found
         def lookup_class_method(class_name, method_name)
-          @class_methods.dig(class_name, method_name)
+          result = @class_methods.dig(class_name, method_name)
+          return result if result
+          return nil unless @code_index
+
+          @code_index.ancestors_of(class_name).each do |ancestor|
+            next if ancestor == class_name
+
+            result = @class_methods.dig(ancestor, method_name)
+            return result if result
+          end
+          nil
         end
 
         # Get method return type (convenience method matching old SignatureProvider API)
@@ -358,21 +387,33 @@ module TypeGuessr
         end
 
         # Register a single gem instance method
-        # Skips if RBS entry already exists (RBS takes priority)
-        def register_gem_method(class_name, method_name, return_type, params = [])
+        # Skips if RBS entry already exists (RBS takes priority).
+        # force: true overwrites existing GemMethodEntry (but never MethodEntry/RBS).
+        def register_gem_method(class_name, method_name, return_type, params = [], force: false)
           @instance_methods[class_name] ||= {}
-          return if @instance_methods[class_name].key?(method_name)
 
-          @instance_methods[class_name][method_name] = GemMethodEntry.new(return_type, params)
+          existing = @instance_methods[class_name][method_name]
+          if existing
+            return if existing.is_a?(MethodEntry)
+            return unless force
+          end
+
+          @instance_methods[class_name][method_name] = GemMethodEntry.new(return_type, params, skip_stdlib_rbs: force)
         end
 
         # Register a single gem class method
-        # Skips if RBS entry already exists (RBS takes priority)
-        def register_gem_class_method(class_name, method_name, return_type, params = [])
+        # Skips if RBS entry already exists (RBS takes priority).
+        # force: true overwrites existing GemMethodEntry (but never MethodEntry/RBS).
+        def register_gem_class_method(class_name, method_name, return_type, params = [], force: false)
           @class_methods[class_name] ||= {}
-          return if @class_methods[class_name].key?(method_name)
 
-          @class_methods[class_name][method_name] = GemMethodEntry.new(return_type, params)
+          existing = @class_methods[class_name][method_name]
+          if existing
+            return if existing.is_a?(MethodEntry)
+            return unless force
+          end
+
+          @class_methods[class_name][method_name] = GemMethodEntry.new(return_type, params, skip_stdlib_rbs: force)
         end
 
         # Bulk load gem cache data into the registry
@@ -447,9 +488,25 @@ module TypeGuessr
           env = RBS::Environment.from_loader(loader).resolve_type_names
           builder = RBS::DefinitionBuilder.new(env: env)
 
+          # Collect type parameter names for generic classes (e.g., "Set" => [:A])
+          # This lets RBSConverter build ClassInstance with named type_params
+          class_type_params = {}
+          env.class_decls.each do |type_name, entry|
+            params = entry.type_params.map(&:name)
+            class_type_params[type_name.to_s.delete_prefix("::")] = params if params.any?
+          end
+          @converter = Converter::RBSConverter.new(class_type_params)
+
           env.class_decls.each_key do |type_name|
             load_class_definitions(type_name, builder)
           end
+
+          env.class_alias_decls.each_value do |entry|
+            alias_name = entry.decl.new_name.to_s.delete_prefix("::")
+            original_name = entry.decl.old_name.to_s.delete_prefix("::")
+            @instance_methods[alias_name] ||= @instance_methods[original_name] if @instance_methods[original_name]
+            @class_methods[alias_name] ||= @class_methods[original_name] if @class_methods[original_name]
+          end
         rescue StandardError => e
           Logger.error("Failed to preload RBS environment", e)
         end

data/lib/type_guessr/core/registry.rb

@@ -0,0 +1,6 @@
+# frozen_string_literal: true
+
+require_relative "registry/method_registry"
+require_relative "registry/instance_variable_registry"
+require_relative "registry/class_variable_registry"
+require_relative "registry/signature_registry"

data/lib/type_guessr/core/type_serializer.rb

@@ -14,7 +14,9 @@ module TypeGuessr
         when Types::Unknown
           { "_type" => "Unknown" }
         when Types::ClassInstance
-          { "_type" => "ClassInstance", "name" => type.name }
+          h = { "_type" => "ClassInstance", "name" => type.name }
+          h["type_params"] = type.type_params.to_h { |k, v| [k.to_s, serialize(v)] } if type.type_params&.any?
+          h
         when Types::SingletonType
           { "_type" => "SingletonType", "name" => type.name }
         when Types::TupleType
@@ -49,19 +51,21 @@ module TypeGuessr
       module_function def deserialize(hash)
         case hash["_type"]
         when "Unguessed" then Types::Unguessed.instance
-        when "Unknown"          then Types::Unknown.instance
-        when "ClassInstance"    then Types::ClassInstance.for(hash["name"])
-        when "SingletonType"    then Types::SingletonType.new(hash["name"])
-        when "ArrayType"        then Types::ArrayType.new(deserialize(hash["element_type"]))
-        when "TupleType"        then Types::TupleType.new(hash["element_types"].map { |t| deserialize(t) })
-        when "HashType"         then Types::HashType.new(deserialize(hash["key_type"]), deserialize(hash["value_type"]))
-        when "RangeType"        then Types::RangeType.new(deserialize(hash["element_type"]))
-        when "HashShape"        then Types::HashShape.new(hash["fields"].to_h { |k, v| [k.to_sym, deserialize(v)] })
-        when "Union"            then Types::Union.new(hash["types"].map { |t| deserialize(t) })
-        when "TypeVariable"     then Types::TypeVariable.new(hash["name"].to_sym)
-        when "SelfType"         then Types::SelfType.instance
-        when "ForwardingArgs"   then Types::ForwardingArgs.instance
-        when "MethodSignature"  then deserialize_method_signature(hash)
+        when "Unknown" then Types::Unknown.instance
+        when "ClassInstance"
+          type_params = hash["type_params"]&.to_h { |k, v| [k.to_sym, deserialize(v)] }
+          Types::ClassInstance.for(hash["name"], type_params)
+        when "SingletonType" then Types::SingletonType.new(hash["name"])
+        when "ArrayType" then Types::ArrayType.new(deserialize(hash["element_type"]))
+        when "TupleType" then Types::TupleType.new(hash["element_types"].map { |t| deserialize(t) })
+        when "HashType" then Types::HashType.new(deserialize(hash["key_type"]), deserialize(hash["value_type"]))
+        when "RangeType" then Types::RangeType.new(deserialize(hash["element_type"]))
+        when "HashShape" then Types::HashShape.new(hash["fields"].to_h { |k, v| [k.to_sym, deserialize(v)] })
+        when "Union" then Types::Union.new(hash["types"].map { |t| deserialize(t) })
+        when "TypeVariable" then Types::TypeVariable.new(hash["name"].to_sym)
+        when "SelfType" then Types::SelfType.instance
+        when "ForwardingArgs" then Types::ForwardingArgs.instance
+        when "MethodSignature" then deserialize_method_signature(hash)
         else
           raise ArgumentError, "Unknown type: #{hash["_type"]}"
         end

data/lib/type_guessr/core/type_simplifier.rb

@@ -33,8 +33,8 @@ module TypeGuessr
         # 1. Single element: unwrap
         return types.first if types.size == 1
 
-        # 2. Filter to most general types (remove children when parent is present)
-        types = filter_to_most_general_types(types) if @code_index
+        # 2. Remove subclass types when parent class is also present
+        types = remove_subclass_types(types) if @code_index
 
         # 3. Check again after filtering
         return types.first if types.size == 1
@@ -45,10 +45,10 @@ module TypeGuessr
         Types::Union.new(types)
       end
 
-      # Filter out types whose ancestor is also in the list
+      # Remove types whose ancestor is also in the list (e.g., remove Integer when Numeric is present)
       # @param types [Array<Types::Type>] List of types
-      # @return [Array<Types::Type>] Filtered list with only the most general types
-      private def filter_to_most_general_types(types)
+      # @return [Array<Types::Type>] Filtered list with subclass types removed
+      private def remove_subclass_types(types)
         # Extract class names from ClassInstance types
         class_names = types.filter_map do |t|
           t.name if t.is_a?(Types::ClassInstance)

data/lib/type_guessr/core/types.rb

@@ -36,7 +36,7 @@ module TypeGuessr
         # Get type variable substitutions for this type
         # Used for substituting type variables in block parameters
         # @return [Hash{Symbol => Type}] type variable substitutions (e.g., { Elem: Integer, K: Symbol, V: String })
-        def type_variable_substitutions
+        def type_parameter_bindings
           {}
         end
 
@@ -68,42 +68,73 @@ module TypeGuessr
       end
 
       # ClassInstance - instance of a class
+      # @param type_params [Hash{Symbol => Type}, nil] type parameters (e.g., { A: String } for Set[String])
       class ClassInstance < Type
         CACHE = {} # rubocop:disable Style/MutableConstant
+        GENERIC_CACHE = {} # rubocop:disable Style/MutableConstant
 
         # Factory method that caches instances for reuse
         # @param name [String] The class name
+        # @param type_params [Hash{Symbol => Type}, nil] type parameters
         # @return [ClassInstance] Cached or new instance
-        def self.for(name)
-          CACHE[name] ||= new(name).freeze
+        def self.for(name, type_params = nil)
+          if type_params.nil?
+            CACHE[name] ||= new(name).freeze
+          else
+            key = [name, type_params]
+            GENERIC_CACHE[key] ||= new(name, type_params).freeze
+          end
         end
 
-        attr_reader :name
+        attr_reader :name, :type_params
 
-        def initialize(name)
+        def initialize(name, type_params = nil)
           super()
           @name = name
+          @type_params = type_params&.freeze
         end
 
         def eql?(other)
-          super && @name == other.name
+          super && @name == other.name && @type_params == other.type_params
         end
 
         def hash
-          [self.class, @name].hash
+          [self.class, @name, @type_params].hash
         end
 
         def to_s
-          case @name
-          when "NilClass" then "nil"
-          when "TrueClass" then "true"
-          when "FalseClass" then "false"
-          else @name
+          base = case @name
+                 when "NilClass" then "nil"
+                 when "TrueClass" then "true"
+                 when "FalseClass" then "false"
+                 else @name
+                 end
+          if @type_params&.any?
+            "#{base}[#{@type_params.values.join(", ")}]"
+          else
+            base
           end
         end
 
         def inspect
-          "#<ClassInstance:#{@name}>"
+          if @type_params&.any?
+            "#<ClassInstance:#{@name}[#{@type_params.map { |k, v| "#{k}=#{v.inspect}" }.join(", ")}]>"
+          else
+            "#<ClassInstance:#{@name}>"
+          end
+        end
+
+        def type_parameter_bindings
+          @type_params || {}
+        end
+
+        def substitute(substitutions)
+          return self if @type_params.nil? || @type_params.empty?
+
+          new_params = @type_params.transform_values { |t| t.substitute(substitutions) }
+          return self if new_params == @type_params
+
+          ClassInstance.new(@name, new_params)
         end
 
         def rbs_class_name
@@ -164,15 +195,17 @@ module TypeGuessr
         def to_s
           if bool_type?
             "bool"
+          elsif nullable_bool_type?
+            "bool?"
           elsif optional_type?
-            "?#{non_nil_type}"
+            "#{non_nil_type}?"
           else
             @types.map(&:to_s).sort.join(" | ")
           end
         end
 
         def inspect
-          "#<Union:#{@types.map(&:to_s).join("|")}>"
+          "#<Union:#{@types.join("|")}>"
         end
 
         def substitute(substitutions)
@@ -190,6 +223,15 @@ module TypeGuessr
           has_true && has_false
         end
 
+        private def nullable_bool_type?
+          return false unless @types.size == 3
+
+          has_true = @types.any? { |t| t.is_a?(ClassInstance) && t.name == "TrueClass" }
+          has_false = @types.any? { |t| t.is_a?(ClassInstance) && t.name == "FalseClass" }
+          has_nil = @types.any? { |t| nil_type?(t) }
+          has_true && has_false && has_nil
+        end
+
         private def optional_type?
           @types.size == 2 && @types.any? { |t| nil_type?(t) }
         end
@@ -225,7 +267,7 @@ module TypeGuessr
         private def simplify_if_unknown_present(types)
           return types if types.size <= 1
 
-          has_unknown = types.any? { |t| t.is_a?(Unknown) }
+          has_unknown = types.any?(Unknown)
           has_unknown ? [Unknown.instance] : types
         end
 
@@ -270,7 +312,7 @@ module TypeGuessr
           "Array"
         end
 
-        def type_variable_substitutions
+        def type_parameter_bindings
           { Elem: @element_type }
         end
       end
@@ -321,7 +363,7 @@ module TypeGuessr
           "Array"
         end
 
-        def type_variable_substitutions
+        def type_parameter_bindings
           return { Elem: Unknown.instance } if @element_types.empty?
 
           unique = @element_types.uniq
@@ -368,7 +410,7 @@ module TypeGuessr
           "Hash"
         end
 
-        def type_variable_substitutions
+        def type_parameter_bindings
           { K: @key_type, V: @value_type }
         end
       end
@@ -409,7 +451,7 @@ module TypeGuessr
           "Range"
         end
 
-        def type_variable_substitutions
+        def type_parameter_bindings
           { Elem: @element_type }
         end
       end
@@ -469,7 +511,7 @@ module TypeGuessr
           "Hash"
         end
 
-        def type_variable_substitutions
+        def type_parameter_bindings
           key_type = ClassInstance.for("Symbol")
           value_types = @fields.values.uniq
           value_type = if value_types.empty?
@@ -574,7 +616,7 @@ module TypeGuessr
         end
 
         def to_s
-          params_str = @params.map(&:to_s).join(", ")
+          params_str = @params.join(", ")
           "(#{params_str}) -> #{@return_type}"
         end
 

data/lib/type_guessr/core.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+# Configuration
+require_relative "core/config"
+
+# Core type system
+require_relative "core/types"
+
+# Node infrastructure
+require_relative "core/node_key_generator"
+require_relative "core/node_context_helper"
+require_relative "core/ir"
+require_relative "core/index"
+
+# Registries
+require_relative "core/registry"
+
+# Converters
+require_relative "core/converter"
+
+# Inference
+require_relative "core/inference"
+
+# Utilities
+require_relative "core/signature_builder"
+require_relative "core/type_simplifier"
+require_relative "core/type_serializer"
+require_relative "core/cache"
+require_relative "core/logger"

data/lib/type_guessr/mcp/server.rb

@@ -10,26 +10,7 @@ require "ruby_indexer/ruby_indexer"
 require "ruby_lsp/internal"
 
 # Load core components
-require_relative "../core/types"
-require_relative "../core/node_key_generator"
-require_relative "../core/node_context_helper"
-require_relative "../core/ir/nodes"
-require_relative "../core/index/location_index"
-require_relative "../core/converter/prism_converter"
-require_relative "../core/converter/rbs_converter"
-require_relative "../core/inference/result"
-require_relative "../core/inference/resolver"
-require_relative "../core/registry/signature_registry"
-require_relative "../core/registry/method_registry"
-require_relative "../core/registry/instance_variable_registry"
-require_relative "../core/registry/class_variable_registry"
-require_relative "../core/signature_builder"
-require_relative "../core/type_simplifier"
-require_relative "../core/type_serializer"
-require_relative "../core/logger"
-require_relative "../core/cache/gem_signature_cache"
-require_relative "../core/cache/gem_dependency_resolver"
-require_relative "../core/cache/gem_signature_extractor"
+require_relative "../core"
 
 # Load CodeIndexAdapter, StandaloneRuntime, and FileWatcher
 require_relative "../../ruby_lsp/type_guessr/code_index_adapter"
@@ -87,7 +68,8 @@ module TypeGuessr
 
       private def build_runtime(ruby_index)
         code_index = RubyLsp::TypeGuessr::CodeIndexAdapter.new(ruby_index)
-        signature_registry = Core::Registry::SignatureRegistry.instance
+        signature_registry = Core::Registry::SignatureRegistry.new(code_index: code_index)
+        Core::Registry::SignatureRegistry.instance = signature_registry
 
         method_registry = Core::Registry::MethodRegistry.new(code_index: code_index)
         ivar_registry = Core::Registry::InstanceVariableRegistry.new(code_index: code_index)
@@ -380,49 +362,71 @@ module TypeGuessr
       end
 
       private def build_tools
-        [build_infer_type_tool, build_get_method_signature_tool, build_search_methods_tool]
+        [build_get_method_sources_tool, build_get_method_signatures_tool, build_search_methods_tool]
       end
 
-      private def build_infer_type_tool
+      private def build_get_method_sources_tool
         runtime = @runtime
         to_response = method(:json_response)
 
         ::MCP::Tool.define(
-          name: "infer_type",
-          description: "Infer the type of a variable, expression, or method at a specific location in a Ruby file. " \
-                       "Returns the guessed type based on heuristic analysis.",
+          name: "get_method_sources",
+          description: "Get source code of methods by class and method name. " \
+                       "Returns the full method definition with file path and line number. " \
+                       "Use when you need to read a method's implementation without knowing " \
+                       "the exact file location.",
           input_schema: {
             type: "object",
             properties: {
-              file_path: { type: "string", description: "Absolute path to the Ruby file" },
-              line: { type: "integer", description: "Line number (1-based)" },
-              column: { type: "integer", description: "Column number (0-based)" }
+              methods: {
+                type: "array",
+                items: {
+                  type: "object",
+                  properties: {
+                    class_name: { type: "string", description: "Fully qualified class name (e.g., 'User', 'Admin::User')" },
+                    method_name: { type: "string", description: "Method name (e.g., 'save', 'initialize')" }
+                  },
+                  required: %w[class_name method_name]
+                },
+                description: "Array of {class_name, method_name} to look up"
+              }
             },
-            required: %w[file_path line column]
+            required: %w[methods]
           }
-        ) do |file_path:, line:, column:, **|
-          to_response.call(runtime.infer_at(file_path, line, column))
+        ) do |methods:, **|
+          to_response.call(runtime.method_sources(methods))
         end
       end
 
-      private def build_get_method_signature_tool
+      private def build_get_method_signatures_tool
         runtime = @runtime
         to_response = method(:json_response)
 
         ::MCP::Tool.define(
-          name: "get_method_signature",
-          description: "Get the inferred signature of a method defined in the project. " \
-                       "Returns parameter types and return type.",
+          name: "get_method_signatures",
+          description: "Get inferred signatures for multiple methods in one call. " \
+                       "More efficient than calling get_method_signature repeatedly " \
+                       "when investigating a class or module.",
           input_schema: {
             type: "object",
             properties: {
-              class_name: { type: "string", description: "Fully qualified class name (e.g., 'User', 'Admin::User')" },
-              method_name: { type: "string", description: "Method name (e.g., 'save', 'initialize')" }
+              methods: {
+                type: "array",
+                items: {
+                  type: "object",
+                  properties: {
+                    class_name: { type: "string", description: "Fully qualified class name" },
+                    method_name: { type: "string", description: "Method name" }
+                  },
+                  required: %w[class_name method_name]
+                },
+                description: "Array of {class_name, method_name} to look up"
+              }
             },
-            required: %w[class_name method_name]
+            required: %w[methods]
           }
-        ) do |class_name:, method_name:, **|
-          to_response.call(runtime.method_signature(class_name, method_name))
+        ) do |methods:, **|
+          to_response.call(runtime.method_signatures(methods))
         end
       end
 
@@ -432,17 +436,22 @@ module TypeGuessr
 
         ::MCP::Tool.define(
           name: "search_methods",
-          description: "Search for method definitions in the project. " \
-                       "Supports patterns like 'User#save', 'save', or 'Admin::*'.",
+          description: "Search for method definitions across the project with type-aware results. " \
+                       "Returns matching methods with their inferred signatures. Unlike grep, results " \
+                       "include class hierarchy context and inferred types. Supports patterns: " \
+                       "'User#save' (specific), 'save' (all classes), 'Admin::*' (namespace). " \
+                       "Use when exploring an unfamiliar codebase or finding which classes implement " \
+                       "a given method.",
           input_schema: {
             type: "object",
             properties: {
-              query: { type: "string", description: "Search query (e.g., 'User#save', 'save', 'initialize')" }
+              query: { type: "string", description: "Search query (e.g., 'User#save', 'save', 'initialize')" },
+              include_signatures: { type: "boolean", description: "Include inferred method signatures in results (default: false)" }
             },
             required: %w[query]
           }
-        ) do |query:, **|
-          to_response.call(runtime.search_methods(query))
+        ) do |query:, include_signatures: false, **|
+          to_response.call(runtime.search_methods(query, include_signatures: include_signatures))
         end
       end