Merge pull request #62228 from hborla/pack-element-generic-environment

[GenericEnvironment] Include original parameter packs in opened pack element signatures.

Author: hborla

include/swift/AST/GenericEnvironment.h

@@ -64,9 +64,10 @@ struct OpenedExistentialEnvironmentData {
   UUID uuid;
 };
 
-/// Extra data in a generic environment for an opened element.
+/// Extra data in a generic environment for an opened pack element.
 struct OpenedElementEnvironmentData {
   UUID uuid;
+  SubstitutionMap outerSubstitutions;
 };
 
 /// Describes the mapping between archetypes and interface types for the
@@ -139,7 +140,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,
+                              SubstitutionMap outerSubs);
 
   friend ArchetypeType;
   friend QueryInterfaceTypeSubstitutions;
@@ -181,9 +183,23 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   /// create a generic environment.
   SubstitutionMap getOpaqueSubstitutions() const;
 
+  /// Retrieve the substitutions for the outer generic parameters of an
+  /// opened pack element generic environment.
+  SubstitutionMap getPackElementContextSubstitutions() const;
+
   /// Retrieve the UUID for an opened element environment.
   UUID getOpenedElementUUID() const;
 
+  using PackElementBinding =
+      std::pair<ElementArchetypeType *, PackArchetypeType *>;
+
+  /// Retrieve the bindings for the opened pack element archetypes in this
+  /// generic environment to the pack archetypes that contain them.
+  ///
+  /// \param bindings The vector to populate with the pack element bindings.
+  void getPackElementBindings(
+      SmallVectorImpl<PackElementBinding> &bindings) const;
+
   /// Create a new, primary generic environment.
   static GenericEnvironment *forPrimary(GenericSignature signature);
 
@@ -206,8 +222,11 @@ 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 outerSubs The substitution map containing archetypes from the
+  /// outer generic context.
   static GenericEnvironment *
-  forOpenedElement(GenericSignature signature, UUID uuid);
+  forOpenedElement(GenericSignature signature, UUID uuid,
+                   SubstitutionMap outerSubs);
 
   /// Make vanilla new/delete illegal.
   void *operator new(size_t Bytes) = delete;
@@ -219,8 +238,9 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
     return Mem; 
   }
 
-  /// For an opaque archetype environment, apply the substitutions.
-  Type maybeApplyOpaqueTypeSubstitutions(Type type) const;
+  /// For an opaque or pack element archetype environment, apply the
+  /// substitutions.
+  Type maybeApplyOuterContextSubstitutions(Type type) const;
 
   /// Compute the canonical interface type within this environment.
   Type getCanonicalInterfaceType(Type interfaceType);

include/swift/AST/Types.h

@@ -6188,14 +6188,6 @@ class GenericTypeParamType : public SubstitutableType {
   /// \endcode
   bool isParameterPack() const;
 
-  /// Returns a new GenericTypeParamType with the same depth and index
-  /// as this one, with the type parameter pack bit set.
-  GenericTypeParamType *asParameterPack(ASTContext &ctx) const;
-
-  /// Returns a new GenericTypeParamType with the same depth and index
-  /// as this one, removing the type parameter pack bit.
-  GenericTypeParamType *asScalar(ASTContext &ctx) const;
-
   // Implement isa/cast/dyncast/etc.
   static bool classof(const TypeBase *T) {
     return T->getKind() == TypeKind::GenericTypeParam;

lib/AST/ASTContext.cpp

@@ -5016,7 +5016,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,
+                                     SubstitutionMap outerSubs) {
   auto &ctx = signature->getASTContext();
 
   auto &openedElementEnvironments =
@@ -5038,7 +5039,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,
+                                                  outerSubs);
 
   openedElementEnvironments[uuid] = genericEnv;
 
@@ -5621,41 +5623,76 @@ ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig) {
   }
 #endif
 
-  SmallVector<GenericTypeParamType *, 2> genericParams;
+  // The pack element signature includes all type parameters and requirements
+  // from the outer context, plus a new set of type parameters representing
+  // open pack elements and their corresponding element requirements.
+
+  llvm::SmallMapVector<GenericTypeParamType *,
+                       GenericTypeParamType *, 2> packElementParams;
+  SmallVector<GenericTypeParamType *, 2> genericParams(
+      baseGenericSig.getGenericParams().begin(), baseGenericSig.getGenericParams().end());
   SmallVector<Requirement, 2> requirements;
 
