Merge pull request #67310 from slavapestov/transform-type-parameter-packs-5.9

Sema: Fix ASTPrinter and opened element environments to handle nested pack expansion types [5.9]

Author: slavapestov

include/swift/AST/Type.h

@@ -316,6 +316,14 @@ class Type {
       llvm::function_ref<llvm::Optional<Type>(TypeBase *, TypePosition)> fn)
       const;
 
+  /// Transform free pack element references, that is, those not captured by a
+  /// pack expansion.
+  ///
+  /// This is the 'map' counterpart to TypeBase::getTypeParameterPacks().
+  Type transformTypeParameterPacks(
+      llvm::function_ref<llvm::Optional<Type>(SubstitutableType *)> fn)
+      const;
+
   /// Look through the given type and its children and apply fn to them.
   void visit(llvm::function_ref<void (Type)> fn) const {
     findIf([&fn](Type t) -> bool {

lib/AST/ASTContext.cpp

@@ -5856,13 +5856,13 @@ ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig,
   }
 
   auto eraseParameterPackRec = [&](Type type) -> Type {
-    return type.transformRec([&](Type t) -> llvm::Optional<Type> {
-      if (auto *paramType = t->getAs<GenericTypeParamType>()) {
+    return type.transformTypeParameterPacks([&](SubstitutableType *t) -> llvm::Optional<Type> {
+      if (auto *paramType = dyn_cast<GenericTypeParamType>(t)) {
         if (packElementParams.find(paramType) != packElementParams.end()) {
           return Type(packElementParams[paramType]);
         }
 
-        return t;
+        return Type(t);
       }
       return llvm::None;
     });

lib/AST/ASTPrinter.cpp

@@ -804,7 +804,7 @@ class PrintAST : public ASTVisitor<PrintAST> {
 
   void printType(Type T) { printTypeWithOptions(T, Options); }
 
-  void printTransformedTypeWithOptions(Type T, PrintOptions options) {
+  Type getTransformedType(Type T) {
     if (CurrentType && Current && CurrentType->mayHaveMembers()) {
       auto *M = Current->getDeclContext()->getParentModule();
       SubstitutionMap subMap;
@@ -825,9 +825,16 @@ class PrintAST : public ASTVisitor<PrintAST> {
       }
 
       T = T.subst(subMap, SubstFlags::DesugarMemberTypes);
+    }
+
+    return T;
+  }
 
+  void printTransformedTypeWithOptions(Type T, PrintOptions options) {
+    T = getTransformedType(T);
+
+    if (CurrentType && Current && CurrentType->mayHaveMembers())
       options.TransformContext = TypeTransformContext(CurrentType);
-    }
 
     printTypeWithOptions(T, options);
   }
@@ -1833,6 +1840,11 @@ void PrintAST::printSingleDepthOfGenericSignature(
 }
 
 void PrintAST::printRequirement(const Requirement &req) {
+  SmallVector<Type, 2> rootParameterPacks;
+  getTransformedType(req.getFirstType())
+      ->getTypeParameterPacks(rootParameterPacks);
+  bool isPackRequirement = !rootParameterPacks.empty();
+
   switch (req.getKind()) {
   case RequirementKind::SameShape:
     Printer << "(repeat (";
@@ -1842,22 +1854,21 @@ void PrintAST::printRequirement(const Requirement &req) {
     Printer << ")) : Any";
     return;
   case RequirementKind::Layout:
-    if (req.getFirstType()->hasParameterPack())
+    if (isPackRequirement)
       Printer << "repeat ";
     printTransformedType(req.getFirstType());
     Printer << " : ";
     req.getLayoutConstraint()->print(Printer, Options);
     return;
   case RequirementKind::Conformance:
   case RequirementKind::Superclass:
-    if (req.getFirstType()->hasParameterPack())
+    if (isPackRequirement)
       Printer << "repeat ";
     printTransformedType(req.getFirstType());
     Printer << " : ";
     break;
   case RequirementKind::SameType:
-    if (req.getFirstType()->hasParameterPack() ||
-        req.getSecondType()->hasParameterPack())
+    if (isPackRequirement)
       Printer << "repeat ";
     printTransformedType(req.getFirstType());
     Printer << " == ";

lib/AST/ParameterPack.cpp

@@ -25,6 +25,64 @@
 
 using namespace swift;
 
+/// FV(PackExpansionType(Pattern, Count), N) = FV(Pattern, N+1)
+/// FV(PackElementType(Param, M), N) = FV(Param, 0) if M >= N, {} otherwise
+/// FV(Param, N) = {Param}
+static Type transformTypeParameterPacksRec(
+    Type t, llvm::function_ref<llvm::Optional<Type>(SubstitutableType *)> fn,
+    unsigned expansionLevel) {
+  return t.transformWithPosition(
+      TypePosition::Invariant,
+      [&](TypeBase *t, TypePosition p) -> llvm::Optional<Type> {
+
+    // If we're already inside N levels of PackExpansionType,  and we're
+    // walking into another PackExpansionType, a type parameter pack
+    // reference now needs level (N+1) to be free.
+    if (auto *expansionType = dyn_cast<PackExpansionType>(t)) {
+      auto countType = expansionType->getCountType();
+      auto patternType = expansionType->getPatternType();
+      auto newPatternType = transformTypeParameterPacksRec(
+          patternType, fn, expansionLevel + 1);
+      if (patternType.getPointer() != newPatternType.getPointer())
+        return Type(PackExpansionType::get(patternType, countType));
+
+      return Type(expansionType);
+    }
+
+    // A PackElementType with level N reaches past N levels of
+    // nested PackExpansionType. So a type parameter pack reference
+    // therein is free if N is greater than or equal to our current
+    // expansion level.
+    if (auto *eltType = dyn_cast<PackElementType>(t)) {
+      if (eltType->getLevel() >= expansionLevel) {
+        return transformTypeParameterPacksRec(eltType->getPackType(), fn,
+                                              /*expansionLevel=*/0);
+      }
+
+      return Type(eltType);
+    }
+
+    // A bare type parameter pack is like a PackElementType with level 0.
+    if (auto *paramType = dyn_cast<SubstitutableType>(t)) {
+      if (expansionLevel == 0 &&
+          (isa<PackArchetypeType>(paramType) ||
+           (isa<GenericTypeParamType>(paramType) &&
+            cast<GenericTypeParamType>(paramType)->isParameterPack()))) {
+        return fn(paramType);
+      }
+
+      return Type(paramType);
+    }
+
+    return llvm::None;
+  });
+}
+
+Type Type::transformTypeParameterPacks(
+    llvm::function_ref<llvm::Optional<Type>(SubstitutableType *)> fn) const {
+  return transformTypeParameterPacksRec(*this, fn, /*expansionLevel=*/0);
+}
+
 namespace {
 
 /// Collects all unique pack type parameters referenced from the pattern type,

lib/Sema/ConstraintSystem.cpp

@@ -1024,8 +1024,11 @@ Type ConstraintSystem::openType(Type type, OpenedTypeMap &replacements,
       // that gets introduced by the interface type, see
       // \c openUnboundGenericType for more details.
       if (auto *packTy = type->getAs<PackType>()) {
-        if (auto expansion = packTy->unwrapSingletonPackExpansion())
-          type = expansion->getPatternType();
+        if (auto expansionTy = packTy->unwrapSingletonPackExpansion()) {
+          auto patternTy = expansionTy->getPatternType();
+          if (patternTy->isTypeParameter())
+            return openType(patternTy, replacements, locator);
+        }
       }
 
       if (auto *expansion = type->getAs<PackExpansionType>()) {

test/ModuleInterface/pack_expansion_type.swift

@@ -2,6 +2,8 @@
 // RUN: %target-swift-emit-module-interface(%t/PackExpansionType.swiftinterface) %s -module-name PackExpansionType -disable-availability-checking
 // RUN: %FileCheck %s < %t/PackExpansionType.swiftinterface
 
+/// Requirements
+
 // CHECK: #if compiler(>=5.3) && $ParameterPacks
 // CHECK-NEXT: public func variadicFunction<each T, each U>(t: repeat each T, u: repeat each U) -> (repeat (each T, each U)) where (repeat (each T, each U)) : Any
 public func variadicFunction<each T, each U>(t: repeat each T, u: repeat each U) -> (repeat (each T, each U)) {}
@@ -24,6 +26,13 @@ public struct VariadicType<each T> {
 // CHECK: }
 // CHECK-NEXT: #endif
 
+// The second requirement should not be prefixed with 'repeat'
+// CHECK: public struct SameTypeReq<T, each U> where T : Swift.Sequence, T.Element == PackExpansionType.VariadicType<repeat each U> {
+public struct SameTypeReq<T: Sequence, each U> where T.Element == VariadicType<repeat each U> {}
+// CHECK: }
+
+/// Pack expansion types
+
 // CHECK: public func returnsVariadicType() -> PackExpansionType.VariadicType<>
 public func returnsVariadicType() -> VariadicType< > {}
 

validation-test/compiler_crashers_2_fixed/rdar112108253.swift

@@ -0,0 +1,17 @@
+// RUN: %target-swift-frontend -emit-ir %s -disable-availability-checking
+
+public protocol P {
+  associatedtype A
+}
+
+public struct G<each T> {}
+
+public protocol Q {
+  init()
+}
+
+public struct S<T: P, each U: P>: Q where repeat T.A == G<repeat each U.A> {
+  public init() {}
+  public init(predicate: repeat each U) {}
+}
+