Resolve rbs type aliases and interfaces from their definitions

lib/repl_type_completor/type_analyzer.rb

@@ -1152,7 +1152,7 @@ def method_call(receiver, method_name, args, kwargs, block, scope, name_match: t
       methods = Types.rbs_methods receiver, method_name.to_sym, args, kwargs, !!block
       block_called = false
       type_breaks = methods.map do |method, given_params, method_params|
-        receiver_vars = receiver.is_a?(Types::InstanceType) ? receiver.named_params : {}
+        receiver_vars = receiver.is_a?(Types::InstanceType) || receiver.is_a?(Types::InterfaceType) ? receiver.named_params : {}
         free_vars = method.type.free_variables - receiver_vars.keys.to_set
         vars = receiver_vars.merge Types.match_free_variables(free_vars, method_params, given_params)
         if block && method.block && method.block.type.respond_to?(:required_positionals)

lib/repl_type_completor/types.rb

@@ -8,6 +8,7 @@
 module ReplTypeCompletor
   module Types
     OBJECT_TO_TYPE_SAMPLE_SIZE = 50
+    EXPANSION_NESTING_LIMIT = 3
 
     singleton_class.attr_reader :rbs_builder, :rbs_load_error
 
@@ -75,6 +76,11 @@ def self.rbs_absolute_type_name(name)
     def self.rbs_search_method(klass, method_name, singleton)
       return unless rbs_builder
 
+      if klass.is_a? InterfaceType
+        definition = klass.definition
+        return definition && definition.methods[method_name]
+      end
+
       klass.ancestors.each do |ancestor|
         next unless (name = Methods::MODULE_NAME_METHOD.bind_call(ancestor))
 
@@ -93,6 +99,8 @@ def self.method_return_type(type, method_name)
           [t, t.module_or_class, true]
         in InstanceType
           [t, t.klass, false]
+        in InterfaceType
+          [t, t, false]
         end
       end
       types = receivers.flat_map do |receiver_type, klass, singleton|
@@ -115,6 +123,8 @@ def self.accessor_method_return_type(type, method_name)
           t.module_or_class
         in InstanceType
           t.instances
+        in InterfaceType
+          nil
         end
       end.flatten
       instances = instances.sample(OBJECT_TO_TYPE_SAMPLE_SIZE) if instances.size > OBJECT_TO_TYPE_SAMPLE_SIZE
@@ -136,6 +146,8 @@ def self.rbs_methods(type, method_name, args_types, kwargs_type, has_block)
           [t, t.module_or_class, true]
         in InstanceType
           [t, t.klass, false]
+        in InterfaceType
+          [t, t, false]
         end
       end
       has_splat = args_types.include?(nil)
@@ -193,7 +205,14 @@ def self.intersect?(a, b)
       end
 
       aa, bb = [atypes, btypes].map {|types| (types[InstanceType] || []).map(&:klass) }
-      (aa.flat_map(&:ancestors) & bb).any?
+      return true if (aa.flat_map(&:ancestors) & bb).any?
+
+      [[atypes[InterfaceType], b], [btypes[InterfaceType], a]].any? do |interfaces, other|
+        interfaces&.any? do |interface|
+          methods = interface.methods
+          !methods.empty? && other.types.any? { |t| (methods - t.methods).empty? }
+        end
+      end
     end
 
     def self.type_from_object(object)
@@ -340,6 +359,38 @@ def inspect_without_params
       end
     end
 
+    class InterfaceType
+      attr_reader :name, :args
+
+      def initialize(name, args = [])
+        @name = name
+        @args = args
+      end
+
+      def definition
+        return @definition if defined?(@definition)
+
+        @definition = (Types.rbs_builder&.build_interface(@name) rescue nil)
+      end
+
+      def transform() = yield(self)
+      def methods() = definition ? definition.methods.keys : []
+      def all_methods() = methods
+
+      def named_params
+        definition ? definition.type_params.zip(@args).to_h.compact : {}
+      end
+
+      def constants() = []
+      def types() = [self]
+      def nillable?() = false
+      def nonnillable() = self
+
+      def inspect
+        args.empty? ? name.name.to_s : "#{name.name}[#{args.map(&:inspect).join(', ')}]"
+      end
+    end
+
     NIL = InstanceType.new NilClass
     OBJECT = InstanceType.new Object
     TRUE = InstanceType.new TrueClass
@@ -364,6 +415,7 @@ class UnionType
       def initialize(*types)
         @types = []
         singleton_types = []
+        interface_types = {}
         instance_types = {}
         collect = -> type do
           case type
@@ -377,10 +429,12 @@ def initialize(*types)
             end
           in SingletonType
             singleton_types << type
+          in InterfaceType
+            interface_types[[type.name, type.args]] ||= type
           end
         end
         types.each(&collect)
-        @types = singleton_types.uniq + instance_types.map do |klass, (params, instances)|
+        @types = singleton_types.uniq + interface_types.values + instance_types.map do |klass, (params, instances)|
           params = params.map { |v| UnionType[*v] }
           InstanceType.new(klass, params, instances)
         end
@@ -422,7 +476,15 @@ def self.array_of(*types)
       InstanceType.new(Array, [type])
     end
 
-    def self.from_rbs_type(return_type, self_type, extra_vars = {})
+    # Alias may validly recurse through type args (`type json = Integer | Array[json]`)
+    # and can expand exponentially (`type b = a | [a]` chains), so expansion is depth limited.
+    def self.expand_alias_type(rbs_type, nesting)
+      return if nesting >= EXPANSION_NESTING_LIMIT
+
+      rbs_builder.expand_alias2 rbs_type.name, rbs_type.args rescue nil
+    end
+
+    def self.from_rbs_type(return_type, self_type, extra_vars = {}, nesting = 0)
       case return_type
       when RBS::Types::Bases::Self
         self_type
@@ -434,12 +496,13 @@ def self.from_rbs_type(return_type, self_type, extra_vars = {})
         self_type.transform do |type|
           case type
           in SingletonType
-            InstanceType.new(self_type.module_or_class.is_a?(Class) ? Class : Module)
+            InstanceType.new(type.module_or_class.is_a?(Class) ? Class : Module)
           in InstanceType
             SingletonType.new type.klass
+          in InterfaceType
+            CLASS
           end
         end
-        UnionType[*types]
       when RBS::Types::Bases::Bool
         BOOLEAN
       when RBS::Types::Bases::Instance
@@ -454,11 +517,11 @@ def self.from_rbs_type(return_type, self_type, extra_vars = {})
         end
       when RBS::Types::Union, RBS::Types::Intersection
         # Intersection is unsupported. fallback to union type
-        UnionType[*return_type.types.map { from_rbs_type _1, self_type, extra_vars }]
+        UnionType[*return_type.types.map { from_rbs_type _1, self_type, extra_vars, nesting }]
       when RBS::Types::Proc
         PROC
       when RBS::Types::Tuple
-        elem = UnionType[*return_type.types.map { from_rbs_type _1, self_type, extra_vars }]
+        elem = UnionType[*return_type.types.map { from_rbs_type _1, self_type, extra_vars, nesting }]
         InstanceType.new(Array, [elem])
       when RBS::Types::Record
         InstanceType.new(Hash, [SYMBOL, OBJECT])
@@ -467,37 +530,28 @@ def self.from_rbs_type(return_type, self_type, extra_vars = {})
       when RBS::Types::Variable
         if extra_vars.key? return_type.name
           extra_vars[return_type.name]
-        elsif self_type.is_a? InstanceType
+        elsif self_type.is_a?(InstanceType) || self_type.is_a?(InterfaceType)
           self_type.named_params[return_type.name] || OBJECT
         elsif self_type.is_a? UnionType
           types = self_type.types.filter_map do |t|
-            t.named_params[return_type.name] if t.is_a? InstanceType
+            t.named_params[return_type.name] if t.is_a?(InstanceType) || t.is_a?(InterfaceType)
           end
           UnionType[*types]
         else
           OBJECT
         end
       when RBS::Types::Optional
-        UnionType[from_rbs_type(return_type.type, self_type, extra_vars), NIL]
+        UnionType[from_rbs_type(return_type.type, self_type, extra_vars, nesting), NIL]
       when RBS::Types::Alias
-        case return_type.name.name
-        when :int
-          INTEGER
-        when :boolish
-          BOOLEAN
-        when :string
-          STRING
-        else
-          # TODO: ???
-          OBJECT
-        end
+        expanded = expand_alias_type(return_type, nesting)
+        expanded ? from_rbs_type(expanded, self_type, extra_vars, nesting + 1) : OBJECT
       when RBS::Types::Interface
-        # unimplemented
-        OBJECT
+        args = return_type.args.map { from_rbs_type _1, self_type, extra_vars, nesting }
+        InterfaceType.new return_type.name, args
       when RBS::Types::ClassInstance
         klass = return_type.name.to_namespace.path.reduce(Object) { _1.const_get _2 }
         if return_type.args
-          params = return_type.args.map { from_rbs_type _1, self_type, extra_vars }
+          params = return_type.args.map { from_rbs_type _1, self_type, extra_vars, nesting }
         end
         InstanceType.new(klass, params || [])
       else
@@ -517,24 +571,24 @@ def self.match_free_variables(vars, types, values)
       accumulator.transform_values { UnionType[*_1] }
     end
 
-    def self._match_free_variable(vars, rbs_type, value, accumulator)
+    def self._match_free_variable(vars, rbs_type, value, accumulator, nesting = 0)
       case [rbs_type, value]
       in [RBS::Types::Variable,]
         (accumulator[rbs_type.name] ||= []) << value if vars.include? rbs_type.name
       in [RBS::Types::ClassInstance, InstanceType]
         names = rbs_builder.build_singleton(rbs_type.name).type_params
         names.zip(rbs_type.args).each do |name, arg|
           v = value.named_params[name]
-          _match_free_variable vars, arg, v, accumulator if v
+          _match_free_variable vars, arg, v, accumulator, nesting if v
         end
       in [RBS::Types::Tuple, InstanceType] if value.klass == Array
         v = value.params[0]
         rbs_type.types.each do |t|
-          _match_free_variable vars, t, v, accumulator
+          _match_free_variable vars, t, v, accumulator, nesting
         end
       in [RBS::Types::Record, InstanceType] if value.klass == Hash
         # TODO
-      in [RBS::Types::Interface,]
+      in [RBS::Types::Interface,] if nesting < EXPANSION_NESTING_LIMIT
         definition = rbs_builder.build_interface rbs_type.name
         convert = {}
         definition.type_params.zip(rbs_type.args).each do |from, arg|
@@ -546,13 +600,20 @@ def self._match_free_variable(vars, rbs_type, value, accumulator)
           return_type = method_return_type value, method_name
           method.defs.each do |method_def|
             interface_return_type = method_def.type.type.return_type
-            _match_free_variable convert, interface_return_type, return_type, ac
+            _match_free_variable convert, interface_return_type, return_type, ac, nesting + 1
           end
         end
         convert.each do |from, to|
           values = ac[from]
           (accumulator[to] ||= []).concat values if values
         end
+      in [RBS::Types::Union,]
+        rbs_type.types.each do |t|
+          _match_free_variable vars, t, value, accumulator, nesting
+        end
+      in [RBS::Types::Alias,]
+        expanded = expand_alias_type(rbs_type, nesting)
+        _match_free_variable vars, expanded, value, accumulator, nesting + 1 if expanded
       else
       end
     end

test/repl_type_completor/test_type_analyze.rb

@@ -774,6 +774,8 @@ def test_block_args
     def test_array_aref
       assert_call('[1][0..].', include: [Array, NilClass], exclude: Integer)
       assert_call('[1][0].', include: Integer, exclude: [Array, NilClass])
+      # Float matches the `int` index param by duck typing (Float#to_int)
+      assert_call('[1][0.5].', include: Integer, exclude: Array)
       assert_call('[1].[](0).', include: Integer, exclude: [Array, NilClass])
       assert_call('[1].[](0){}.', include: Integer, exclude: [Array, NilClass])
     end

test/repl_type_completor/test_types.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require 'repl_type_completor'
+require 'tmpdir'
 require_relative './helper'
 
 module TestReplTypeCompletor
@@ -94,6 +95,104 @@ def foobar; end
       assert_include type.all_methods, :rand
     end
 
+    def type_from_rbs(rbs_string)
+      ReplTypeCompletor::Types.load_rbs_builder unless ReplTypeCompletor::Types.rbs_builder
+      rbs_type = RBS::Parser.parse_type(rbs_string)
+      ReplTypeCompletor::Types.from_rbs_type(rbs_type, ReplTypeCompletor::Types::OBJECT)
+    end
+
+    def test_interface_type
+      to_int = type_from_rbs('::_ToInt')
+      assert_equal '_ToInt', to_int.inspect
+      assert_equal [:to_int], to_int.methods
+      assert ReplTypeCompletor::Types.intersect?(ReplTypeCompletor::Types::FLOAT, to_int)
+      assert ReplTypeCompletor::Types.intersect?(to_int, ReplTypeCompletor::Types::FLOAT)
+      refute ReplTypeCompletor::Types.intersect?(ReplTypeCompletor::Types::STRING, to_int)
+
+      to_ary = type_from_rbs('::_ToAry[::Integer]')
+      assert_equal '_ToAry[Integer]', to_ary.inspect
+      return_type = ReplTypeCompletor::Types.method_return_type(to_ary, :to_ary)
+      assert_equal Array, return_type.klass
+      assert_equal Integer, return_type.params[0].klass
+    end
+
+    def test_alias_type_expansion
+      int_type = type_from_rbs('::int')
+      assert_equal 'Integer | _ToInt', int_type.inspect
+    end
+
+    def with_isolated_rbs_env(rbs_source)
+      Dir.mktmpdir do |dir|
+        File.write File.join(dir, 'test.rbs'), rbs_source
+        loader = RBS::EnvironmentLoader.new core_root: nil
+        loader.add path: Pathname(dir)
+        env = RBS::Environment.from_loader(loader)
+        builder = RBS::DefinitionBuilder.new env: env.resolve_type_names
+        original_builder = ReplTypeCompletor::Types.rbs_builder
+        begin
+          ReplTypeCompletor::Types.instance_variable_set :@rbs_builder, builder
+          yield
+        ensure
+          ReplTypeCompletor::Types.instance_variable_set :@rbs_builder, original_builder
+        end
+      end
+    end
+
+    def test_recursive_alias_type_expansion
+      rbs_source = <<~RBS
+        type json = Integer | Array[json] | Hash[String, json]
+        type unguarded = unguarded | Integer
+        type opt = [opt]? | Integer
+        type a = Integer | [a]
+        type b = a | [a]
+        type c = b | [b]
+        type d = c | [c]
+        type e = d | [d]
+        interface _Generic[T]
+          def get: () -> T
+        end
+        type generic_rec = _Generic[generic_rec] | Integer
+      RBS
+      with_isolated_rbs_env rbs_source do
+        json_type = type_from_rbs('::json')
+        assert_equal [Array, Hash, Integer], json_type.types.map(&:klass).sort_by(&:name)
+        json_elem = json_type.types.find { _1.klass == Array }.params[0]
+        assert_include json_elem.types.map(&:klass), Integer
+
+        # Invalid in RBS (RecursiveTypeAliasError by `rbs validate`) but loadable
+        assert_include type_from_rbs('::unguarded').types.map(&:klass), Integer
+
+        # Recursion through optional and interface type args
+        assert_include type_from_rbs('::opt').types.map(&:klass), Integer
+        generic_rec_type = type_from_rbs('::generic_rec')
+        assert_include generic_rec_type.types.grep(ReplTypeCompletor::Types::InstanceType).map(&:klass), Integer
+
+        # Exponentially expanding alias chain
+        assert_include type_from_rbs('::e').types.map(&:klass), Array
+      end
+    end
+
+    def test_cyclic_generic_interface_match
+      rbs_source = <<~RBS
+        interface _CycA[T]
+          def a: () -> _CycB[T]
+        end
+        interface _CycB[T]
+          def b: () -> _CycA[T]
+        end
+      RBS
+      with_isolated_rbs_env rbs_source do
+        var = RBS::Types::Variable.new(name: :X, location: nil)
+        cyclic = RBS::Types::Interface.new(
+          name: ReplTypeCompletor::Types.rbs_absolute_type_name('_CycA'),
+          args: [var],
+          location: nil
+        )
+        matched = ReplTypeCompletor::Types.match_free_variables([:X], [cyclic], [ReplTypeCompletor::Types::INTEGER])
+        assert_kind_of Hash, matched
+      end
+    end
+
     def test_basic_object_methods
       bo = BasicObject.new
       def bo.foobar; end