[GenericEnvironment] Only include opened pack elements within a given shape class in an opened element generic environment.

include/swift/AST/ASTContext.h

@@ -1369,7 +1369,8 @@ class ASTContext final {
   ///
   /// This drops the parameter pack bit from each generic parameter,
   /// and converts same-element requirements to same-type requirements.
-  CanGenericSignature getOpenedElementSignature(CanGenericSignature baseGenericSig);
+  CanGenericSignature getOpenedElementSignature(CanGenericSignature baseGenericSig,
+                                                CanType shapeClass);
 
   GenericSignature getOverrideGenericSignature(const ValueDecl *base,
                                                const ValueDecl *derived);

include/swift/AST/GenericEnvironment.h

@@ -67,6 +67,7 @@ struct OpenedExistentialEnvironmentData {
 /// Extra data in a generic environment for an opened pack element.
 struct OpenedElementEnvironmentData {
   UUID uuid;
+  CanType shapeClass;
   SubstitutionMap outerSubstitutions;
 };
 
@@ -140,7 +141,8 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
       Type existential, GenericSignature parentSig, UUID uuid);
 
   /// Private constructor for opened element environments.
-  explicit GenericEnvironment(GenericSignature signature, UUID uuid,
+  explicit GenericEnvironment(GenericSignature signature,
+                              UUID uuid, CanType shapeClass,
                               SubstitutionMap outerSubs);
 
   friend ArchetypeType;
@@ -187,6 +189,9 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   /// opened pack element generic environment.
   SubstitutionMap getPackElementContextSubstitutions() const;
 
+  /// Retrieve the shape equivalence class for an opened element environment.
+  CanType getOpenedElementShapeClass() const;
+
   /// Retrieve the UUID for an opened element environment.
   UUID getOpenedElementUUID() const;
 
@@ -222,10 +227,12 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   /// signature of the context whose element type is being opened, but with
   /// the pack parameter bit erased from one or more generic parameters
   /// \param uuid The unique identifier for this opened element
+  /// \param shapeClass The shape equivalence class for the originating packs.
   /// \param outerSubs The substitution map containing archetypes from the
   /// outer generic context.
   static GenericEnvironment *
-  forOpenedElement(GenericSignature signature, UUID uuid,
+  forOpenedElement(GenericSignature signature,
+                   UUID uuid, CanType shapeClass,
                    SubstitutionMap outerSubs);
 
   /// Make vanilla new/delete illegal.

include/swift/Sema/ConstraintSystem.h

@@ -3987,7 +3987,8 @@ class ConstraintSystem {
   void addPackElementEnvironment(PackExpansionExpr *expr);
 
   /// Get the opened element generic environment for the given locator.
-  GenericEnvironment *getPackElementEnvironment(ConstraintLocator *locator);
+  GenericEnvironment *getPackElementEnvironment(ConstraintLocator *locator,
+                                                CanType shapeClass);
 
   /// Retrieve the constraint locator for the given anchor and
   /// path, uniqued and automatically infer the summary flags

lib/AST/ASTContext.cpp

@@ -356,9 +356,10 @@ struct ASTContext::Implementation {
   llvm::DenseMap<std::pair<CanType, const GenericSignatureImpl *>, CanGenericSignature>
       ExistentialSignatures;
 
-  /// The element signature for a generic signature, constructed by dropping
-  /// the parameter pack bit from generic parameters.
-  llvm::DenseMap<const GenericSignatureImpl *,
+  /// The element signature for a generic signature, which contains a clone
+  /// of the context generic signature with new type parameters and requirements
+  /// for opened pack elements in the given shape equivalence class.
+  llvm::DenseMap<std::pair<CanType, const GenericSignatureImpl *>,
                  CanGenericSignature> ElementSignatures;
 
   /// Overridden declarations.
@@ -5036,7 +5037,8 @@ GenericEnvironment::forOpenedExistential(
 
 /// Create a new generic environment for an element archetype.
 GenericEnvironment *
-GenericEnvironment::forOpenedElement(GenericSignature signature, UUID uuid,
+GenericEnvironment::forOpenedElement(GenericSignature signature,
+                                     UUID uuid, CanType shapeClass,
                                      SubstitutionMap outerSubs) {
   auto &ctx = signature->getASTContext();
 
@@ -5059,7 +5061,8 @@ GenericEnvironment::forOpenedElement(GenericSignature signature, UUID uuid,
                                   OpenedElementEnvironmentData, Type>(
       0, 0, 1, numGenericParams);
   void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
-  auto *genericEnv = new (mem) GenericEnvironment(signature, uuid,
+  auto *genericEnv = new (mem) GenericEnvironment(signature,
+                                                  uuid, shapeClass,
                                                   outerSubs);
 
   openedElementEnvironments[uuid] = genericEnv;
@@ -5624,9 +5627,11 @@ ASTContext::getOpenedExistentialSignature(Type type, GenericSignature parentSig)
 }
 
 CanGenericSignature
-ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig) {
+ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig,
+                                      CanType shapeClass) {
   auto &sigs = getImpl().ElementSignatures;
-  auto found = sigs.find(baseGenericSig.getPointer());
+  auto key = std::make_pair(shapeClass, baseGenericSig.getPointer());
+  auto found = sigs.find(key);
   if (found != sigs.end())
     return found->second;
 
@@ -5660,6 +5665,11 @@ ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig) {
     if (!paramType->isParameterPack())
       continue;
 
+    // Only include opened element parameters for packs in the given
+    // shape equivalence class.
+    if (!baseGenericSig->haveSameShape(paramType, shapeClass->mapTypeOutOfContext()))
+      continue;
+
     auto *elementParam = GenericTypeParamType::get(/*isParameterPack*/false,
                                                    packElementDepth,
                                                    packElementParams.size(),
@@ -5671,7 +5681,7 @@ ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig) {
   auto eraseParameterPackRec = [&](Type type) -> Type {
     return type.transformRec([&](Type t) -> Optional<Type> {
       if (auto *paramType = t->getAs<GenericTypeParamType>()) {
-        if (paramType->isParameterPack()) {
+        if (packElementParams.find(paramType) != packElementParams.end()) {
           return Type(packElementParams[paramType]);
         }
 
@@ -5718,7 +5728,7 @@ ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig) {
   auto elementSig = buildGenericSignature(
       *this, GenericSignature(), genericParams, requirements)
           .getCanonicalSignature();
-  sigs[baseGenericSig.getPointer()] = elementSig;
+  sigs[key] = elementSig;
   return elementSig;
 }
 

lib/AST/GenericEnvironment.cpp

@@ -104,6 +104,12 @@ GenericEnvironment::getPackElementContextSubstitutions() const {
   return environmentData->outerSubstitutions;
 }
 
+CanType GenericEnvironment::getOpenedElementShapeClass() const {
+  assert(getKind() == Kind::OpenedElement);
+  auto environmentData = getTrailingObjects<OpenedElementEnvironmentData>();
+  return environmentData->shapeClass;
+}
+
 Type GenericEnvironment::getOpenedExistentialType() const {
   assert(getKind() == Kind::OpenedExistential);
   return getTrailingObjects<OpenedExistentialEnvironmentData>()->existential;
@@ -127,7 +133,9 @@ UUID GenericEnvironment::getOpenedElementUUID() const {
 
 void GenericEnvironment::getPackElementBindings(
     SmallVectorImpl<PackElementBinding> &bindings) const {
-  auto packElements = getGenericSignature().getInnermostGenericParams();
+  auto sig = getGenericSignature();
+  auto shapeClass = getOpenedElementShapeClass();
+  auto packElements = sig.getInnermostGenericParams();
   auto packElementDepth = packElements.front()->getDepth();
   auto elementIt = packElements.begin();
 
@@ -141,6 +149,11 @@ void GenericEnvironment::getPackElementBindings(
     if (!genericParam->isParameterPack())
       continue;
 
+    // Only include opened element parameters for packs in the given
+    // shape equivalence class.
+    if (!sig->haveSameShape(genericParam, shapeClass->mapTypeOutOfContext()))
+      continue;
+
     assert(elementIt != packElements.end());
     auto *elementArchetype =
         mapTypeIntoContext(*elementIt++)->getAs<ElementArchetypeType>();
@@ -185,12 +198,13 @@ GenericEnvironment::GenericEnvironment(
                           Type());
 }
 
-GenericEnvironment::GenericEnvironment(
-      GenericSignature signature, UUID uuid, SubstitutionMap outerSubs)
+GenericEnvironment::GenericEnvironment(GenericSignature signature,
+                                       UUID uuid, CanType shapeClass,
+                                       SubstitutionMap outerSubs)
   : SignatureAndKind(signature, Kind::OpenedElement)
 {
   new (getTrailingObjects<OpenedElementEnvironmentData>())
-    OpenedElementEnvironmentData{uuid, outerSubs};
+    OpenedElementEnvironmentData{uuid, shapeClass, outerSubs};
 
   // Clear out the memory that holds the context types.
   std::uninitialized_fill(getContextTypes().begin(), getContextTypes().end(),
@@ -556,6 +570,7 @@ GenericEnvironment::mapPackTypeIntoElementContext(Type type) const {
   assert(!type->hasArchetype());
 
   auto sig = getGenericSignature();
+  auto shapeClass = getOpenedElementShapeClass();
   QueryInterfaceTypeSubstitutions substitutions(this);
 
   llvm::SmallDenseMap<GenericParamKey,
@@ -570,6 +585,9 @@ GenericEnvironment::mapPackTypeIntoElementContext(Type type) const {
     if (!genericParam->isParameterPack())
       continue;
 
+    if (!sig->haveSameShape(genericParam, shapeClass->mapTypeOutOfContext()))
+      continue;
+
     auto elementIndex = elementParamForPack.size();
     elementParamForPack[{genericParam}] = packElements[elementIndex];
   }
@@ -592,11 +610,25 @@ GenericEnvironment::mapPackTypeIntoElementContext(Type type) const {
 Type
 GenericEnvironment::mapElementTypeIntoPackContext(Type type) const {
   assert(getKind() == Kind::Primary);
-  assert(!type->hasArchetype());
+
+  // We need to pass in an archetype to get the shape class from its
+  // generic environment.
+  assert(type->hasElementArchetype());
+
+  ElementArchetypeType *element = nullptr;
+  type.visit([&](Type type) {
+    auto archetype = type->getAs<ElementArchetypeType>();
+    if (!element && archetype)
+      element = archetype;
+  });
 
   auto sig = getGenericSignature();
+  auto *elementEnv = element->getGenericEnvironment();
+  auto shapeClass = elementEnv->getOpenedElementShapeClass();
   QueryInterfaceTypeSubstitutions substitutions(this);
 
+  type = type->mapTypeOutOfContext();
+
   llvm::SmallDenseMap<GenericParamKey, GenericTypeParamType *>
       packParamForElement;
   auto elementDepth =
@@ -606,6 +638,9 @@ GenericEnvironment::mapElementTypeIntoPackContext(Type type) const {
     if (!genericParam->isParameterPack())
       continue;
 
+    if (!sig->haveSameShape(genericParam, shapeClass->mapTypeOutOfContext()))
+      continue;
+
     GenericParamKey elementKey(/*isParameterPack*/false,
                                /*depth*/elementDepth,
                                /*index*/packParamForElement.size());

lib/IRGen/GenPack.cpp

@@ -174,9 +174,10 @@ static void emitPackExpansionType(IRGenFunction &IGF,
   auto genericSig = genericEnv->getGenericSignature().getCanonicalSignature();
 
   // Create an opened element signature and environment.
-  auto elementSig = IGF.IGM.Context.getOpenedElementSignature(genericSig);
+  auto elementSig = IGF.IGM.Context.getOpenedElementSignature(
+      genericSig, expansionTy.getCountType());
   auto *elementEnv = GenericEnvironment::forOpenedElement(
-      elementSig, UUID::fromTime(), subMap);
+      elementSig, UUID::fromTime(), expansionTy.getCountType(), subMap);
 
   // Open each pack archetype.
   for (auto patternPackType : patternPacks) {
@@ -325,4 +326,4 @@ void irgen::cleanupTypeMetadataPack(IRGenFunction &IGF,
     IGF.Builder.CreateLifetimeEnd(pack.getAddress(),
                                   IGF.IGM.getPointerSize() * (*elementCount));
   }
-}
+}

lib/Sema/CSApply.cpp

@@ -3818,11 +3818,16 @@ namespace {
     }
 
     Expr *visitPackExpansionExpr(PackExpansionExpr *expr) {
+      simplifyExprType(expr);
+
+      // Set the opened pack element environment for this pack expansion.
+      auto expansionTy = cs.getType(expr)->castTo<PackExpansionType>();
       auto *locator = cs.getConstraintLocator(expr);
-      auto *environment = cs.getPackElementEnvironment(locator);
+      auto *environment = cs.getPackElementEnvironment(locator,
+          expansionTy->getCountType()->getCanonicalType());
       expr->setGenericEnvironment(environment);
 
-      return simplifyExprType(expr);
+      return expr;
     }
 
     Expr *visitPackElementExpr(PackElementExpr *expr) {
@@ -7119,8 +7124,7 @@ Expr *ExprRewriter::coerceToType(Expr *expr, Type toType,
     auto *pattern = coerceToType(expansion->getPatternExpr(),
                                  toElementType, locator);
     auto *packEnv = cs.DC->getGenericEnvironmentOfContext();
-    auto patternType = packEnv->mapElementTypeIntoPackContext(
-        toElementType->mapTypeOutOfContext());
+    auto patternType = packEnv->mapElementTypeIntoPackContext(toElementType);
     auto shapeType = toExpansionType->getCountType();
     auto expansionTy = PackExpansionType::get(patternType, shapeType);
 

lib/Sema/CSBindings.cpp

@@ -1423,18 +1423,25 @@ void PotentialBindings::infer(Constraint *constraint) {
     if (elementVar == TypeVar && !packVar) {
       // Produce a potential binding to the opened element archetype corresponding
       // to the pack type.
+      auto shapeClass = packType->getReducedShape();
       packType = packType->mapTypeOutOfContext();
-      auto *elementEnv = CS.getPackElementEnvironment(constraint->getLocator());
+      auto *elementEnv = CS.getPackElementEnvironment(constraint->getLocator(),
+                                                      shapeClass);
       auto elementType = elementEnv->mapPackTypeIntoElementContext(packType);
       addPotentialBinding({elementType, AllowedBindingKind::Exact, constraint});
 
       break;
     } else if (packVar == TypeVar && !elementVar) {
       // Produce a potential binding to the pack archetype corresponding to
       // the opened element type.
+      Type patternType;
       auto *packEnv = CS.DC->getGenericEnvironmentOfContext();
-      elementType = elementType->mapTypeOutOfContext();
-      auto patternType = packEnv->mapElementTypeIntoPackContext(elementType);
+      if (!elementType->hasElementArchetype()) {
+        patternType = elementType;
+      } else {
+        patternType = packEnv->mapElementTypeIntoPackContext(elementType);
+      }
+
       addPotentialBinding({patternType, AllowedBindingKind::Exact, constraint});
 
       break;

lib/Sema/CSSimplify.cpp

@@ -8704,8 +8704,8 @@ ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
 
   // This constraint only exists to vend bindings.
   auto *packEnv = DC->getGenericEnvironmentOfContext();
-  if (packType->isEqual(packEnv->mapElementTypeIntoPackContext
-                        (elementType->mapTypeOutOfContext()))) {
+  if ((!elementType->hasElementArchetype() && packType->isEqual(elementType)) ||
+      packType->isEqual(packEnv->mapElementTypeIntoPackContext(elementType))) {
     return SolutionKind::Solved;
   } else {
     return SolutionKind::Error;

lib/Sema/ConstraintSystem.cpp

@@ -643,18 +643,20 @@ void ConstraintSystem::addPackElementEnvironment(PackExpansionExpr *expr) {
 }
 
 GenericEnvironment *
-ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator) {
+ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator,
+                                            CanType shapeClass) {
   auto result = PackExpansionEnvironments.find(locator);
   if (result == PackExpansionEnvironments.end())
     return nullptr;
 
   auto uuid = result->second;
   auto &ctx = getASTContext();
   auto elementSig = ctx.getOpenedElementSignature(
-      DC->getGenericSignatureOfContext().getCanonicalSignature());
+      DC->getGenericSignatureOfContext().getCanonicalSignature(), shapeClass);
   auto *contextEnv = DC->getGenericEnvironmentOfContext();
   auto contextSubs = contextEnv->getForwardingSubstitutionMap();
-  return GenericEnvironment::forOpenedElement(elementSig, uuid, contextSubs);
+  return GenericEnvironment::forOpenedElement(elementSig, uuid, shapeClass,
+                                              contextSubs);
 }
 
 /// Extend the given depth map by adding depths for all of the subexpressions

lib/Serialization/Deserialization.cpp

@@ -1261,14 +1261,14 @@ Expected<GenericEnvironment *> ModuleFile::getGenericEnvironmentChecked(
 
   unsigned kind;
   GenericSignatureID parentSigID;
-  TypeID existentialID;
+  TypeID existentialOrShapeID;
   SubstitutionMapID subsID;
-  GenericEnvironmentLayout::readRecord(scratch, kind, existentialID,
+  GenericEnvironmentLayout::readRecord(scratch, kind, existentialOrShapeID,
                                        parentSigID, subsID);
 
-  auto existentialTypeOrError = getTypeChecked(existentialID);
-  if (!existentialTypeOrError)
-    return existentialTypeOrError.takeError();
+  auto existentialOrShapeTypeOrError = getTypeChecked(existentialOrShapeID);
+  if (!existentialOrShapeTypeOrError)
+    return existentialOrShapeTypeOrError.takeError();
 
   auto parentSigOrError = getGenericSignatureChecked(parentSigID);
   if (!parentSigOrError)
@@ -1282,12 +1282,14 @@ Expected<GenericEnvironment *> ModuleFile::getGenericEnvironmentChecked(
   switch (GenericEnvironmentKind(kind)) {
   case GenericEnvironmentKind::OpenedExistential:
     genericEnv = GenericEnvironment::forOpenedExistential(
-        existentialTypeOrError.get(), parentSigOrError.get(), UUID::fromTime());
+        existentialOrShapeTypeOrError.get(), parentSigOrError.get(), UUID::fromTime());
     break;
 
   case GenericEnvironmentKind::OpenedElement:
     genericEnv = GenericEnvironment::forOpenedElement(
-        parentSigOrError.get(), UUID::fromTime(), contextSubsOrError.get());
+        parentSigOrError.get(), UUID::fromTime(),
+        existentialOrShapeTypeOrError.get()->getCanonicalType(),
+        contextSubsOrError.get());
   }
 
   envOffset = genericEnv;