Extract per-declaration type param alignment into entry-level align_params

lib/rbs/definition_builder.rb

@@ -124,10 +124,13 @@ def define_instance(definition, type_name, subst, define_class_vars:)
       end
 
       entry = env.class_decls[type_name] or raise "Unknown name for build_instance: #{type_name}"
-      args = entry.type_params.map {|param| Types::Variable.new(name: param.name, location: param.location) }
 
       entry.each_decl do |decl|
-        subst_ = subst + Substitution.build(decl.type_params.each.map(&:name), args)
+        if align_params = entry.align_params(decl)
+          subst_ = subst + align_params
+        else
+          subst_ = subst
+        end
 
         decl.members.each do |member|
           case member

lib/rbs/definition_builder/ancestor_builder.rb

@@ -177,18 +177,11 @@ def validate_super_class!(type_name, entry)
 
         return if with_super_classes.size <= 1
 
-        entry_param_names = entry.type_params.map(&:name)
-
         super_types = with_super_classes.map do |decl|
           super_class = decl.super_class or raise
           args = super_class.args
 
-          decl_param_names = decl.type_params.map(&:name)
-          unless decl_param_names == entry_param_names || args.empty?
-            align_params = Substitution.build(
-              decl_param_names,
-              entry.type_params.map {|param| Types::Variable.new(name: param.name, location: param.location) }
-            )
+          if align_params = entry.align_params(decl)
             args = args.map {|type| type.sub(align_params) }
           end
 
