[Sema] Key opened pack element generic on the shape class of the expansion.

Author: hborla

include/swift/AST/ASTContext.h

@@ -1367,7 +1367,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/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.
@@ -5599,9 +5600,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;
 
@@ -5693,7 +5696,7 @@ ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig) {
   auto elementSig = buildGenericSignature(
       *this, GenericSignature(), genericParams, requirements)
           .getCanonicalSignature();
-  sigs[baseGenericSig.getPointer()] = elementSig;
+  sigs[key] = elementSig;
   return elementSig;
 }
 

lib/IRGen/GenPack.cpp

@@ -174,7 +174,8 @@ 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);
 
@@ -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) {

lib/Sema/CSBindings.cpp

@@ -1424,7 +1424,8 @@ void PotentialBindings::infer(Constraint *constraint) {
       // Produce a potential binding to the opened element archetype corresponding
       // to the pack type.
       packType = packType->mapTypeOutOfContext();
-      auto *elementEnv = CS.getPackElementEnvironment(constraint->getLocator());
+      auto *elementEnv = CS.getPackElementEnvironment(constraint->getLocator(),
+                                                      packType->getReducedShape());
       auto elementType = elementEnv->mapPackTypeIntoElementContext(packType);
       addPotentialBinding({elementType, AllowedBindingKind::Exact, constraint});
 

lib/Sema/ConstraintSystem.cpp

@@ -643,15 +643,16 @@ 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);