type-guessr 0.0.1 → 0.0.3

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

@@ -0,0 +1,543 @@
+# frozen_string_literal: true
+
+require "rbs"
+require_relative "../types"
+require_relative "../logger"
+require_relative "../converter"
+require_relative "../type_serializer"
+
+module TypeGuessr
+  module Core
+    module Registry
+      # Preloads stdlib RBS signatures and provides O(1) hash lookup
+      # Singleton to ensure RBS is loaded only once across all usages
+      class SignatureRegistry
+        class << self
+          attr_accessor :instance
+        end
+
+        # Represents a method signature entry from RBS
+        # Handles overload resolution and type conversion
+        class MethodEntry
+          # @param method_types [Array<RBS::MethodType>] RBS method type definitions
+          # @param converter [Converter::RBSConverter] RBS to internal type converter
+          def initialize(method_types, converter)
+            @method_types = method_types
+            @converter = converter
+          end
+
+          # Get the return type with overload resolution
+          # @param arg_types [Array<Types::Type>] argument types for overload matching
+          # @return [Types::Type] the return type
+          def return_type(arg_types = [])
+            best_match = find_best_overload(arg_types)
+            @converter.convert(best_match.type.return_type)
+          end
+
+          # Get block parameter types for this method
+          # @return [Array<Types::Type>] array of block parameter types (empty if no block)
+          def block_param_types
+            return @block_param_types if defined?(@block_param_types)
+
+            @block_param_types = compute_block_param_types
+          end
+
+          # Get method-level type parameter names from the best matching overload
+          # @param arg_types [Array<Types::Type>] argument types for overload matching
+          # @return [Array<Symbol>] type parameter names (e.g., [:U], [:X], [])
+          def type_params(arg_types = [])
+            find_best_overload(arg_types).type_params.map(&:name)
+          end
+
+          # Get the type variable used as the block return type
+          # Returns the method-level type variable that appears in the block's return type,
+          # allowing Resolver to substitute it with the actual block body type.
+          # @param arg_types [Array<Types::Type>] argument types for overload matching
+          # @return [Symbol, nil] type variable name (e.g., :U, :X) or nil if no block/no type var
+          def block_return_type_var(arg_types = [])
+            best = find_best_overload(arg_types)
+            return nil unless best.block
+
+            method_type_param_names = best.type_params.to_set(&:name)
+            return nil if method_type_param_names.empty?
+
+            extract_type_var_from_return(best.block.type.return_type, method_type_param_names)
+          end
+
+          # Get raw method signatures for display
+          # @return [Array<RBS::MethodType>] raw RBS method types
+          def signatures
+            @method_types
+          end
+
+          # Get formatted signature strings for display
+          # @return [Array<String>] human-readable method signatures
+          def signature_strings
+            @method_types.map(&:to_s)
+          end
+
+          private def compute_block_param_types
+            # Find the signature with a block
+            sig_with_block = @method_types.find(&:block)
+            return [] unless sig_with_block
+
+            extract_block_param_types(sig_with_block)
+          end
+
+          private def extract_block_param_types(method_type)
+            return [] unless method_type.block
+
+            block_func = method_type.block.type
+            return [] unless block_func.is_a?(RBS::Types::Function)
+
+            block_func.required_positionals.flat_map do |param|
+              # Handle Tuple types (e.g., [K, V] in Hash#each) by flattening
+              if param.type.is_a?(RBS::Types::Tuple)
+                param.type.types.map { |t| @converter.convert(t) }
+              else
+                [@converter.convert(param.type)]
+              end
+            end
+          end
+
+          # Extract a method-level type variable from the block return type
+          # Walks the RBS type tree to find a Variable whose name is in the method's type_params.
+          # Handles Union types like `nil | false | U` (filter_map).
+          # @param rbs_type [RBS::Types::t] the block return type from RBS
+          # @param method_type_param_names [Set<Symbol>] method-level type param names
+          # @return [Symbol, nil] the type variable name, or nil if not found
+          private def extract_type_var_from_return(rbs_type, method_type_param_names)
+            case rbs_type
+            when RBS::Types::Variable
+              method_type_param_names.include?(rbs_type.name) ? rbs_type.name : nil
+            when RBS::Types::Union
+              # Search union members for a type variable (e.g., nil | false | U)
+              rbs_type.types.each do |t|
+                found = extract_type_var_from_return(t, method_type_param_names)
+                return found if found
+              end
+              nil
+            end
+          end
+
+          # Find the best matching overload for given argument types
+          # @param arg_types [Array<Types::Type>] argument types
+          # @return [RBS::MethodType] best matching method type (first if no match)
+          private def find_best_overload(arg_types)
+            return @method_types.first if arg_types.empty?
+
+            # Score each overload
+            scored = @method_types.map do |method_type|
+              score = calculate_overload_score(method_type, arg_types)
+              [method_type, score]
+            end
+
+            # Return best scoring overload, or first if all scores are 0
+            best = scored.max_by { |_mt, score| score }
+            best[1].positive? ? best[0] : @method_types.first
+          end
+
+          # Calculate match score for an overload
+          private def calculate_overload_score(method_type, arg_types)
+            func = method_type.type
+            return 0 unless func.is_a?(RBS::Types::Function)
+
+            required = func.required_positionals
+            optional = func.optional_positionals
+            rest = func.rest_positionals
+
+            # Check argument count
+            min_args = required.size
+            max_args = rest ? Float::INFINITY : required.size + optional.size
+            return 0 unless arg_types.size.between?(min_args, max_args)
+
+            # Score each argument match
+            score = 0
+            arg_types.each_with_index do |arg_type, i|
+              param = if i < required.size
+                        required[i]
+                      elsif i < required.size + optional.size
+                        optional[i - required.size]
+                      else
+                        rest
+                      end
+
+              break unless param
+
+              score += type_match_score(arg_type, param.type)
+            end
+
+            score
+          end
+
+          # Calculate match score between our type and RBS parameter type
+          private def type_match_score(our_type, rbs_type)
+            case rbs_type
+            when RBS::Types::ClassInstance
+              class_name = rbs_type.name.to_s.delete_prefix("::")
+              return 2 if types_match_class?(our_type, class_name)
+
+              0
+            when RBS::Types::Union
+              max_score = rbs_type.types.map { |t| type_match_score(our_type, t) }.max || 0
+              max_score.positive? ? 1 : 0
+            else
+              # Unknown RBS type - give weak match to avoid penalizing
+              1
+            end
+          end
+
+          private def types_match_class?(our_type, class_name)
+            case our_type
+            when Types::ClassInstance
+              our_type.name == class_name
+            when Types::ArrayType, Types::TupleType
+              class_name == "Array"
+            when Types::HashShape
+              class_name == "Hash"
+            else
+              false
+            end
+          end
+        end
+
+        # Represents a cached gem method entry (pre-inferred, no RBS)
+        # Implements the same interface as MethodEntry for transparent lookup
+        class GemMethodEntry
+          attr_reader :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
+
+          def block_param_types
+            []
+          end
+
+          def type_params(_arg_types = [])
+            []
+          end
+
+          def block_return_type_var(_arg_types = [])
+            nil
+          end
+
+          def signatures
+            []
+          end
+
+          # Get formatted signature strings for display
+          # @return [Array<String>] human-readable method signatures
+          def signature_strings
+            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(code_index: nil)
+          @code_index = code_index
+          @instance_methods = {} # { "String" => { "upcase" => MethodEntry } }
+          @class_methods = {}    # { "File" => { "read" => MethodEntry } }
+          @converter = Converter::RBSConverter.new
+          @preloaded = false
+          @on_demand_inferrer = nil
+        end
+
+        # Preload stdlib RBS signatures
+        # @return [self]
+        def preload
+          return self if @preloaded
+
+          load_stdlib_rbs
+          @preloaded = true
+          self
+        end
+
+        # Check if preloading is complete
+        # @return [Boolean]
+        def preloaded?
+          @preloaded
+        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)
+          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)
+          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)
+        # Triggers on-demand inference if the entry has Unguessed return type.
+        # @param class_name [String] the class name
+        # @param method_name [String] the method name
+        # @param arg_types [Array<Types::Type>] argument types for overload matching
+        # @return [Types::Type] the return type (Unknown if not found)
+        def get_method_return_type(class_name, method_name, arg_types = [])
+          entry = lookup(class_name, method_name)
+          return Types::Unknown.instance unless entry
+
+          result = entry.return_type(arg_types)
+          if result.is_a?(Types::Unguessed)
+            try_on_demand_inference(class_name, method_name, :instance)
+            entry = lookup(class_name, method_name)
+            return Types::Unknown.instance unless entry
+
+            result = entry.return_type(arg_types)
+            return Types::Unknown.instance if result.is_a?(Types::Unguessed)
+          end
+          result
+        end
+
+        # Get class method return type (convenience method matching old SignatureProvider API)
+        # Triggers on-demand inference if the entry has Unguessed return type.
+        # @param class_name [String] the class name
+        # @param method_name [String] the method name
+        # @param arg_types [Array<Types::Type>] argument types for overload matching
+        # @return [Types::Type] the return type (Unknown if not found)
+        def get_class_method_return_type(class_name, method_name, arg_types = [])
+          entry = lookup_class_method(class_name, method_name)
+          return Types::Unknown.instance unless entry
+
+          result = entry.return_type(arg_types)
+          if result.is_a?(Types::Unguessed)
+            try_on_demand_inference(class_name, method_name, :class)
+            entry = lookup_class_method(class_name, method_name)
+            return Types::Unknown.instance unless entry
+
+            result = entry.return_type(arg_types)
+            return Types::Unknown.instance if result.is_a?(Types::Unguessed)
+          end
+          result
+        end
+
+        # Get block parameter types for a method
+        # @param class_name [String] the receiver class name
+        # @param method_name [String] the method name
+        # @return [Array<Types::Type>] array of block parameter types (empty if no block)
+        def get_block_param_types(class_name, method_name)
+          entry = lookup(class_name, method_name)
+          return [] unless entry
+
+          entry.block_param_types
+        end
+
+        # Get method signatures for hover display
+        # @param class_name [String] the class name
+        # @param method_name [String] the method name
+        # @return [Array<Signature>] wrapped signatures for compatibility
+        def get_method_signatures(class_name, method_name)
+          entry = lookup(class_name, method_name)
+          return [] unless entry
+
+          entry.signatures.map { |mt| Signature.new(mt) }
+        end
+
+        # Get class method signatures for hover display
+        # @param class_name [String] the class name
+        # @param method_name [String] the method name
+        # @return [Array<Signature>] wrapped signatures for compatibility
+        def get_class_method_signatures(class_name, method_name)
+          entry = lookup_class_method(class_name, method_name)
+          return [] unless entry
+
+          entry.signatures.map { |mt| Signature.new(mt) }
+        end
+
+        # Register a single gem instance method
+        # 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] ||= {}
+
+          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).
+        # 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] ||= {}
+
+          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
+        # @param cache_data [Hash] { class_name => { method_name => { "return_type" => ..., "params" => [...] } } }
+        # @param kind [Symbol] :instance or :class
+        def load_gem_cache(cache_data, kind: :instance)
+          target = kind == :class ? @class_methods : @instance_methods
+
+          cache_data.each do |class_name, methods|
+            target[class_name] ||= {}
+            methods.each do |method_name, entry_data|
+              next if target[class_name].key?(method_name) # RBS priority
+
+              return_type = TypeSerializer.deserialize(entry_data["return_type"])
+              params = (entry_data["params"] || []).map do |p|
+                Types::ParamSignature.new(
+                  name: p["name"].to_sym,
+                  kind: p["kind"].to_sym,
+                  type: TypeSerializer.deserialize(p["type"])
+                )
+              end
+              target[class_name][method_name] = GemMethodEntry.new(return_type, params)
+            end
+          end
+        end
+
+        # Replace Unguessed GemMethodEntry entries with inferred results.
+        # Only replaces entries that are GemMethodEntry with Unguessed return type.
+        # @param cache_data [Hash] { class_name => { method_name => { "return_type" => ..., "params" => [...] } } }
+        # @param kind [Symbol] :instance or :class
+        def replace_unguessed_entries(cache_data, kind: :instance)
+          target = kind == :class ? @class_methods : @instance_methods
+
+          cache_data.each do |class_name, methods|
+            next unless target[class_name]
+
+            methods.each do |method_name, entry_data|
+              existing = target[class_name][method_name]
+              next unless existing.is_a?(GemMethodEntry)
+              next unless existing.return_type.is_a?(Types::Unguessed)
+
+              return_type = TypeSerializer.deserialize(entry_data["return_type"])
+              params = (entry_data["params"] || []).map do |p|
+                Types::ParamSignature.new(
+                  name: p["name"].to_sym,
+                  kind: p["kind"].to_sym,
+                  type: TypeSerializer.deserialize(p["type"])
+                )
+              end
+              target[class_name][method_name] = GemMethodEntry.new(return_type, params)
+            end
+          end
+        end
+
+        # Wrapper for RBS method type (for compatibility with existing code)
+        class Signature
+          attr_reader :method_type
+
+          def initialize(method_type)
+            @method_type = method_type
+          end
+        end
+
+        private def try_on_demand_inference(class_name, method_name, kind)
+          return unless @on_demand_inferrer
+
+          @on_demand_inferrer.call(class_name, method_name, kind)
+        end
+
+        private def load_stdlib_rbs
+          loader = RBS::EnvironmentLoader.new
+          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
+
+        private def load_class_definitions(type_name, builder)
+          class_name = type_name.to_s.delete_prefix("::")
+
+          # Load instance methods
+          begin
+            definition = builder.build_instance(type_name)
+            @instance_methods[class_name] = build_method_entries(definition)
+          rescue RBS::NoTypeFoundError, RBS::NoSuperclassFoundError, RBS::NoMixinFoundError
+            # Skip classes with missing dependencies
+          end
+
+          # Load class methods (singleton)
+          begin
+            definition = builder.build_singleton(type_name)
+            @class_methods[class_name] = build_method_entries(definition)
+          rescue RBS::NoTypeFoundError, RBS::NoSuperclassFoundError, RBS::NoMixinFoundError
+            # Skip classes with missing dependencies
+          end
+        end
+
+        private def build_method_entries(definition)
+          # RBS methods hash uses Symbol keys, but we use String keys for lookup
+          definition.methods.to_h do |method_name, method_def|
+            [method_name.to_s, MethodEntry.new(method_def.method_types, @converter)]
+          end
+        end
+      end
+    end
+  end
+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/signature_builder.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+require_relative "types"
+
+module TypeGuessr
+  module Core
+    # Builds MethodSignature from DefNode using Resolver for type inference.
+    # Extracts the param formatting + type inference logic that was previously
+    # embedded in the LSP hover layer, making it reusable across contexts.
+    class SignatureBuilder
+      def initialize(resolver)
+        @resolver = resolver
+      end
+
+      # @param def_node [IR::DefNode] Method definition node
+      # @return [Types::MethodSignature] Structured method signature
+      def build_from_def_node(def_node)
+        params = build_param_signatures(def_node.params)
+        return_type = @resolver.infer(def_node).type
+        Types::MethodSignature.new(params, return_type)
+      end
+
+      private def build_param_signatures(param_nodes)
+        return [] if param_nodes.nil? || param_nodes.empty?
+
+        param_nodes.map { |p| build_param_signature(p) }
+      end
+
+      private def build_param_signature(param_node)
+        type = @resolver.infer(param_node).type
+        Types::ParamSignature.new(
+          name: param_node.name,
+          kind: param_node.kind,
+          type: type
+        )
+      end
+    end
+  end
+end

