[GenericEnvironment] For opened pack element environments, only include element type

parameters whose originating packs are in a given shape equivalence class.

Author: hborla

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.

lib/AST/ASTContext.cpp

@@ -5012,7 +5012,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();
 
@@ -5035,7 +5036,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;
@@ -5638,6 +5640,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(),
@@ -5649,7 +5656,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]);
         }
 

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

@@ -177,7 +177,7 @@ static void emitPackExpansionType(IRGenFunction &IGF,
   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) {

lib/Sema/CSApply.cpp

@@ -7124,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,19 +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(),
-                                                      packType->getReducedShape());
+                                                      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

@@ -8621,8 +8621,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 (packType->isEqual(packEnv->mapElementTypeIntoPackContext(elementType)) ||
+      (!elementType->hasElementArchetype() && packType->isEqual(elementType))) {
     return SolutionKind::Solved;
   } else {
     return SolutionKind::Error;

lib/Sema/ConstraintSystem.cpp

@@ -655,7 +655,8 @@ ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator,
       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

@@ -1259,14 +1259,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)
@@ -1280,12 +1280,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;

lib/Serialization/ModuleFormat.h

@@ -1759,7 +1759,7 @@ namespace decls_block {
   using GenericEnvironmentLayout = BCRecordLayout<
     GENERIC_ENVIRONMENT,
     BCFixed<1>,                  // GenericEnvironmentKind
-    TypeIDField,                 // existential type
+    TypeIDField,                 // existential type or shape class
     GenericSignatureIDField,     // parent signature
     SubstitutionMapIDField       // substitution map
   >;

lib/Serialization/Serialization.cpp

@@ -1529,14 +1529,15 @@ void Serializer::writeASTBlockEntity(const GenericEnvironment *genericEnv) {
 
   case GenericEnvironment::Kind::OpenedElement: {
     auto kind = GenericEnvironmentKind::OpenedElement;
+    auto shapeClassID = addTypeRef(genericEnv->getOpenedElementShapeClass());
     auto parentSig = genericEnv->getGenericSignature();
     auto parentSigID = addGenericSignatureRef(parentSig);
     auto contextSubs = genericEnv->getPackElementContextSubstitutions();
     auto subsID = addSubstitutionMapRef(contextSubs);
 
     auto genericEnvAbbrCode = DeclTypeAbbrCodes[GenericEnvironmentLayout::Code];
     GenericEnvironmentLayout::emitRecord(Out, ScratchRecord, genericEnvAbbrCode,
-                                         unsigned(kind), /*existentialTypeID=*/0,
+                                         unsigned(kind), shapeClassID,
                                          parentSigID, subsID);
     return;
   }