AST: Clone generic parameter list of tuple extension from extended typealias

slavapestov · slavapestov · commit d5cdfb2cf2a7 · 2023-09-05T23:21:27.000-04:00
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -5178,21 +5178,19 @@ class ProtocolDecl final : public NominalTypeDecl {
 /// and conformances for tuple types.
 ///
 /// - The declared interface type is the special TheTupleType singleton.
-/// - The generic parameter list has one pack generic parameter, <Elements...>
-/// - The generic signature has no requirements, <Elements...>
+/// - The generic parameter list has one pack generic parameter, <each Element>
+/// - The generic signature has no requirements, <each Element>
 /// - The self interface type is the tuple type containing a single pack
-///   expansion, (Elements...).
+///   expansion, (repeat each Element).
 class BuiltinTupleDecl final : public NominalTypeDecl {
-  TupleType *TupleSelfType = nullptr;
-
 public:
   BuiltinTupleDecl(Identifier Name, DeclContext *Parent);
 
   SourceRange getSourceRange() const {
     return SourceRange();
   }
 
-  TupleType *getTupleSelfType() const;
+  TupleType *getTupleSelfType(const ExtensionDecl *owner) const;
 
   // Implement isa/cast/dyncast/etc.
   static bool classof(const Decl *D) {
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -2353,6 +2353,8 @@ ERROR(extension_access_with_conformances,none,
 ERROR(experimental_tuple_extension,none,
       "tuple extensions are experimental",
       ())
+ERROR(tuple_extension_wrong_type,none,
+      "tuple extension must be written as extension of %0", (Type))
 ERROR(extension_metatype,none,
       "cannot extend a metatype %0", (Type))
 ERROR(extension_placeholder,none,
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -7861,27 +7861,24 @@ Type DeclContext::getSelfTypeInContext() const {
   return mapTypeIntoContext(getSelfInterfaceType());
 }
 
-TupleType *BuiltinTupleDecl::getTupleSelfType() const {
-  if (TupleSelfType)
-    return TupleSelfType;
-
+TupleType *BuiltinTupleDecl::getTupleSelfType(const ExtensionDecl *owner) const {
   auto &ctx = getASTContext();
 
-  // Get the generic parameter type 'Elements'.
-  auto paramType = getGenericParams()->getParams()[0]
-      ->getDeclaredInterfaceType();
+  // Get the generic parameter type 'each T'.
+  auto *genericParams = owner->getGenericParams();
+  assert(genericParams != nullptr);
+  assert(genericParams->getParams().size() == 1);
+  assert(genericParams->getOuterParameters() == nullptr);
+  auto paramType = genericParams->getParams()[0]->getDeclaredInterfaceType();
 
-  // Build the pack expansion type 'Elements...'.
+  // Build the pack expansion type 'repeat each T'.
   Type packExpansionType = PackExpansionType::get(paramType, paramType);
 
-  // Build the one-element tuple type '(Elements...)'.
+  // Build the one-element tuple type '(repeat each T)'.
   SmallVector<TupleTypeElt, 1> elts;
   elts.push_back(packExpansionType);
 
-  const_cast<BuiltinTupleDecl *>(this)->TupleSelfType =
-      TupleType::get(elts, ctx);
-
-  return TupleSelfType;
+  return TupleType::get(elts, ctx);
 }
 
 /// Retrieve the interface type of 'self' for the given context.
@@ -7890,7 +7887,7 @@ Type DeclContext::getSelfInterfaceType() const {
 
   if (auto *nominalDecl = getSelfNominalTypeDecl()) {
     if (auto *builtinTupleDecl = dyn_cast<BuiltinTupleDecl>(nominalDecl))
-      return builtinTupleDecl->getTupleSelfType();
+      return builtinTupleDecl->getTupleSelfType(cast<ExtensionDecl>(this));
 
     if (isa<ProtocolDecl>(nominalDecl)) {
       auto *genericParams = nominalDecl->getGenericParams();
diff --git a/lib/AST/NameLookup.cpp b/lib/AST/NameLookup.cpp
@@ -3279,10 +3279,10 @@ bool TypeRepr::isProtocolOrProtocolComposition(DeclContext *dc){
 static GenericParamList *
 createExtensionGenericParams(ASTContext &ctx,
                              ExtensionDecl *ext,
-                             NominalTypeDecl *nominal) {
+                             DeclContext *source) {
   // Collect generic parameters from all outer contexts.
   SmallVector<GenericParamList *, 2> allGenericParams;
-  nominal->forEachGenericContext([&](GenericParamList *gpList) {
+  source->forEachGenericContext([&](GenericParamList *gpList) {
     allGenericParams.push_back(gpList->clone(ext));
   });
 
@@ -3295,6 +3295,28 @@ createExtensionGenericParams(ASTContext &ctx,
   return toParams;
 }
 
+/// If the extended type is a generic typealias whose underlying type is
+/// a tuple, the extension inherits the generic paramter list from the
+/// typealias.
+static GenericParamList *
+createTupleExtensionGenericParams(ASTContext &ctx,
+                                  ExtensionDecl *ext,
+                                  TypeRepr *extendedTypeRepr) {
+  DirectlyReferencedTypeDecls referenced =
+    directReferencesForTypeRepr(ctx.evaluator, ctx,
+                                extendedTypeRepr,
+                                ext->getParent());
+
+  if (referenced.size() != 1 || !isa<TypeAliasDecl>(referenced[0]))
+    return nullptr;
+
+  auto *typeAlias = cast<TypeAliasDecl>(referenced[0]);
+  if (!typeAlias->isGeneric())
+    return nullptr;
+
+  return createExtensionGenericParams(ctx, ext, typeAlias);
+}
+
 CollectedOpaqueReprs swift::collectOpaqueTypeReprs(TypeRepr *r, ASTContext &ctx,
                                                    DeclContext *d) {
   class Walker : public ASTWalker {
@@ -3416,9 +3438,9 @@ GenericParamListRequest::evaluate(Evaluator &evaluator, GenericContext *value) c
   if (auto *tupleDecl = dyn_cast<BuiltinTupleDecl>(value)) {
     auto &ctx = value->getASTContext();
 
-    // Builtin.TheTupleType has a single pack generic parameter: <Elements...>
+    // Builtin.TheTupleType has a single pack generic parameter: <each Element>
     auto *genericParam = GenericTypeParamDecl::createImplicit(
-        tupleDecl->getDeclContext(), ctx.Id_Elements, /*depth*/ 0, /*index*/ 0,
+        tupleDecl->getDeclContext(), ctx.Id_Element, /*depth*/ 0, /*index*/ 0,
         /*isParameterPack*/ true);
 
     return GenericParamList::create(ctx, SourceLoc(), genericParam,
@@ -3433,6 +3455,21 @@ GenericParamListRequest::evaluate(Evaluator &evaluator, GenericContext *value) c
     if (!nominal) {
       return nullptr;
     }
+
+    // For a tuple extension, the generic parameter list comes from the
+    // extended type alias.
+    if (isa<BuiltinTupleDecl>(nominal)) {
+      if (auto *extendedTypeRepr = ext->getExtendedTypeRepr()) {
+        auto *genericParams = createTupleExtensionGenericParams(
+            ctx, ext, extendedTypeRepr);
+        if (genericParams)
+          return genericParams;
+
+        // Otherwise, just clone the generic parameter list of the
+        // Builtin.TheTupleType. We'll diagnose later.
+      }
+    }
+
     auto *genericParams = createExtensionGenericParams(ctx, ext, nominal);
 
     // Protocol extensions need an inheritance clause due to how name lookup
diff --git a/lib/Sema/TypeCheckDecl.cpp b/lib/Sema/TypeCheckDecl.cpp
@@ -2979,21 +2979,24 @@ ExtendedTypeRequest::evaluate(Evaluator &eval, ExtensionDecl *ext) const {
       PlaceholderType::get,
       /*packElementOpener*/ nullptr);
 
-  const auto extendedType = resolution.resolveType(extendedRepr);
+  auto extendedType = resolution.resolveType(extendedRepr);
 
   if (extendedType->hasError())
     return error();
 
   // Hack to allow extending a generic typealias.
   if (auto *unboundGeneric = extendedType->getAs<UnboundGenericType>()) {
     if (auto *aliasDecl = dyn_cast<TypeAliasDecl>(unboundGeneric->getDecl())) {
-      auto extendedNominal =
-          aliasDecl->getUnderlyingType()->getAnyNominal();
-      if (extendedNominal)
+      auto underlyingType = aliasDecl->getUnderlyingType();
+      if (auto extendedNominal = underlyingType->getAnyNominal()) {
         return TypeChecker::isPassThroughTypealias(
                    aliasDecl, extendedNominal)
                    ? extendedType
                    : extendedNominal->getDeclaredType();
+      }
+
+      if (underlyingType->is<TupleType>())
+        extendedType = underlyingType;
     }
   }
 
@@ -3006,8 +3009,9 @@ ExtendedTypeRequest::evaluate(Evaluator &eval, ExtensionDecl *ext) const {
     return error();
   }
 
-  // Cannot extend function types, tuple types, etc.
-  if (!extendedType->getAnyNominal() &&
+  // Cannot extend function types, metatypes, existentials, etc.
+  if (!extendedType->is<TupleType>() &&
+      !extendedType->getAnyNominal() &&
       !extendedType->is<ParameterizedProtocolType>()) {
     diags.diagnose(ext->getLoc(), diag::non_nominal_extension, extendedType)
          .highlight(extendedRepr->getSourceRange());
diff --git a/lib/Sema/TypeCheckDeclPrimary.cpp b/lib/Sema/TypeCheckDeclPrimary.cpp
@@ -3493,9 +3493,17 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
     auto *nominal = ED->getExtendedNominal();
 
     // Diagnose experimental tuple extensions.
-    if (nominal && isa<BuiltinTupleDecl>(nominal) &&
-        !getASTContext().LangOpts.hasFeature(Feature::TupleConformances)) {
-      ED->diagnose(diag::experimental_tuple_extension);
+    if (nominal && isa<BuiltinTupleDecl>(nominal)) {
+      if (!getASTContext().LangOpts.hasFeature(Feature::TupleConformances)) {
+        ED->diagnose(diag::experimental_tuple_extension);
+      }
+
+      if (extType->is<TupleType>()) {
+        auto selfType = ED->getSelfInterfaceType();
+        if (!extType->isEqual(selfType)) {
+          ED->diagnose(diag::tuple_extension_wrong_type, selfType);
+        }
+      }
     }
 
     if (nominal == nullptr) {
diff --git a/lib/Sema/TypeCheckGeneric.cpp b/lib/Sema/TypeCheckGeneric.cpp
@@ -571,6 +571,10 @@ static void collectAdditionalExtensionRequirements(
   if (type->is<ErrorType>())
     return;
 
+  // Tuple extensions cannot introduce requirements via this mechanism.
+  if (type->is<TupleType>())
+    return;
+
   // Find the nominal type declaration and its parent type.
   if (type->is<ProtocolCompositionType>())
     type = type->getCanonicalType();
@@ -775,7 +779,11 @@ GenericSignatureRequest::evaluate(Evaluator &evaluator,
     collectAdditionalExtensionRequirements(ext->getExtendedType(), sameTypeReqs);
 
     // Re-use the signature of the type being extended by default.
-    if (sameTypeReqs.empty() && !ext->getTrailingWhereClause()) {
+    // For tuple extensions, always build a new signature to get
+    // the right sugared types, since we don't want to expose the
+    // name of the generic parameter of BuiltinTupleDecl itself.
+    if (sameTypeReqs.empty() && !ext->getTrailingWhereClause() &&
+        !isa<BuiltinTupleDecl>(ext->getExtendedNominal())) {
       return parentSig;
     }
 
diff --git a/test/Generics/tuple-conformances.swift b/test/Generics/tuple-conformances.swift
@@ -1,18 +1,23 @@
-// RUN: %target-typecheck-verify-swift -enable-experimental-feature TupleConformances -parse-stdlib
+// RUN: %target-typecheck-verify-swift -enable-experimental-feature TupleConformances
 
+// Because of -enable-experimental-feature TupleConformances
 // REQUIRES: asserts
 
-import Swift
-
 protocol P {
   associatedtype A
   associatedtype B
 
   func f()
 }
 
-extension Builtin.TheTupleType: P where repeat each Elements: P {
-  typealias A = (repeat (each Elements).A)
+extension () {}
+// expected-error@-1 {{tuple extension must be written as extension of '(repeat each Element)'}}
+// FIXME: Inaccurate
+
+typealias Tuple<each Element> = (repeat each Element)
+
+extension Tuple: P where repeat each Element: P {
+  typealias A = (repeat (each Element).A)
   typealias B = Float
   func f() {}
 }
@@ -37,7 +42,7 @@ func useConformance() {
 
 ////
 
-extension Builtin.TheTupleType: Equatable where repeat each Elements: Equatable {
+extension Tuple: Equatable where repeat each Element: Equatable {
   // FIXME: Hack
   @_disfavoredOverload
   public static func ==(lhs: Self, rhs: Self) -> Bool {
@@ -51,7 +56,7 @@ extension Builtin.TheTupleType: Equatable where repeat each Elements: Equatable
   }
 }
 
-extension Builtin.TheTupleType: Hashable where repeat each Elements: Hashable {
+extension Tuple: Hashable where repeat each Element: Hashable {
   public func hash(into hasher: inout Hasher) {
     repeat (each self).hash(into: &hasher)
   }
diff --git a/test/IRGen/tuple_conformances.swift b/test/IRGen/tuple_conformances.swift
@@ -1,10 +1,8 @@
-// RUN: %target-swift-frontend -emit-ir -parse-stdlib -primary-file %s -enable-experimental-feature TupleConformances -parse-as-library | %FileCheck %s
+// RUN: %target-swift-frontend -emit-ir -primary-file %s -enable-experimental-feature TupleConformances -parse-as-library | %FileCheck %s
 
 // -enable-experimental-feature requires an asserts build
 // REQUIRES: asserts
 
-import Swift
-
 protocol P {
   func defaultRequirement()
   static func staticRequirement()
@@ -14,7 +12,9 @@ extension P {
   func defaultRequirement() {}
 }
 
-extension Builtin.TheTupleType: P where repeat each Elements: P {
+typealias Tuple<each Element> = (repeat each Element)
+
+extension Tuple: P where repeat each Element: P {
   static func staticRequirement() {}
 }
 
