[NFC] Store the interface shape type in an opened element environment

include/swift/AST/ASTContext.h

@@ -1385,7 +1385,7 @@ 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,
-                                                CanType shapeClass);
+                                                CanGenericTypeParamType shapeClass);
 
   GenericSignature getOverrideGenericSignature(const ValueDecl *base,
                                                const ValueDecl *derived);

include/swift/AST/DiagnosticsParse.def

@@ -1676,6 +1676,8 @@ ERROR(pack_element_attribute_expected_rparen,none,
       "expected ')' after id value for 'pack_element' attribute", ())
 ERROR(multiple_open_pack_element,none,
       "multiple 'open_pack_element' instructions with same UUID", ())
+ERROR(opened_shape_class_not_pack_param,none,
+      "shape class type must be a pack type parameter", ())
 
 // unowned
 ERROR(attr_unowned_invalid_specifier,none,

include/swift/AST/GenericEnvironment.h

@@ -71,7 +71,7 @@ struct OpenedExistentialEnvironmentData {
 /// Extra data in a generic environment for an opened pack element.
 struct OpenedElementEnvironmentData {
   UUID uuid;
-  CanType shapeClass;
+  CanGenericTypeParamType shapeClass;
   SubstitutionMap outerSubstitutions;
 };
 
@@ -146,7 +146,8 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
 
   /// Private constructor for opened element environments.
   explicit GenericEnvironment(GenericSignature signature,
-                              UUID uuid, CanType shapeClass,
+                              UUID uuid,
+                              CanGenericTypeParamType shapeClass,
                               SubstitutionMap outerSubs);
 
   friend ArchetypeType;
@@ -194,7 +195,8 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   SubstitutionMap getPackElementContextSubstitutions() const;
 
   /// Retrieve the shape equivalence class for an opened element environment.
-  CanType getOpenedElementShapeClass() const;
+  /// This is always a pack parameter.
+  CanGenericTypeParamType getOpenedElementShapeClass() const;
 
   /// Retrieve the UUID for an opened element environment.
   UUID getOpenedElementUUID() const;
@@ -238,12 +240,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 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, CanType shapeClass,
+                   UUID uuid, CanGenericTypeParamType shapeClass,
                    SubstitutionMap outerSubs);
 
   /// Make vanilla new/delete illegal.

lib/AST/ASTContext.cpp

@@ -5144,7 +5144,8 @@ GenericEnvironment::forOpenedExistential(
 /// Create a new generic environment for an element archetype.
 GenericEnvironment *
 GenericEnvironment::forOpenedElement(GenericSignature signature,
-                                     UUID uuid, CanType shapeClass,
+                                     UUID uuid,
+                                     CanGenericTypeParamType shapeClass,
                                      SubstitutionMap outerSubs) {
   auto &ctx = signature->getASTContext();
 
@@ -5735,7 +5736,7 @@ ASTContext::getOpenedExistentialSignature(Type type, GenericSignature parentSig)
 
 CanGenericSignature
 ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig,
-                                      CanType shapeClass) {
+                                      CanGenericTypeParamType shapeClass) {
   auto &sigs = getImpl().ElementSignatures;
   auto key = std::make_pair(shapeClass, baseGenericSig.getPointer());
   auto found = sigs.find(key);
@@ -5774,7 +5775,7 @@ ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig,
 
     // Only include opened element parameters for packs in the given
     // shape equivalence class.
-    if (!baseGenericSig->haveSameShape(paramType, shapeClass->mapTypeOutOfContext()))
+    if (!baseGenericSig->haveSameShape(paramType, shapeClass))
       continue;
 
     auto *elementParam = GenericTypeParamType::get(/*isParameterPack*/false,

lib/AST/ASTPrinter.cpp

@@ -6765,8 +6765,7 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
     auto sig = subs.getGenericSignature();
     auto params = sig.getGenericParams();
 
-    auto elementShapeClass =
-      env->getOpenedElementShapeClass()->mapTypeOutOfContext();
+    auto elementShapeClass = env->getOpenedElementShapeClass();
 
     // The element archetypes are at a depth one past the max depth
     // of the base signature.

lib/AST/GenericEnvironment.cpp

@@ -104,7 +104,8 @@ GenericEnvironment::getPackElementContextSubstitutions() const {
   return environmentData->outerSubstitutions;
 }
 
-CanType GenericEnvironment::getOpenedElementShapeClass() const {
+CanGenericTypeParamType
+GenericEnvironment::getOpenedElementShapeClass() const {
   assert(getKind() == Kind::OpenedElement);
   auto environmentData = getTrailingObjects<OpenedElementEnvironmentData>();
   return environmentData->shapeClass;
@@ -160,7 +161,7 @@ void GenericEnvironment::forEachPackElementGenericTypeParam(
 
     // Only include opened element parameters for packs in the given
     // shape equivalence class.
-    if (!sig->haveSameShape(genericParam, shapeClass->mapTypeOutOfContext()))
+    if (!sig->haveSameShape(genericParam, shapeClass))
       continue;
 
     function(genericParam);
@@ -218,7 +219,8 @@ GenericEnvironment::GenericEnvironment(
 }
 
 GenericEnvironment::GenericEnvironment(GenericSignature signature,
-                                       UUID uuid, CanType shapeClass,
+                                       UUID uuid,
+                                       CanGenericTypeParamType shapeClass,
                                        SubstitutionMap outerSubs)
   : SignatureAndKind(signature, Kind::OpenedElement)
 {
@@ -627,7 +629,7 @@ GenericEnvironment::mapPackTypeIntoElementContext(Type type) const {
     if (!genericParam->isParameterPack())
       continue;
 
-    if (!sig->haveSameShape(genericParam, shapeClass->mapTypeOutOfContext()))
+    if (!sig->haveSameShape(genericParam, shapeClass))
       continue;
 
     auto elementIndex = elementParamForPack.size();
@@ -680,7 +682,7 @@ GenericEnvironment::mapElementTypeIntoPackContext(Type type) const {
     if (!genericParam->isParameterPack())
       continue;
 
-    if (!sig->haveSameShape(genericParam, shapeClass->mapTypeOutOfContext()))
+    if (!sig->haveSameShape(genericParam, shapeClass))
       continue;
 
     GenericParamKey elementKey(/*isParameterPack*/false,
@@ -767,16 +769,19 @@ OpenedElementContext::createForContextualExpansion(ASTContext &ctx,
          "must be given a contextual type");
 
   // Get the outer generic signature and environment.
-  auto *genericEnv = cast<ArchetypeType>(expansionType.getCountType())
-                         ->getGenericEnvironment();
+  auto countArchetype = cast<ArchetypeType>(expansionType.getCountType());
+  auto *genericEnv = countArchetype->getGenericEnvironment();
   auto subMap = genericEnv->getForwardingSubstitutionMap();
 
+  auto countType = cast<GenericTypeParamType>(
+      countArchetype->getInterfaceType()->getCanonicalType());
+
   auto genericSig = genericEnv->getGenericSignature().getCanonicalSignature();
   // Create an opened element signature and environment.
   auto elementSig = ctx.getOpenedElementSignature(
-      genericSig, expansionType.getCountType());
+      genericSig, countType);
   auto *elementEnv = GenericEnvironment::forOpenedElement(
-      elementSig, UUID::fromTime(), expansionType.getCountType(), subMap);
+      elementSig, UUID::fromTime(), countType, subMap);
 
   return {elementEnv, elementSig};
 }

lib/SIL/IR/SILInstructions.cpp

@@ -2379,7 +2379,6 @@ OpenPackElementInst *OpenPackElementInst::create(
                                      PackType *packSubstitution) {
     collector.collect(packSubstitution->getCanonicalType());
   });
-  collector.collect(env->getOpenedElementShapeClass());
   collector.addTo(typeDependentOperands, F);
 
   SILType type = SILType::getSILTokenType(F.getASTContext());

lib/SIL/IR/SILPrinter.cpp

@@ -2312,7 +2312,13 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
           << " of " << subs.getGenericSignature()
           << " at ";
     printSubstitutions(subs);
-    *this << ", shape $" << env->getOpenedElementShapeClass()
+    // The shape class in the opened environment is a canonical interface
+    // type, which won't resolve in the generic signature we just printed.
+    // Map it back to the sugared generic parameter.
+    auto sugaredShapeClass =
+      subs.getGenericSignature()->getSugaredType(
+        env->getOpenedElementShapeClass());
+    *this << ", shape $" << sugaredShapeClass
           << ", uuid \"" << env->getOpenedElementUUID() << "\"";
   }
   void visitPackElementGetInst(PackElementGetInst *I) {

lib/SIL/Parser/ParseSIL.cpp

@@ -3387,14 +3387,23 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
     // Parse the shape class that should be opened.  This is a contextual
     // type within the signature we just parsed.
     CanType shapeClass;
+    SourceLoc shapeClassLoc;
     if (!P.consumeIf(tok::comma) ||
         parseVerbatim("shape") ||
         P.parseToken(tok::sil_dollar,
                      diag::expected_tok_in_sil_instr, "$") ||
-        parseASTType(shapeClass, openedGenericsSig, openedGenerics,
-                     /*wantContextualType*/ true))
+        parseASTType(shapeClass, shapeClassLoc, openedGenericsSig,
+                     openedGenerics, /*wantContextualType*/ true))
       return true;
 
+    // Map it out of context.  It should be a type pack parameter.
+    shapeClass = shapeClass->mapTypeOutOfContext()->getCanonicalType();
+    auto shapeParam = dyn_cast<GenericTypeParamType>(shapeClass);
+    if (!shapeParam || !shapeParam->isParameterPack()) {
+      P.diagnose(shapeClassLoc, diag::opened_shape_class_not_pack_param);
+      return true;
+    }
+
     // Parse the UUID for the opening.
     UUID uuid;
     if (!P.consumeIf(tok::comma) ||
@@ -3406,10 +3415,10 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
     // the opened elements to the signature we parsed above.
     auto openedElementSig =
       P.Context.getOpenedElementSignature(
-        openedGenericsSig.getCanonicalSignature(), shapeClass);
+        openedGenericsSig.getCanonicalSignature(), shapeParam);
 
     auto openedEnv = GenericEnvironment::forOpenedElement(openedElementSig,
-                         uuid, shapeClass, openedSubMap);
+                         uuid, shapeParam, openedSubMap);
 
     auto openInst = B.createOpenPackElement(InstLoc, Val, openedEnv);
     ResultVal = openInst;

lib/Sema/ConstraintSystem.cpp

@@ -667,13 +667,16 @@ ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator,
       !shapeClass->isEqual(uuidAndShape.second))
     return nullptr;
 
+  auto shapeParam = cast<GenericTypeParamType>(
+      shapeClass->mapTypeOutOfContext()->getCanonicalType());
+
   auto &ctx = getASTContext();
   auto elementSig = ctx.getOpenedElementSignature(
-      DC->getGenericSignatureOfContext().getCanonicalSignature(), shapeClass);
+      DC->getGenericSignatureOfContext().getCanonicalSignature(), shapeParam);
   auto *contextEnv = DC->getGenericEnvironmentOfContext();
   auto contextSubs = contextEnv->getForwardingSubstitutionMap();
   return GenericEnvironment::forOpenedElement(elementSig, uuidAndShape.first,
-                                              shapeClass, contextSubs);
+                                              shapeParam, contextSubs);
 }
 
 /// Extend the given depth map by adding depths for all of the subexpressions

lib/Serialization/Deserialization.cpp

@@ -1292,7 +1292,8 @@ Expected<GenericEnvironment *> ModuleFile::getGenericEnvironmentChecked(
   case GenericEnvironmentKind::OpenedElement:
     genericEnv = GenericEnvironment::forOpenedElement(
         parentSigOrError.get(), UUID::fromTime(),
-        existentialOrShapeTypeOrError.get()->getCanonicalType(),
+        cast<GenericTypeParamType>(
+          existentialOrShapeTypeOrError.get()->getCanonicalType()),
         contextSubsOrError.get());
   }