Fix the mapping of pack types into opened element environments.

First, we need this to work on both lowered and unlowered types,
so Type::subst is problematic: it'll assert if it sees a type
like SILFunctionType.  In this case, the substitution is simple
enough that that's never a problem, but Type::subst doesn't know
that, and the assertion is generally a good one.

Second, we need this to not recurse into nested pack expansions.

Third, we need this to not mess around with any existing element
archetypes we might see in the type, so mapping in and out of
context is not really okay.

Fortunately, because we're mapping between structures (pack and
element archetypes) that are guaranteed to have the same
constraints, this transformation is really easy and we can just
do it with transformRec.

Author: rjmccall

include/swift/AST/GenericEnvironment.h

@@ -294,10 +294,26 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
 
   /// Map a contextual type containing parameter packs to a contextual
   /// type in the opened element generic context.
+  ///
+  /// This operation only makes sense if the generic environment that the
+  /// pack archetypes are contextual in matches the generic signature
+  /// of this environment.  That will be true for opened element
+  /// environments coming straight out of the type checker, such as
+  /// the one in a PackExpansionExpr, or opened element environments
+  /// created directly from the current environment.  It is not
+  /// reliable for opened element environments in arbitrary SIL functions.
   Type mapContextualPackTypeIntoElementContext(Type type) const;
 
   /// Map a contextual type containing parameter packs to a contextual
   /// type in the opened element generic context.
+  ///
+  /// This operation only makes sense if the generic environment that the
+  /// pack archetypes are contextual in matches the generic signature
+  /// of this environment.  That will be true for opened element
+  /// environments coming straight out of the type checker, such as
+  /// the one in a PackExpansionExpr, or opened element environments
+  /// created directly from the current environment.  It is not
+  /// reliable for opened element environments in arbitrary SIL functions.
   CanType mapContextualPackTypeIntoElementContext(CanType type) const;
 
   /// Map a type containing pack element type parameters to a contextual

lib/AST/GenericEnvironment.cpp

@@ -148,6 +148,53 @@ UUID GenericEnvironment::getOpenedElementUUID() const {
   return getTrailingObjects<OpenedElementEnvironmentData>()->uuid;
 }
 