data/lib/type_guessr/core/type_serializer.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+require_relative "types"
+
+module TypeGuessr
+  module Core
+    # Converts Type objects to/from JSON-compatible Hashes.
+    # Used by GemSignatureCache for persisting inferred method signatures.
+    module TypeSerializer
+      module_function def serialize(type)
+        case type
+        when Types::Unguessed
+          { "_type" => "Unguessed" }
+        when Types::Unknown
+          { "_type" => "Unknown" }
+        when Types::ClassInstance
+          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
+          { "_type" => "TupleType", "element_types" => type.element_types.map { |t| serialize(t) } }
+        when Types::ArrayType
+          { "_type" => "ArrayType", "element_type" => serialize(type.element_type) }
+        when Types::HashShape
+          { "_type" => "HashShape", "fields" => type.fields.transform_keys(&:to_s).transform_values { |v| serialize(v) } }
+        when Types::HashType
+          { "_type" => "HashType", "key_type" => serialize(type.key_type), "value_type" => serialize(type.value_type) }
+        when Types::RangeType
+          { "_type" => "RangeType", "element_type" => serialize(type.element_type) }
+        when Types::Union
+          { "_type" => "Union", "types" => type.types.map { |t| serialize(t) } }
+        when Types::TypeVariable
+          { "_type" => "TypeVariable", "name" => type.name.to_s }
+        when Types::SelfType
+          { "_type" => "SelfType" }
+        when Types::ForwardingArgs
+          { "_type" => "ForwardingArgs" }
+        when Types::MethodSignature
+          { "_type" => "MethodSignature",
+            "return_type" => serialize(type.return_type),
+            "params" => type.params.map { |p| serialize_param(p) } }
+        end
+      end
+
+      # Deserialize a Hash back to a Type object
+      # @param hash [Hash] JSON-compatible hash with "_type" discriminator
+      # @return [Types::Type] The deserialized type
+      # @raise [ArgumentError] if "_type" is unknown
+      module_function def deserialize(hash)
+        case hash["_type"]
+        when "Unguessed" then Types::Unguessed.instance
+        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
+      end
+
+      # @param param [Types::ParamSignature]
+      # @return [Hash]
+      module_function def serialize_param(param)
+        { "name" => param.name.to_s, "kind" => param.kind.to_s, "type" => serialize(param.type) }
+      end
+
+      # @param hash [Hash]
+      # @return [Types::MethodSignature]
+      module_function def deserialize_method_signature(hash)
+        params = hash["params"].map do |p|
+          Types::ParamSignature.new(
+            name: p["name"].to_sym,
+            kind: p["kind"].to_sym,
+            type: deserialize(p["type"])
+          )
+        end
+        Types::MethodSignature.new(params, deserialize(hash["return_type"]))
+      end
+
+      private_class_method :serialize_param, :deserialize_method_signature
+    end
+  end
+end