[GenericEnvironment] Add helper methods to map pack interface types to

element archetypes, and element interface types to pack archetypes.

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

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/Sema/CSApply.cpp

@@ -7095,32 +7095,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 = genericParam->asScalar(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 = elementParam->asParameterPack(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

@@ -8432,7 +8432,6 @@ ConstraintSystem::SolutionKind
 ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
                                                   TypeMatchOptions flags,
                                                   ConstraintLocatorBuilder locator) {
-  ASTContext &ctx = getASTContext();
   auto elementType = simplifyType(first, flags);
   auto *loc = getConstraintLocator(locator);
 
@@ -8449,17 +8448,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 = elementParam->asParameterPack(ctx);
-    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,26 +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 genericParam->asScalar(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);