+namespace {
+
+struct FindOpenedElementParam {
+  ArrayRef<Type> openedPacks;
+  TypeArrayView<GenericTypeParamType> packElementParams;
+
+  FindOpenedElementParam(const GenericEnvironment *env,
+                         ArrayRef<Type> openedPacks)
+    : openedPacks(openedPacks),
+      packElementParams(
+        env->getGenericSignature().getInnermostGenericParams()) {
+    assert(openedPacks.size() == packElementParams.size());
+  }
+
+  GenericTypeParamType *operator()(Type packParam) {
+    for (auto i : indices(openedPacks)) {
+      if (openedPacks[i]->isEqual(packParam))
+        return packElementParams[i];
+    }
+    llvm_unreachable("parameter was not an opened pack parameter");
+  }
+};
+
+struct FindElementArchetypeForOpenedPackParam {
+  FindOpenedElementParam findElementParam;
+  QueryInterfaceTypeSubstitutions getElementArchetype;
+
+  FindElementArchetypeForOpenedPackParam(const GenericEnvironment *env,
+                                         ArrayRef<Type> openedPacks)
+    : findElementParam(env, openedPacks), getElementArchetype(env) {}
+
+
+  ElementArchetypeType *operator()(Type interfaceType) {
+    assert(interfaceType->isTypeParameter());
+    if (auto member = interfaceType->getAs<DependentMemberType>()) {
+      auto baseArchetype = (*this)(member->getBase());
+      return baseArchetype->getNestedType(member->getAssocType())
+               ->castTo<ElementArchetypeType>();
+    }
+    assert(interfaceType->is<GenericTypeParamType>());
+    return getElementArchetype(findElementParam(interfaceType))
+             ->castTo<ElementArchetypeType>();
+  }
+};
+
+}
+
 void GenericEnvironment::forEachPackElementArchetype(
           llvm::function_ref<void(ElementArchetypeType *)> function) const {
   auto packElements = getGenericSignature().getInnermostGenericParams();
@@ -613,25 +660,51 @@ Type GenericEnvironment::mapTypeIntoContext(GenericTypeParamType *type) const {
 
 Type
 GenericEnvironment::mapContextualPackTypeIntoElementContext(Type type) const {
+  assert(getKind() == Kind::OpenedElement);
+  assert(!type->hasTypeParameter() && "expected contextual type");
+
   if (!type->hasArchetype()) return type;
 
-  // FIXME: this is potentially wrong if there are multiple
-  // openings in play at once, because we really shouldn't touch
-  // other element archetypes.
-  return mapPackTypeIntoElementContext(type->mapTypeOutOfContext());
+  auto sig = getGenericSignature();
+  auto shapeClass = getOpenedElementShapeClass();
+
+  FindElementArchetypeForOpenedPackParam
+    findElementArchetype(this, getOpenedPackParams());
+
+  return type.transformRec([&](TypeBase *ty) -> Optional<Type> {
+    // We're only directly substituting pack archetypes.
+    auto archetype = ty->getAs<PackArchetypeType>();
+    if (!archetype) {
+      // Don't recurse into nested pack expansions.
+      if (ty->is<PackExpansionType>())
+        return Type(ty);
+
+      // Recurse into any other type.
+      return None;
+    }
+
+    auto rootArchetype = cast<PackArchetypeType>(archetype->getRoot());
+
+    // TODO: assert that the generic environment of the pack archetype
+    // matches the signature that was originally opened to make this
+    // environment.  Unfortunately, that isn't a trivial check because of
+    // the extra opened-element parameters.
+
+    // If the archetype isn't the shape that was opened by this
+    // environment, ignore it.
+    auto rootParam = cast<GenericTypeParamType>(
+      rootArchetype->getInterfaceType().getPointer());
+    assert(rootParam->isParameterPack());
+    if (!sig->haveSameShape(rootParam, shapeClass))
+      return Type(ty);
+
+    return Type(findElementArchetype(archetype->getInterfaceType()));
+  });
 }
 
 CanType
 GenericEnvironment::mapContextualPackTypeIntoElementContext(CanType type) const {
-  if (!type->hasArchetype()) return type;
-
-  // FIXME: this is potentially wrong if there are multiple
-  // openings in play at once, because we really shouldn't touch
-  // other element archetypes.
-  // FIXME: if we do this properly, there's no way for this rewrite
-  // to produce a non-canonical type.
-  return mapPackTypeIntoElementContext(type->mapTypeOutOfContext())
-           ->getCanonicalType();
+  return CanType(mapContextualPackTypeIntoElementContext(Type(type)));
 }
 
 Type
@@ -641,40 +714,34 @@ GenericEnvironment::mapPackTypeIntoElementContext(Type type) const {
 
   auto sig = getGenericSignature();
   auto shapeClass = getOpenedElementShapeClass();
-  QueryInterfaceTypeSubstitutions substitutions(this);
 
-  llvm::SmallDenseMap<GenericParamKey,
-                      GenericTypeParamType *> elementParamForPack;
-  auto packElements = sig.getInnermostGenericParams();
-  auto elementDepth = packElements.front()->getDepth();
-
-  for (auto *genericParam : sig.getGenericParams()) {
-    if (genericParam->getDepth() == elementDepth)
-      break;
-
-    if (!genericParam->isParameterPack())
-      continue;
-
-    if (!sig->haveSameShape(genericParam, shapeClass))
-      continue;
-
-    auto elementIndex = elementParamForPack.size();
-    elementParamForPack[{genericParam}] = packElements[elementIndex];
-  }
+  FindElementArchetypeForOpenedPackParam
+    findElementArchetype(this, getOpenedPackParams());
 
   // 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 type.transformRec([&](TypeBase *ty) -> Optional<Type> {
+    // We're only directly substituting pack parameters.
+    if (!ty->isTypeParameter()) {
+      // Don't recurse into nested pack expansions; just map it into
+      // context.
+      if (ty->is<PackExpansionType>())
+        return mapTypeIntoContext(ty);
+
+      // Recurse into any other type.
+      return None;
+    }
 
-    if (auto *elementParam = elementParamForPack[{genericParam}])
-      return substitutions(elementParam);
+    // Just do normal mapping for types that are not rooted in
+    // opened type parameters.
+    auto rootParam = ty->getRootGenericParam();
+    if (!rootParam->isParameterPack() ||
+        !sig->haveSameShape(rootParam, shapeClass))
+      return mapTypeIntoContext(ty);
 
-    return substitutions(genericParam);
-  }, LookUpConformanceInSignature(sig.getPointer()));
+    return Type(findElementArchetype(ty));
+  });
 }
 
 Type