+  auto packElementDepth =
+      baseGenericSig.getInnermostGenericParams().front()->getDepth() + 1;
+
   for (auto paramType : baseGenericSig.getGenericParams()) {
-    genericParams.push_back(paramType->asScalar(*this));
+    if (!paramType->isParameterPack())
+      continue;
+
+    auto *elementParam = GenericTypeParamType::get(/*isParameterPack*/false,
+                                                   packElementDepth,
+                                                   packElementParams.size(),
+                                                   *this);
+    genericParams.push_back(elementParam);
+    packElementParams[paramType] = elementParam;
   }
 
   auto eraseParameterPackRec = [&](Type type) -> Type {
     return type.transformRec([&](Type t) -> Optional<Type> {
-      if (auto *paramType = t->getAs<GenericTypeParamType>())
-        return Type(paramType->asScalar(*this));
+      if (auto *paramType = t->getAs<GenericTypeParamType>()) {
+        if (paramType->isParameterPack()) {
+          return Type(packElementParams[paramType]);
+        }
+
+        return t;
+      }
       return None;
     });
   };
 
   for (auto requirement : baseGenericSig.getRequirements()) {
+    requirements.push_back(requirement);
+
+    // If this requirement contains parameter packs, create a new requirement
+    // for the corresponding pack element.
     switch (requirement.getKind()) {
     case RequirementKind::SameShape:
       // Drop same-shape requirements from the element signature.
       break;
     case RequirementKind::Conformance:
     case RequirementKind::Superclass:
-    case RequirementKind::SameType:
-      requirements.emplace_back(
-          requirement.getKind(),
-          eraseParameterPackRec(requirement.getFirstType()),
-          eraseParameterPackRec(requirement.getSecondType()));
+    case RequirementKind::SameType: {
+      auto firstType = eraseParameterPackRec(requirement.getFirstType());
+      auto secondType = eraseParameterPackRec(requirement.getSecondType());
+      if (firstType->isEqual(requirement.getFirstType()) &&
+          secondType->isEqual(requirement.getSecondType()))
+        break;
+
+      requirements.emplace_back(requirement.getKind(),
+                                firstType, secondType);
       break;
-    case RequirementKind::Layout:
-      requirements.emplace_back(
-          requirement.getKind(),
-          eraseParameterPackRec(requirement.getFirstType()),
-          requirement.getLayoutConstraint());
+    }
+    case RequirementKind::Layout: {
+      auto firstType = eraseParameterPackRec(requirement.getFirstType());
+      if (firstType->isEqual(requirement.getFirstType()))
+        break;
+
+      requirements.emplace_back(requirement.getKind(), firstType,
+                                requirement.getLayoutConstraint());
       break;
     }
+    }
   }
 
   auto elementSig = buildGenericSignature(

lib/AST/ASTVerifier.cpp

@@ -196,7 +196,7 @@ class Verifier : public ASTWalker {
   SmallVector<ScopeLike, 4> Scopes;
 
   /// The stack of generic contexts.
-  using GenericLike = llvm::PointerUnion<DeclContext *, GenericSignature>;
+  using GenericLike = llvm::PointerUnion<DeclContext *, GenericEnvironment *>;
   SmallVector<GenericLike, 2> Generics;
 
   /// The stack of optional evaluations active at this point.
@@ -634,10 +634,21 @@ class Verifier : public ASTWalker {
 
           auto genericCtx = Generics.back();
           GenericSignature genericSig;
-          if (auto *genericDC = genericCtx.dyn_cast<DeclContext *>())
+          if (auto *genericDC = genericCtx.dyn_cast<DeclContext *>()) {
             genericSig = genericDC->getGenericSignatureOfContext();
-          else
-            genericSig = genericCtx.get<GenericSignature>();
+          } else {
+            auto *genericEnv = genericCtx.get<GenericEnvironment *>();
+            genericSig = genericEnv->getGenericSignature();
+
+            // Check whether this archetype is a substitution from the
+            // outer generic context of an opened element environment.
+            if (genericEnv->getKind() == GenericEnvironment::Kind::OpenedElement) {
+              auto contextSubs = genericEnv->getPackElementContextSubstitutions();
+              QuerySubstitutionMap isInContext{contextSubs};
+              if (isInContext(root->getInterfaceType()->castTo<GenericTypeParamType>()))
+                return false;
+            }
+          }
 
           if (genericSig.getPointer() != archetypeSig.getPointer()) {
             Out << "Archetype " << root->getString() << " not allowed "
@@ -824,7 +835,7 @@ class Verifier : public ASTWalker {
       if (!shouldVerify(cast<Expr>(expr)))
         return false;
 
-      Generics.push_back(expr->getGenericEnvironment()->getGenericSignature());
+      Generics.push_back(expr->getGenericEnvironment());
 
       for (auto *placeholder : expr->getOpaqueValues()) {
         assert(!OpaqueValues.count(placeholder));
@@ -837,8 +848,8 @@ class Verifier : public ASTWalker {
     void verifyCheckedAlways(PackExpansionExpr *E) {
       // Remove the element generic environment before verifying
       // the pack expansion type, which contains pack archetypes.
-      assert(Generics.back().get<GenericSignature>().getPointer() ==
-             E->getGenericEnvironment()->getGenericSignature().getPointer());
+      assert(Generics.back().get<GenericEnvironment *>() ==
+             E->getGenericEnvironment());
       Generics.pop_back();
       verifyCheckedAlwaysBase(E);
     }