@@ -486,10 +479,7 @@ def mixin_ancestors0(decl, type_name, align_params:, included_modules:, included
 
       def mixin_ancestors(entry, type_name, included_modules:, included_interfaces:, extended_modules:, prepended_modules:, extended_interfaces:)
         entry.each_decl do |decl|
-          align_params = Substitution.build(
-            decl.type_params.each.map(&:name),
-            entry.type_params.map {|param| Types::Variable.new(name: param.name, location: param.location) }
-          )
+          align_params = entry.align_params(decl)
 
           mixin_ancestors0(decl,
                            type_name,

lib/rbs/definition_builder/method_builder.rb

@@ -104,7 +104,7 @@ def build_instance(type_name)
             type = Types::ClassInstance.new(name: type_name, args: args, location: nil)
             Methods.new(type: type).tap do |methods|
               entry.each_decl do |decl|
-                subst = Substitution.build(decl.type_params.each.map(&:name), args)
+                subst = entry.align_params(decl)
                 case decl
                 when AST::Declarations::Base
                   each_rbs_member_with_accessibility(decl.members) do |member, accessibility|
@@ -115,22 +115,22 @@ def build_instance(type_name)
                         build_method(
                           methods,
                           type,
-                          member: member.update(overloads: member.overloads.map {|overload| overload.sub(subst) }),
+                          member: subst ? member.update(overloads: member.overloads.map {|overload| overload.sub(subst) }) : member,
                           accessibility: member.visibility || accessibility
                         )
                       when :singleton_instance
                         build_method(
                           methods,
                           type,
-                          member: member.update(overloads: member.overloads.map {|overload| overload.sub(subst) }),
+                          member: subst ? member.update(overloads: member.overloads.map {|overload| overload.sub(subst) }) : member,
                           accessibility: :private
                         )
                       end
                     when AST::Members::AttrReader, AST::Members::AttrWriter, AST::Members::AttrAccessor
                       if member.kind == :instance
                         build_attribute(methods,
                                         type,
-                                        member: member.update(type: member.type.sub(subst)),
+                                        member: subst ? member.update(type: member.type.sub(subst)) : member,
                                         accessibility: member.visibility || accessibility)
                       end
                     when AST::Members::Alias

lib/rbs/environment/class_entry.rb

@@ -64,6 +64,18 @@ def validate_type_params
           end
         end
       end
+
+      def align_params(decl)
+        entry_params = type_params
+        decl_param_names = decl.type_params.map(&:name)
+
+        return nil if decl_param_names == entry_params.map(&:name)
+
+        Substitution.build(
+          decl_param_names,
+          entry_params.map {|param| Types::Variable.new(name: param.name, location: param.location) }
+        )
+      end
     end
   end
 end

lib/rbs/environment/module_entry.rb

@@ -41,23 +41,13 @@ def type_params
       end
 
       def self_types
-        params = type_params
-        param_names = params.map(&:name)
-
         each_decl.flat_map do |decl|
           self_types = decl.self_types
-          decl_param_names = decl.type_params.map(&:name)
+          subst = align_params(decl)
 
-          if self_types.empty? || decl_param_names == param_names
+          if self_types.empty? || subst.nil?
             self_types
           else
-            # The declaration uses different type parameter names from the primary declaration.
-            # Rename the type variables in the self types, so that they are aligned to `#type_params`.
-            subst = Substitution.build(
-              decl_param_names,
-              params.map {|param| Types::Variable.new(name: param.name, location: param.location) }
-            )
-
             self_types.map do |self_type|
               AST::Declarations::Module::Self.new(
                 name: self_type.name,
@@ -69,6 +59,18 @@ def self_types
         end.uniq
       end
 
+      def align_params(decl)
+        entry_params = type_params
+        decl_param_names = decl.type_params.map(&:name)
+
+        return nil if decl_param_names == entry_params.map(&:name)
+
+        Substitution.build(
+          decl_param_names,
+          entry_params.map {|param| Types::Variable.new(name: param.name, location: param.location) }
+        )
+      end
+
       def validate_type_params
         unless context_decls.empty?
           first_decl, *rest_decls = each_decl.to_a

sig/environment/class_entry.rbs

@@ -45,6 +45,12 @@ module RBS
       # * Raises `GenericParameterMismatchError` if incompatible declaration is detected.
       #
       def validate_type_params: () -> void
+
+      # Returns a substitution that renames the type parameters of the declaration to the entry's type parameters (`#type_params`)
+      #
+      # Returns `nil` if the declaration uses the same type parameter names as `#type_params`.
+      #
+      def align_params: (declaration | ModuleEntry::declaration) -> Substitution?
     end
   end
 end

sig/environment/module_entry.rbs

@@ -54,6 +54,12 @@ module RBS
       # declarations, but `#location` points to the original declaration.
       #
       def self_types: () -> Array[AST::Declarations::Module::Self]
+
+      # Returns a substitution that renames the type parameters of the declaration to the entry's type parameters (`#type_params`)
+      #
+      # Returns `nil` if the declaration uses the same type parameter names as `#type_params`.
+      #
+      def align_params: (declaration | ClassEntry::declaration) -> Substitution?
     end
   end
 end

test/rbs/environment_test.rb

@@ -361,6 +361,43 @@ module Foo[C] : _Animal[Integer]
     end
   end
 
+  def test_module_entry_align_params
+    _, _, decls = RBS::Parser.parse_signature(<<EOF)
+module Foo[A, B]
+end
+
+module Foo[X, Y]
+end
+
+module Foo[X]
+end
+EOF
+
+    Environment::ModuleEntry.new(type_name("::Foo")).tap do |entry|
+      entry << [nil, decls[0]]
+      entry << [nil, decls[1]]
+
+      # Aligned to the primary declaration's names, so no substitution is needed
+      assert_nil entry.align_params(decls[0])
+
+      entry.align_params(decls[1]).tap do |subst|
+        subst or raise
+        assert_equal RBS::Types::Variable.new(name: :A, location: nil), subst[RBS::Types::Variable.new(name: :X, location: nil)]
+        assert_equal RBS::Types::Variable.new(name: :B, location: nil), subst[RBS::Types::Variable.new(name: :Y, location: nil)]
+      end
+    end
+
+    Environment::ModuleEntry.new(type_name("::Foo")).tap do |entry|
+      entry << [nil, decls[0]]
+      entry << [nil, decls[2]]
+
+      # The type params validation runs before the alignment, so the arity mismatch is detected first
+      assert_raises RBS::GenericParameterMismatchError do
+        entry.align_params(decls[2])
+      end
+    end
+  end
+
   def test_absolute_type
     env = Environment.new