[Type checker] Allow extensions of typealiases naming generic specializations

lib/Sema/TypeCheckDecl.cpp

@@ -4476,26 +4476,21 @@ static bool isPassThroughTypealias(TypeAliasDecl *typealias) {
 
 /// Form the interface type of an extension from the raw type and the
 /// extension's list of generic parameters.
-static Type formExtensionInterfaceType(TypeChecker &tc, ExtensionDecl *ext,
-                                       Type type,
-                                       GenericParamList *genericParams,
-                                       bool &mustInferRequirements) {
+static Type formExtensionInterfaceType(
+                         TypeChecker &tc, ExtensionDecl *ext,
+                         Type type,
+                         GenericParamList *genericParams,
+                         SmallVectorImpl<std::pair<Type, Type>> &sameTypeReqs,
+                         bool &mustInferRequirements) {
   if (type->is<ErrorType>())
     return type;
 
   // Find the nominal type declaration and its parent type.
-  Type parentType;
-  GenericTypeDecl *genericDecl;
-  if (auto unbound = type->getAs<UnboundGenericType>()) {
-    parentType = unbound->getParent();
-    genericDecl = unbound->getDecl();
-  } else {
-    if (type->is<ProtocolCompositionType>())
-      type = type->getCanonicalType();
-    auto nominalType = type->castTo<NominalType>();
-    parentType = nominalType->getParent();
-    genericDecl = nominalType->getDecl();
-  }
+  if (type->is<ProtocolCompositionType>())
+    type = type->getCanonicalType();
+
+  Type parentType = type->getNominalParent();
+  GenericTypeDecl *genericDecl = type->getAnyGeneric();
 
   // Reconstruct the parent, if there is one.
   if (parentType) {
@@ -4505,7 +4500,7 @@ static Type formExtensionInterfaceType(TypeChecker &tc, ExtensionDecl *ext,
                                  : genericParams;
     parentType =
       formExtensionInterfaceType(tc, ext, parentType, parentGenericParams,
-                                 mustInferRequirements);
+                                 sameTypeReqs, mustInferRequirements);
   }
 
   // Find the nominal type.
@@ -4523,9 +4518,20 @@ static Type formExtensionInterfaceType(TypeChecker &tc, ExtensionDecl *ext,
     resultType = NominalType::get(nominal, parentType,
                                   nominal->getASTContext());
   } else {
+    auto currentBoundType = type->getAs<BoundGenericType>();
+
     // Form the bound generic type with the type parameters provided.
+    unsigned gpIndex = 0;
     for (auto gp : *genericParams) {
-      genericArgs.push_back(gp->getDeclaredInterfaceType());
+      SWIFT_DEFER { ++gpIndex; };
+
+      auto gpType = gp->getDeclaredInterfaceType();
+      genericArgs.push_back(gpType);
+
+      if (currentBoundType) {
+        sameTypeReqs.push_back({gpType,
+                                currentBoundType->getGenericArgs()[gpIndex]});
+      }
     }
 
     resultType = BoundGenericType::get(nominal, parentType, genericArgs);
@@ -4562,8 +4568,9 @@ checkExtensionGenericParams(TypeChecker &tc, ExtensionDecl *ext, Type type,
 
   // Form the interface type of the extension.
   bool mustInferRequirements = false;
+  SmallVector<std::pair<Type, Type>, 4> sameTypeReqs;
   Type extInterfaceType =
-    formExtensionInterfaceType(tc, ext, type, genericParams,
+    formExtensionInterfaceType(tc, ext, type, genericParams, sameTypeReqs,
                                mustInferRequirements);
 
   // Local function used to infer requirements from the extended type.
@@ -4575,18 +4582,34 @@ checkExtensionGenericParams(TypeChecker &tc, ExtensionDecl *ext, Type type,
                               extInterfaceType,
                               nullptr,
                               source);
+
+    for (const auto &sameTypeReq : sameTypeReqs) {
+      builder.addRequirement(
+        Requirement(RequirementKind::SameType, sameTypeReq.first,
+                    sameTypeReq.second),
+        source, ext->getModuleContext());
+    }
   };
 
   // Validate the generic type signature.
   auto *env = tc.checkGenericEnvironment(genericParams,
                                          ext->getDeclContext(), nullptr,
                                          /*allowConcreteGenericParams=*/true,
                                          ext, inferExtendedTypeReqs,
-                                         mustInferRequirements);
+                                         (mustInferRequirements ||
+                                            !sameTypeReqs.empty()));
 
   return { env, extInterfaceType };
 }
 
+static bool isNonGenericTypeAliasType(Type type) {
+  // A non-generic typealias can extend a specialized type.
+  if (auto *aliasType = dyn_cast<NameAliasType>(type.getPointer()))
+    return aliasType->getDecl()->getGenericContextDepth() == (unsigned)-1;
+
+  return false;
+}
+
 static void validateExtendedType(ExtensionDecl *ext, TypeChecker &tc) {
   // If we didn't parse a type, fill in an error type and bail out.
   if (!ext->getExtendedTypeLoc().getTypeRepr()) {
@@ -4630,20 +4653,22 @@ static void validateExtendedType(ExtensionDecl *ext, TypeChecker &tc) {
     return;
   }
 
-  // Cannot extend a bound generic type.
-  if (extendedType->isSpecialized()) {
-    tc.diagnose(ext->getLoc(), diag::extension_specialization,
-             extendedType->getAnyNominal()->getName())
+  // Cannot extend function types, tuple types, etc.
+  if (!extendedType->getAnyNominal()) {
+    tc.diagnose(ext->getLoc(), diag::non_nominal_extension, extendedType)
       .highlight(ext->getExtendedTypeLoc().getSourceRange());
     ext->setInvalid();
     ext->getExtendedTypeLoc().setInvalidType(tc.Context);
     return;
   }
 
-  // Cannot extend function types, tuple types, etc.
-  if (!extendedType->getAnyNominal()) {
-    tc.diagnose(ext->getLoc(), diag::non_nominal_extension, extendedType)
-      .highlight(ext->getExtendedTypeLoc().getSourceRange());
+  // Cannot extend a bound generic type, unless it's referenced via a
+  // non-generic typealias type.
+  if (extendedType->isSpecialized() &&
+      !isNonGenericTypeAliasType(extendedType)) {
+    tc.diagnose(ext->getLoc(), diag::extension_specialization,
+                extendedType->getAnyNominal()->getName())
+    .highlight(ext->getExtendedTypeLoc().getSourceRange());
     ext->setInvalid();
     ext->getExtendedTypeLoc().setInvalidType(tc.Context);
     return;

test/decl/ext/generic.swift

@@ -30,7 +30,7 @@ extension X<Int, Double, String> {
 
 typealias GGG = X<Int, Double, String>
 
-extension GGG { } // expected-error{{constrained extension must be declared on the unspecialized generic type 'X' with constraints specified by a 'where' clause}}
+extension GGG { } // okay through a typealias
 
 // Lvalue check when the archetypes are not the same.
 struct LValueCheck<T> {
@@ -222,4 +222,4 @@ extension NewGeneric {
   static func newMember() -> NewGeneric {
     return NewGeneric()
   }
-}
+}

test/decl/ext/typealias.swift

@@ -0,0 +1,80 @@
+// RUN: %target-typecheck-verify-swift
+
+struct Foo<T> {
+  var maybeT: T? { return nil }
+}
+
+extension Foo {
+  struct Bar<U, V> {
+    var maybeT: T? { return nil }
+    var maybeU: U? { return nil }
+    var maybeV: V? { return nil }
+
+    struct Inner {
+      var maybeT: T? { return nil }
+      var maybeU: U? { return nil }
+      var maybeV: V? { return nil }
+    }
+  }
+}
+
+typealias FooInt = Foo<Int>
+
+extension FooInt {
+  func goodT() -> Int {
+    return maybeT!
+  }
+
+  func badT() -> Float {
+    return maybeT! // expected-error{{cannot convert return expression of type 'Int' to return type 'Float'}}
+  }
+}
+
+typealias FooIntBarFloatDouble = Foo<Int>.Bar<Float, Double>
+
+extension FooIntBarFloatDouble {
+  func goodT() -> Int {
+    return maybeT!
+  }
+  func goodU() -> Float {
+    return maybeU!
+  }
+  func goodV() -> Double {
+    return maybeV!
+  }
+
+  func badT() -> Float {
+    return maybeT! // expected-error{{cannot convert return expression of type 'Int' to return type 'Float'}}
+  }
+  func badU() -> Int {
+    return maybeU! // expected-error{{cannot convert return expression of type 'Float' to return type 'Int'}}
+  }
+  func badV() -> Int {
+    return maybeV! // expected-error{{cannot convert return expression of type 'Double' to return type 'Int'}}
+  }
+}
+
+typealias FooIntBarFloatDoubleInner = Foo<Int>.Bar<Float, Double>.Inner
+
+extension FooIntBarFloatDoubleInner {
+  func goodT() -> Int {
+    return maybeT!
+  }
+  func goodU() -> Float {
+    return maybeU!
+  }
+  func goodV() -> Double {
+    return maybeV!
+  }
+
+  func badT() -> Float {
+    return maybeT! // expected-error{{cannot convert return expression of type 'Int' to return type 'Float'}}
+  }
+  func badU() -> Int {
+    return maybeU! // expected-error{{cannot convert return expression of type 'Float' to return type 'Int'}}
+  }
+  func badV() -> Int {
+    return maybeV! // expected-error{{cannot convert return expression of type 'Double' to return type 'Int'}}
+  }
+}
+