[ConstraintSystem] Open pack element types using a PackElementOf constraint

when resolving pack reference type reprs inside pack expansion expressions.

Author: hborla

include/swift/Sema/ConstraintSystem.h

@@ -6108,6 +6108,30 @@ class HandlePlaceholderType {
   }
 };
 
+/// A function object that opens a given pack type by generating a
+/// \c PackElementOf constraint.
+class OpenPackElementType {
+  ConstraintSystem &cs;
+  ConstraintLocator *locator;
+  GenericEnvironment *elementEnv;
+
+public:
+  explicit OpenPackElementType(ConstraintSystem &cs,
+                               ConstraintLocator *locator,
+                               GenericEnvironment *elementEnv)
+      : cs(cs), locator(locator), elementEnv(elementEnv) {}
+
+  Type operator()(Type packType) const {
+    auto *elementType = cs.createTypeVariable(locator, TVO_CanBindToHole);
+    auto elementLoc = cs.getConstraintLocator(locator,
+        LocatorPathElt::OpenedPackElement(elementEnv));
+
+    cs.addConstraint(ConstraintKind::PackElementOf, elementType,
+                     packType, elementLoc);
+    return elementType;
+  }
+};
+
 /// Compute the shuffle required to map from a given tuple type to
 /// another tuple type.
 ///

lib/Sema/CSGen.cpp

@@ -1416,13 +1416,20 @@ namespace {
     Type
     resolveTypeReferenceInExpression(TypeRepr *repr, TypeResolverContext resCtx,
                                      const ConstraintLocatorBuilder &locator) {
+      TypeResolutionOptions options(resCtx);
+
       // Introduce type variables for unbound generics.
       const auto genericOpener = OpenUnboundGenericType(CS, locator);
       const auto placeholderHandler = HandlePlaceholderType(CS, locator);
-      // FIXME: Open pack elements with a PackElementOf constraint.
       OpenPackElementFn packElementOpener = nullptr;
+      if (!PackElementEnvironments.empty()) {
+        options |= TypeResolutionFlags::AllowPackReferences;
+        packElementOpener = OpenPackElementType(CS,
+            CS.getConstraintLocator(locator), PackElementEnvironments.back());
+      }
+
       const auto result = TypeResolution::resolveContextualType(
-          repr, CS.DC, resCtx, genericOpener, placeholderHandler,
+          repr, CS.DC, options, genericOpener, placeholderHandler,
           packElementOpener);
       if (result->hasError()) {
         CS.recordFix(
@@ -1662,16 +1669,20 @@ namespace {
         ArrayRef<TypeRepr *> specializationArgs) {
       // Resolve each type.
       SmallVector<Type, 2> specializationArgTypes;
-      const auto options =
+      auto options =
           TypeResolutionOptions(TypeResolverContext::InExpression);
       for (auto specializationArg : specializationArgs) {
+        OpenPackElementFn packElementOpener = nullptr;
+        if (!PackElementEnvironments.empty()) {
+          options |= TypeResolutionFlags::AllowPackReferences;
+          packElementOpener = OpenPackElementType(CS, locator, PackElementEnvironments.back());
+        }
         const auto result = TypeResolution::resolveContextualType(
             specializationArg, CurDC, options,
             // Introduce type variables for unbound generics.
             OpenUnboundGenericType(CS, locator),
             HandlePlaceholderType(CS, locator),
-            // FIXME: Open pack elements with a PackElementOf constraint.
-            /*packElementOpener*/ nullptr);
+            packElementOpener);
         if (result->hasError())
           return true;
 
@@ -1728,16 +1739,21 @@ namespace {
           // Bind the specified generic arguments to the type variables in the
           // open type.
           auto *const locator = CS.getConstraintLocator(expr);
-          const auto options =
+          auto options =
               TypeResolutionOptions(TypeResolverContext::InExpression);
           for (size_t i = 0, e = specializations.size(); i < e; ++i) {
+            OpenPackElementFn packElementOpener = nullptr;
+            if (!PackElementEnvironments.empty()) {
+              options |= TypeResolutionFlags::AllowPackReferences;
+              packElementOpener = OpenPackElementType(CS, locator, PackElementEnvironments.back());
+            }
+
             const auto result = TypeResolution::resolveContextualType(
                 specializations[i], CS.DC, options,
                 // Introduce type variables for unbound generics.
                 OpenUnboundGenericType(CS, locator),
                 HandlePlaceholderType(CS, locator),
-                // FIXME: Open pack elements with a PackElementOf constraint.
-                /*packElementOpener*/ nullptr);
+                packElementOpener);
             if (result->hasError())
               return Type();
 
@@ -3001,18 +3017,12 @@ namespace {
 
     Type visitPackElementExpr(PackElementExpr *expr) {
       auto packType = CS.getType(expr->getPackRefExpr());
-      auto *elementType =
-          CS.createTypeVariable(CS.getConstraintLocator(expr),
-                                TVO_CanBindToHole);
-      auto *elementEnv = PackElementEnvironments.back();
-      auto *elementLoc = CS.getConstraintLocator(
-          expr, LocatorPathElt::OpenedPackElement(elementEnv));
 
       // The type of a PackElementExpr is the opened pack element archetype
       // of the pack reference.
-      CS.addConstraint(ConstraintKind::PackElementOf, elementType, packType,
-                       elementLoc);
-      return elementType;
+      OpenPackElementType openPackElement(CS, CS.getConstraintLocator(expr),
+                                          PackElementEnvironments.back());
+      return openPackElement(packType);
     }
 
     Type visitDynamicTypeExpr(DynamicTypeExpr *expr) {