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

Fix various places where we didn't handle "bound" pack references correctly

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

@@ -5867,13 +5867,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

@@ -803,7 +803,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;
@@ -824,9 +824,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);
   }
@@ -1832,6 +1839,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 (";
@@ -1841,22 +1853,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/GenericEnvironment.cpp

@@ -659,6 +659,10 @@ Type GenericEnvironment::mapTypeIntoContext(GenericTypeParamType *type) const {
   return result;
 }
 
+/// So this expects a type written with the archetypes of the original generic
+/// environment, not 'this', the opened element environment, because it is the
+/// original PackArchetypes that become ElementArchetypes. Also this function
+/// does not apply outer substitutions, which might not be what you expect.
 Type
 GenericEnvironment::mapContextualPackTypeIntoElementContext(Type type) const {
   assert(getKind() == Kind::OpenedElement);
@@ -672,77 +676,40 @@ GenericEnvironment::mapContextualPackTypeIntoElementContext(Type type) const {
   FindElementArchetypeForOpenedPackParam
     findElementArchetype(this, getOpenedPackParams());
 
-  return type.transformRec([&](TypeBase *ty) -> llvm::Optional<Type> {
-    // We're only directly substituting pack archetypes.
-    auto archetype = ty->getAs<PackArchetypeType>();
-    if (!archetype) {
-      // Don't recurse into nested pack expansions.
-      if (ty->is<PackExpansionType>())
-        return Type(ty);
+  return type.transformTypeParameterPacks(
+      [&](SubstitutableType *ty) -> llvm::Optional<Type> {
+        if (auto *packArchetype = dyn_cast<PackArchetypeType>(ty)) {
+          auto interfaceType = packArchetype->getInterfaceType();
+          if (sig->haveSameShape(interfaceType, shapeClass))
+            return Type(findElementArchetype(interfaceType));
+        }
 
-      // Recurse into any other type.
-      return llvm::None;
-    }
-
-    auto rootArchetype = cast<PackArchetypeType>(archetype->getRoot());
-
-    // TODO: assert that the generic environment of the pack archetype
-    // matches the signature that was originally opened to make this
-    // environment.  Unfortunately, that isn't a trivial check because of
-    // the extra opened-element parameters.
-
-    // If the archetype isn't the shape that was opened by this
-    // environment, ignore it.
-    auto rootParam = cast<GenericTypeParamType>(
-      rootArchetype->getInterfaceType().getPointer());
-    assert(rootParam->isParameterPack());
-    if (!sig->haveSameShape(rootParam, shapeClass))
-      return Type(ty);
-
-    return Type(findElementArchetype(archetype->getInterfaceType()));
-  });
+        return llvm::None;
+      });
 }
 
 CanType
 GenericEnvironment::mapContextualPackTypeIntoElementContext(CanType type) const {
   return CanType(mapContextualPackTypeIntoElementContext(Type(type)));
 }
 
+/// Unlike mapContextualPackTypeIntoElementContext(), this also applies outer
+/// substitutions, so it behaves like mapTypeIntoContext() in that respect.
 Type
 GenericEnvironment::mapPackTypeIntoElementContext(Type type) const {
   assert(getKind() == Kind::OpenedElement);
   assert(!type->hasArchetype());
 
-  auto sig = getGenericSignature();
-  auto shapeClass = getOpenedElementShapeClass();
-
-  FindElementArchetypeForOpenedPackParam
-    findElementArchetype(this, getOpenedPackParams());
+  if (!type->hasTypeParameter()) return type;
 
-  // Map the interface type to the element type by stripping
-  // away the isParameterPack bit before mapping type parameters
-  // to archetypes.
-  return type.transformRec([&](TypeBase *ty) -> llvm::Optional<Type> {
-    // We're only directly substituting pack parameters.
-    if (!ty->isTypeParameter()) {
-      // Don't recurse into nested pack expansions; just map it into
-      // context.
-      if (ty->is<PackExpansionType>())
-        return mapTypeIntoContext(ty);
-
-      // Recurse into any other type.
-      return llvm::None;
-    }
-
-    // Just do normal mapping for types that are not rooted in
-    // opened type parameters.
-    auto rootParam = ty->getRootGenericParam();
-    if (!rootParam->isParameterPack() ||
-        !sig->haveSameShape(rootParam, shapeClass))
-      return mapTypeIntoContext(ty);
+  // Get a contextual type in the original generic environment, not the
+  // substituted one, which is what mapContextualPackTypeIntoElementContext()
+  // expects.
+  auto contextualType = getPackElementContextSubstitutions()
+    .getGenericSignature().getGenericEnvironment()->mapTypeIntoContext(type);
 
-    return Type(findElementArchetype(ty));
-  });
+  contextualType = mapContextualPackTypeIntoElementContext(contextualType);
+  return maybeApplyOuterContextSubstitutions(contextualType);
 }
 
 Type

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/ClangImporter/ImportName.cpp

@@ -2247,7 +2247,6 @@ ImportedName NameImporter::importNameImpl(const clang::NamedDecl *D,
         if (arg.getKind() == clang::TemplateArgument::Type) {
           auto ty = arg.getAsType().getTypePtr();
           if (auto builtin = dyn_cast<clang::BuiltinType>(ty)) {
-            auto &ctx = swiftCtx;
             if (auto swiftTypeName = getSwiftBuiltinTypeName(builtin)) {
               buffer << *swiftTypeName;
               return;

lib/IRGen/GenPack.cpp

@@ -330,8 +330,8 @@ static void bindElementSignatureRequirementsAtIndex(
               IGF.emitTypeMetadataRef(patternPackArchetype, request);
           auto elementArchetype =
               context.environment
-                  ->mapPackTypeIntoElementContext(
-                      patternPackArchetype->getInterfaceType())
+                  ->mapContextualPackTypeIntoElementContext(
+                      patternPackArchetype)
                   ->getCanonicalType();
           auto *patternPack = response.getMetadata();
           auto elementMetadata = bindMetadataAtIndex(
@@ -346,8 +346,8 @@ static void bindElementSignatureRequirementsAtIndex(
           auto patternPackArchetype = getMappedPackArchetypeType(context, ty);
           auto elementArchetype =
               context.environment
-                  ->mapPackTypeIntoElementContext(
-                      patternPackArchetype->getInterfaceType())
+                  ->mapContextualPackTypeIntoElementContext(
+                      patternPackArchetype)
                   ->getCanonicalType();
           llvm::Value *_metadata = nullptr;
           auto packConformance =

lib/Sema/CSSimplify.cpp

@@ -9162,7 +9162,6 @@ ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
   if (!patternType->hasTypeVariable()) {
     auto *loc = getConstraintLocator(locator);
     auto shapeClass = patternType->getReducedShape();
-    patternType = patternType->mapTypeOutOfContext();
     auto *elementEnv = getPackElementEnvironment(loc, shapeClass);
 
     // Without an opened element environment, we cannot derive the
@@ -9185,7 +9184,7 @@ ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
     }
 
     auto expectedElementTy =
-        elementEnv->mapPackTypeIntoElementContext(patternType);
+        elementEnv->mapContextualPackTypeIntoElementContext(patternType);
     assert(!expectedElementTy->is<PackType>());
 
     addConstraint(ConstraintKind::Equal, elementType, expectedElementTy,

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