Resolve type variables of a reopened declaration by position

lib/typeprof/core/ast/sig_decl.rb

@@ -14,16 +14,20 @@ def initialize(raw_decl, lenv)
         @cpath = AST.resolve_rbs_name(raw_decl.name, lenv)
         # TODO: decl.type_params
         # TODO: decl.super_class.args
+        # TODO?: param.variance, param.unchecked, param.upper_bound
+        @params = raw_decl.type_params.map {|param| param.name }
+        sig_type_params = [@cpath, @params]
+        # The header (default types, self types, and superclass arguments) is
+        # resolved in the outer scope but may refer to the type parameters
+        @header_lenv = LocalEnv.new(@lenv.file_context, lenv.cref, {}, [], sig_type_params:)
         ncref = CRef.new(@cpath, :class, nil, lenv.cref)
-        nlenv = LocalEnv.new(@lenv.file_context, ncref, {}, [])
+        nlenv = LocalEnv.new(@lenv.file_context, ncref, {}, [], sig_type_params:)
         @members = raw_decl.members.map do |member|
           AST.create_rbs_member(member, nlenv)
         end.compact
-        # TODO?: param.variance, param.unchecked, param.upper_bound
-        @params = raw_decl.type_params.map {|param| param.name }
         @params_default_types = raw_decl.type_params.map do |param|
           ty = param.default_type
-          ty ? AST.create_rbs_type(ty, lenv) : nil
+          ty ? AST.create_rbs_type(ty, @header_lenv) : nil
         end
       end
 
@@ -81,7 +85,7 @@ def initialize(raw_decl, lenv)
           cpath = name.namespace.path + [self_type.name.name]
           toplevel = name.namespace.absolute?
           @self_types << [cpath, toplevel]
-          @self_type_args << self_type.args.map {|arg| AST.create_rbs_type(arg, lenv) }
+          @self_type_args << self_type.args.map {|arg| AST.create_rbs_type(arg, @header_lenv) }
         end
       end
 
@@ -138,7 +142,7 @@ def initialize(raw_decl, lenv)
           name = superclass.name
           @superclass_cpath = name.namespace.path + [name.name]
           @superclass_toplevel = name.namespace.absolute?
-          @superclass_args = superclass.args.map {|arg| AST.create_rbs_type(arg, lenv) }
+          @superclass_args = superclass.args.map {|arg| AST.create_rbs_type(arg, @header_lenv) }
         else
           @superclass_cpath = nil
           @superclass_toplevel = nil

lib/typeprof/core/ast/sig_type.rb

@@ -883,18 +883,32 @@ def initialize(raw_decl, lenv)
 
       def attrs = { var: }
 
+      # A reopened declaration may rename the type parameters (e.g., `Enumerable[E]`
+      # and `Enumerable[Elem]`), but the subst is keyed by the names of the module
+      # entity, so fall back to matching them by position
+      def resolve_var(genv, subst)
+        decl_cpath, decl_params = @lenv.sig_type_params
+        return subst[@var] unless decl_params
+        idx = decl_params.index(@var)
+        return subst[@var] unless idx
+        subst[genv.resolve_cpath(decl_cpath).type_params.keys[idx]]
+      end
+
       def covariant_vertex0(genv, changes, vtx, subst)
-        raise "unknown type variable: #{ @var }" unless subst[@var]
-        changes.add_edge(genv, subst[@var], vtx)
+        var_vtx = resolve_var(genv, subst)
+        raise "unknown type variable: #{ @var }" unless var_vtx
+        changes.add_edge(genv, var_vtx, vtx)
       end
 
       def contravariant_vertex0(genv, changes, vtx, subst)
-        raise "unknown type variable: #{ @var }" unless subst[@var]
-        changes.add_edge(genv, Source.new(Type::Var.new(genv, @var, subst[@var])), vtx)
+        var_vtx = resolve_var(genv, subst)
+        raise "unknown type variable: #{ @var }" unless var_vtx
+        changes.add_edge(genv, Source.new(Type::Var.new(genv, @var, var_vtx)), vtx)
       end
 
       def typecheck(genv, changes, vtx, subst)
-        changes.add_edge(genv, vtx.new_vertex(genv, self), subst[@var]) unless vtx == subst[@var]
+        var_vtx = resolve_var(genv, subst)
+        changes.add_edge(genv, vtx.new_vertex(genv, self), var_vtx) unless vtx == var_vtx
         true
       end
 

