Merge pull request #61737 from hborla/generic-environment-helpers

[GenericEnvironment] Add helper methods for mapping pack types to opened element types and vice versa.

Author: hborla

include/swift/AST/GenericEnvironment.h

@@ -239,6 +239,14 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   /// Map a generic parameter type to a contextual type.
   Type mapTypeIntoContext(GenericTypeParamType *type) const;
 
+  /// Map a type containing parameter packs to a contextual type
+  /// in the opened element generic context.
+  Type mapPackTypeIntoElementContext(Type type) const;
+
+  /// Map a type containing pack element type parameters to a contextual
+  /// type in the pack generic context.
+  Type mapElementTypeIntoPackContext(Type type) const;
+
   /// Map the given SIL interface type to a contextual type.
   ///
   /// This operation will also reabstract dependent types according to the

include/swift/AST/Types.h

@@ -6161,6 +6161,14 @@ 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

@@ -5601,20 +5601,14 @@ ASTContext::getOpenedElementSignature(CanGenericSignature baseGenericSig) {
   SmallVector<GenericTypeParamType *, 2> genericParams;
   SmallVector<Requirement, 2> requirements;
 
-  auto eraseParameterPack = [&](GenericTypeParamType *paramType) {
-    return GenericTypeParamType::get(
-        /*isParameterPack=*/false, paramType->getDepth(),
-        paramType->getIndex(), *this);
-  };
-
   for (auto paramType : baseGenericSig.getGenericParams()) {
-    genericParams.push_back(eraseParameterPack(paramType));
+    genericParams.push_back(paramType->asScalar(*this));
   }
 
   auto eraseParameterPackRec = [&](Type type) -> Type {
     return type.transformRec([&](Type t) -> Optional<Type> {
       if (auto *paramType = t->getAs<GenericTypeParamType>())
-        return Type(eraseParameterPack(paramType));
+        return Type(paramType->asScalar(*this));
       return None;
     });
   };

lib/AST/GenericEnvironment.cpp

@@ -494,6 +494,48 @@ Type GenericEnvironment::mapTypeIntoContext(GenericTypeParamType *type) const {
   return result;
 }
 
+Type
+GenericEnvironment::mapPackTypeIntoElementContext(Type type) const {
+  assert(getKind() == Kind::OpenedElement);
+  assert(!type->hasArchetype());
+
+  auto sig = getGenericSignature();
+  ASTContext &ctx = sig->getASTContext();
+  QueryInterfaceTypeSubstitutions substitutions(this);
+
+  // Map the interface type to the element type by stripping
+  // away the isParameterPack bit before mapping type parameters
+  // to archetypes.
+  return type.subst([&](SubstitutableType *type) {
+    auto *genericParam = type->getAs<GenericTypeParamType>();
+    if (!genericParam)
+      return Type();
+
+    return substitutions(genericParam->asScalar(ctx));
+  }, LookUpConformanceInSignature(sig.getPointer()));
+}
+
+Type
+GenericEnvironment::mapElementTypeIntoPackContext(Type type) const {
+  assert(getKind() == Kind::Primary);
+  assert(!type->hasArchetype());
+
+  auto sig = getGenericSignature();
+  ASTContext &ctx = sig->getASTContext();
+  QueryInterfaceTypeSubstitutions substitutions(this);
+
+  // Map element archetypes to the pack archetypes by converting
+  // element types to interface types and adding the isParameterPack
+  // bit. Then, map type parameters to archetypes.
+  return type.subst([&](SubstitutableType *type) {
+    auto *genericParam = type->getAs<GenericTypeParamType>();
+    if (!genericParam)
+      return Type();
+
+    return substitutions(genericParam->asParameterPack(ctx));
+  }, LookUpConformanceInSignature(sig.getPointer()));
+}
+
 SubstitutionMap GenericEnvironment::getForwardingSubstitutionMap() const {
   auto genericSig = getGenericSignature();
   return SubstitutionMap::get(genericSig,

lib/AST/ParameterPack.cpp

@@ -55,6 +55,20 @@ bool GenericTypeParamType::isParameterPack() const {
          GenericTypeParamType::TYPE_SEQUENCE_BIT;
 }
 
+GenericTypeParamType *
+GenericTypeParamType::asParameterPack(ASTContext &ctx) const {
+  return GenericTypeParamType::get(/*isParameterPack*/true,
+                                   getDepth(), getIndex(),
+                                   ctx);
+}
+
+GenericTypeParamType *
+GenericTypeParamType::asScalar(ASTContext &ctx) const {
+  return GenericTypeParamType::get(/*isParameterPack*/false,
+                                   getDepth(), getIndex(),
+                                   ctx);
+}
+
 /// G<{X1, ..., Xn}, {Y1, ..., Yn}>... => {G<X1, Y1>, ..., G<Xn, Yn>}...
 PackExpansionType *PackExpansionType::expand() {
   auto countType = getCountType();
@@ -278,4 +292,4 @@ unsigned ParameterList::getOrigParamIndex(SubstitutionMap subMap,
   subMap.dump(llvm::errs());
   dump(llvm::errs());
   abort();
-}
+}

lib/Sema/CSApply.cpp

@@ -7089,37 +7089,15 @@ Expr *ExprRewriter::coerceToType(Expr *expr, Type toType,
     auto toExpansionType = toType->getAs<PackExpansionType>();
     auto *expansion = dyn_cast<PackExpansionExpr>(expr);
 
-    // FIXME: Duplicated code from `PackReferenceFinder` in PreCheckExpr.
-    auto toElementType = toExpansionType->getPatternType();
-    toElementType =
-        toElementType->mapTypeOutOfContext().transform([&](Type type) -> Type {
-          auto *genericParam = type->getAs<GenericTypeParamType>();
-          if (!genericParam || !genericParam->isParameterPack())
-            return type;
-
-          auto param = GenericTypeParamType::get(/*isParameterPack*/false,
-                                                 genericParam->getDepth(),
-                                                 genericParam->getIndex(), ctx);
-          return expansion->getGenericEnvironment()->mapTypeIntoContext(param);
-        });
+    auto *elementEnv = expansion->getGenericEnvironment();
+    auto toElementType = elementEnv->mapPackTypeIntoElementContext(
+        toExpansionType->getPatternType()->mapTypeOutOfContext());
 
     auto *pattern = coerceToType(expansion->getPatternExpr(),
                                  toElementType, locator);
-
-    // FIXME: Duplicated code from `visitPackExpansionExpr` in CSGen.
-    auto patternType =
-        cs.getType(pattern).transform([&](Type type) -> Type {
-          auto *element = type->getAs<ElementArchetypeType>();
-          if (!element)
-            return type;
-
-          auto *elementParam = element->mapTypeOutOfContext()->getAs<GenericTypeParamType>();
-          auto *pack = GenericTypeParamType::get(/*isParameterPack*/true,
-                                                 elementParam->getDepth(),
-                                                 elementParam->getIndex(),
-                                                 ctx);
-          return cs.DC->mapTypeIntoContext(pack);
-        });
+    auto *packEnv = cs.DC->getGenericEnvironmentOfContext();
+    auto patternType = packEnv->mapElementTypeIntoPackContext(
+        toElementType->mapTypeOutOfContext());
     auto shapeType = toExpansionType->getCountType();
     auto expansionTy = PackExpansionType::get(patternType, shapeType);
 

lib/Sema/CSSimplify.cpp

@@ -8464,20 +8464,11 @@ ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
 
   // Replace opened element archetypes with pack archetypes
   // for the resulting type of the pack expansion.
-  auto patternType = elementType.transform([&](Type type) -> Type {
-    auto *element = type->getAs<ElementArchetypeType>();
-    if (!element)
-      return type;
-
-    auto *elementParam = element->mapTypeOutOfContext()->getAs<GenericTypeParamType>();
-    auto *pack = GenericTypeParamType::get(/*isParameterPack*/true,
-                                           elementParam->getDepth(),
-                                           elementParam->getIndex(),
-                                           this->getASTContext());
-    return this->DC->mapTypeIntoContext(pack);
-  });
-
+  auto *environment = DC->getGenericEnvironmentOfContext();
+  auto patternType = environment->mapElementTypeIntoPackContext(
+      elementType->mapTypeOutOfContext());
   addConstraint(ConstraintKind::Bind, second, patternType, locator);
+
   return SolutionKind::Solved;
 }
 

lib/Sema/PreCheckExpr.cpp

@@ -436,28 +436,15 @@ static Expr *getPackExpansion(DeclContext *dc, Expr *expr, SourceLoc opLoc) {
         if (auto expansionType = decl->getType()->getAs<PackExpansionType>()) {
           auto sourceRange = declRef->getSourceRange();
 
-          // Map the pattern interface type to the element interface type
-          // by making all type parameter packs scalar type parameters.
-          auto patternType = expansionType->getPatternType();
-          auto elementInterfaceType =
-              patternType->mapTypeOutOfContext().transform([&](Type type) -> Type {
-                auto *genericParam = type->getAs<GenericTypeParamType>();
-                if (!genericParam || !genericParam->isParameterPack())
-                  return type;
-
-                return GenericTypeParamType::get(/*isParameterPack*/false,
-                                                 genericParam->getDepth(),
-                                                 genericParam->getIndex(), ctx);
-              });
-
-          // Map the element interface type into the context of the opened
+          // Map the pattern interface type into the context of the opened
           // element signature.
           if (!environment) {
             auto sig = ctx.getOpenedElementSignature(
                 dc->getGenericSignatureOfContext().getCanonicalSignature());
             environment = GenericEnvironment::forOpenedElement(sig, UUID::fromTime());
           }
-          auto elementType = environment->mapTypeIntoContext(elementInterfaceType);
+          auto elementType = environment->mapPackTypeIntoElementContext(
+              expansionType->getPatternType()->mapTypeOutOfContext());
 
           auto *opaqueValue = new (ctx) OpaqueValueExpr(sourceRange, elementType);
           opaqueValues.push_back(opaqueValue);