AST: Rewrite GenericEnvironment::mapElementTypeIntoPackContext()

We don't need to build a DenseMap here. More importantly, this
changes the logic to avoid calling mapTypeOutOfContext() on
element archetypes, instead doing the mapping directly.

Author: slavapestov

lib/AST/GenericEnvironment.cpp

@@ -718,52 +718,58 @@ GenericEnvironment::mapElementTypeIntoPackContext(Type type) const {
   // generic environment.
   assert(type->hasElementArchetype());
 
-  ElementArchetypeType *element = nullptr;
-  type.visit([&](Type type) {
-    auto archetype = type->getAs<ElementArchetypeType>();
-    if (!element && archetype)
-      element = archetype;
-  });
+  GenericEnvironment *elementEnv = nullptr;
 
-  auto sig = getGenericSignature();
-  auto *elementEnv = element->getGenericEnvironment();
-  auto shapeClass = elementEnv->getOpenedElementShapeClass();
-  QueryInterfaceTypeSubstitutions substitutions(this);
+  // Map element archetypes to interface types in the element generic
+  // environment's signature.
+  type = type.subst(
+    [&](SubstitutableType *type) -> Type {
+      auto *archetype = cast<ArchetypeType>(type);
 
-  type = type->mapTypeOutOfContext();
+      if (isa<OpenedArchetypeType>(archetype))
+        return archetype;
 
-  auto interfaceType = element->getInterfaceType();
+      if (isa<ElementArchetypeType>(archetype)) {
+        assert(!elementEnv ||
+               elementEnv == archetype->getGenericEnvironment());
+        elementEnv = archetype->getGenericEnvironment();
+      }
 
-  llvm::SmallDenseMap<GenericParamKey, GenericTypeParamType *>
-      packParamForElement;
-  auto elementDepth = interfaceType->getRootGenericParam()->getDepth();
+      return archetype->getInterfaceType();
+    },
+    MakeAbstractConformanceForGenericType(),
+    SubstFlags::AllowLoweredTypes |
+    SubstFlags::PreservePackExpansionLevel);
 
+  auto shapeClass = elementEnv->getOpenedElementShapeClass();
+
+  llvm::SmallVector<GenericTypeParamType *, 2> members;
+  auto elementDepth = elementEnv->getGenericSignature()->getMaxDepth();
+
+  auto sig = getGenericSignature();
   for (auto *genericParam : sig.getGenericParams()) {
     if (!genericParam->isParameterPack())
       continue;
 
     if (!sig->haveSameShape(genericParam, shapeClass))
       continue;
 
-    GenericParamKey elementKey(/*isParameterPack*/false,
-                               /*depth*/elementDepth,
-                               /*index*/packParamForElement.size());
-    packParamForElement[elementKey] = genericParam;
+    members.push_back(genericParam);
   }
 
-  // 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.
+  // Map element interface types to pack archetypes.
+  QueryInterfaceTypeSubstitutions mapIntoContext(this);
   return type.subst(
       [&](SubstitutableType *type) {
         auto *genericParam = type->getAs<GenericTypeParamType>();
         if (!genericParam)
           return Type();
 
-        if (auto *packParam = packParamForElement[{genericParam}])
-          return substitutions(packParam);
-
-        return substitutions(genericParam);
+        if (genericParam->getDepth() == elementDepth) {
+          genericParam = members[genericParam->getIndex()];
+          assert(genericParam->isParameterPack());
+        }
+        return mapIntoContext(genericParam);
       },
       LookUpConformanceInSignature(sig.getPointer()),
       SubstFlags::PreservePackExpansionLevel);