lib/typeprof/core/env.rb

@@ -283,29 +283,13 @@ def resolve_type_alias(cpath, name)
       mod.get_type_alias(name)
     end
 
-    # Returns the type parameter names of the first generic declaration of cpath
-    def find_type_params(decls, cpath)
-      decls.each do |decl|
-        next unless decl.cpath == cpath
-        next unless decl.respond_to?(:params)
-        params = decl.params
-        return params if params && !params.empty?
-      end
-      nil
-    end
-
     def load_core_rbs(raw_decls, position_encoding)
       file_context = FileContext.new(nil, position_encoding)
       lenv = LocalEnv.new(file_context, CRef::Toplevel, {}, [])
       decls = raw_decls.map do |raw_decl|
         AST.create_rbs_decl(raw_decl, lenv)
       end.compact
 
-      # A module entity has one set of parameter names shared by all its declarations,
-      # so the shim must reuse the core's ones (Array's was `Elem`, and is `E` since RBS 4.1)
-      ary_elem, = find_type_params(decls, [:Array]) || [:Elem]
-      hash_key, hash_val = find_type_params(decls, [:Hash]) || [:K, :V]
-
       decls += AST.parse_rbs("typeprof-rbs-shim.rbs", <<-RBS, position_encoding)
         class Exception
           include _Exception
@@ -314,12 +298,12 @@ class String
           include _ToS
           include _ToStr
         end
-        class Array[#{ ary_elem }]
-          include _ToAry[#{ ary_elem }]
-          include _Each[#{ ary_elem }]
+        class Array[Elem]
+          include _ToAry[Elem]
+          include _Each[Elem]
         end
-        class Hash[#{ hash_key }, #{ hash_val }]
-          include _Each[[#{ hash_key }, #{ hash_val }]]
+        class Hash[K, V]
+          include _Each[[K, V]]
         end
         class Object
           include Hash::_Key
@@ -380,7 +364,7 @@ def code_units_cache
   end
 
   class LocalEnv
-    def initialize(file_context, cref, locals, return_boxes, forward_args = nil)
+    def initialize(file_context, cref, locals, return_boxes, forward_args = nil, sig_type_params: nil)
       @file_context = file_context
       @cref = cref
       @locals = locals
@@ -390,9 +374,11 @@ def initialize(file_context, cref, locals, return_boxes, forward_args = nil)
       @ivar_narrowings = {}
       @strict_const_scope = false
       @forward_args = forward_args
+      # [cpath, names] of the type parameters of the enclosing RBS declaration
+      @sig_type_params = sig_type_params
     end
 
-    attr_reader :file_context, :cref, :locals, :return_boxes, :break_vtx, :next_boxes, :strict_const_scope
+    attr_reader :file_context, :cref, :locals, :return_boxes, :break_vtx, :next_boxes, :strict_const_scope, :sig_type_params
     attr_accessor :module_function, :forward_args
 
     def path = @file_context&.path

scenario/rbs/reopen-type-params1.rb

@@ -0,0 +1,28 @@
+## update: test.rbs
+module Foo[E] : _Each[E]
+  def first_e: () -> E
+end
+
+class C
+  include Foo[Integer]
+  def each: () { (Integer) -> void } -> void
+end
+
+## update: test2.rbs
+# Reopen with a different type parameter name
+module Foo[X] : _Each[X]
+  def first_x: () -> X
+  def with_x: (X) -> X
+end
+
+## update: test.rb
+def f = C.new.first_e
+def g = C.new.first_x
+def h = C.new.with_x(1)
+
+## assert
+class Object
+  def f: -> Integer
+  def g: -> Integer
+  def h: -> Integer
+end

scenario/rbs/reopen-type-params2.rb

@@ -0,0 +1,32 @@
+## update: test.rbs
+class A[E]
+  def a: () -> E
+end
+
+class B[U]
+  @x: U
+  def x: () -> U
+end
+
+## update: test2.rbs
+# Reopen with a different type parameter name
+class B[T] < A[T]
+  @y: T
+  def y: () -> T
+end
+
+class Object
+  def b: () -> B[String]
+end
+
+## update: test.rb
+def f = b.a
+def g = b.x
+def h = b.y
+
+## assert
+class Object
+  def f: -> String
+  def g: -> String
+  def h: -> String
+end