[ConstraintSystem] Add a PackElementOf constraint to delay mapping opened

element pattern types to pattern archetypes until after type variables are
resolved.

Author: hborla

Diff for: include/swift/Sema/Constraint.h

@@ -215,6 +215,9 @@ enum class ConstraintKind : char {
   /// Represents an AST node contained in a body of a function/closure.
   /// It only has an AST node to generate constraints and infer the type for.
   SyntacticElement,
+  /// The first type is the opened pack element type of the second type, which
+  /// is the pattern of a pack expansion type.
+  PackElementOf,
   /// Do not add new uses of this, it only exists to retain compatibility for
   /// rdar://85263844.
   ///
@@ -686,6 +689,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     case ConstraintKind::OneWayBindParam:
     case ConstraintKind::DefaultClosureType:
     case ConstraintKind::UnresolvedMemberChainBase:
+    case ConstraintKind::PackElementOf:
       return ConstraintClassification::Relational;
 
     case ConstraintKind::ValueMember:

Diff for: include/swift/Sema/ConstraintSystem.h

@@ -5501,6 +5501,18 @@ class ConstraintSystem {
                                               TypeMatchOptions flags,
                                               ConstraintLocatorBuilder locator);
 
+  /// Attempt to simplify a PackElementOf constraint.
+  ///
+  /// Solving this constraint is delayed until the element type is fully
+  /// resolved with no type variables. The element type is then mapped out
+  /// of the opened element context and into the context of the surrounding
+  /// function, effecively substituting opened element archetypes with their
+  /// corresponding pack archetypes, and bound to the second type.
+  SolutionKind
+  simplifyPackElementOfConstraint(Type first, Type second,
+                                  TypeMatchOptions flags,
+                                  ConstraintLocatorBuilder locator);
+
   /// Attempt to simplify the ApplicableFunction constraint.
   SolutionKind simplifyApplicableFnConstraint(
       Type type1, Type type2,

Diff for: lib/Sema/CSBindings.cpp

@@ -1396,6 +1396,7 @@ void PotentialBindings::infer(Constraint *constraint) {
   case ConstraintKind::SyntacticElement:
   case ConstraintKind::Conjunction:
   case ConstraintKind::BindTupleOfFunctionParams:
+  case ConstraintKind::PackElementOf:
     // Constraints from which we can't do anything.
     break;
 

Diff for: lib/Sema/CSGen.cpp

@@ -2886,21 +2886,11 @@ namespace {
         CS.setType(binding, type);
       }
 
-      // Replace opened element archetypes with pack archetypes
-      // for the resulting type of the pack expansion.
       auto elementResultType = CS.getType(expr->getPatternExpr());
-      auto patternTy = elementResultType.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(),
-                                               CS.getASTContext());
-        return CS.DC->mapTypeIntoContext(pack);
-      });
+      auto patternTy = CS.createTypeVariable(CS.getConstraintLocator(expr),
+                                             TVO_CanBindToHole);
+      CS.addConstraint(ConstraintKind::PackElementOf, elementResultType,
+                       patternTy, CS.getConstraintLocator(expr));
 
       // FIXME: Use a ShapeOf constraint here.
       Type shapeType;

Diff for: lib/Sema/CSSimplify.cpp

@@ -2424,6 +2424,7 @@ ConstraintSystem::matchTupleTypes(TupleType *tuple1, TupleType *tuple2,
   case ConstraintKind::PropertyWrapper:
   case ConstraintKind::SyntacticElement:
   case ConstraintKind::BindTupleOfFunctionParams:
+  case ConstraintKind::PackElementOf:
     llvm_unreachable("Bad constraint kind in matchTupleTypes()");
   }
 
@@ -2595,6 +2596,7 @@ static bool matchFunctionRepresentations(FunctionType::ExtInfo einfo1,
   case ConstraintKind::PropertyWrapper:
   case ConstraintKind::SyntacticElement:
   case ConstraintKind::BindTupleOfFunctionParams:
+  case ConstraintKind::PackElementOf:
     return true;
   }
 
@@ -3010,6 +3012,7 @@ ConstraintSystem::matchFunctionTypes(FunctionType *func1, FunctionType *func2,
   case ConstraintKind::PropertyWrapper:
   case ConstraintKind::SyntacticElement:
   case ConstraintKind::BindTupleOfFunctionParams:
+  case ConstraintKind::PackElementOf:
     llvm_unreachable("Not a relational constraint");
   }
 
@@ -6362,6 +6365,7 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
     case ConstraintKind::PropertyWrapper:
     case ConstraintKind::SyntacticElement:
     case ConstraintKind::BindTupleOfFunctionParams:
+    case ConstraintKind::PackElementOf:
       llvm_unreachable("Not a relational constraint");
     }
   }
@@ -8424,6 +8428,43 @@ ConstraintSystem::simplifyBindTupleOfFunctionParamsConstraint(
   return SolutionKind::Solved;
 }
 
+ConstraintSystem::SolutionKind
+ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
+                                                  TypeMatchOptions flags,
+                                                  ConstraintLocatorBuilder locator) {
+  auto elementType = simplifyType(first, flags);
+  auto *loc = getConstraintLocator(locator);
+
+  if (elementType->hasTypeVariable()) {
+    if (!flags.contains(TMF_GenerateConstraints))
+      return SolutionKind::Unsolved;
+
+    addUnsolvedConstraint(
+        Constraint::create(*this, ConstraintKind::PackElementOf,
+                           first, second, loc));
+
+    return SolutionKind::Solved;
+  }
+
+  // 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);
+  });
+
+  addConstraint(ConstraintKind::Bind, second, patternType, locator);
+  return SolutionKind::Solved;
+}
+
 static bool isForKeyPathSubscript(ConstraintSystem &cs,
                                   ConstraintLocator *locator) {
   if (!locator || !locator->getAnchor())
@@ -13786,6 +13827,9 @@ ConstraintSystem::addConstraintImpl(ConstraintKind kind, Type first,
     return simplifyBindTupleOfFunctionParamsConstraint(first, second, subflags,
                                                        locator);
 
+  case ConstraintKind::PackElementOf:
+    return simplifyPackElementOfConstraint(first, second, subflags, locator);
+
   case ConstraintKind::ValueMember:
   case ConstraintKind::UnresolvedValueMember:
   case ConstraintKind::ValueWitness:
@@ -14355,6 +14399,11 @@ ConstraintSystem::simplifyConstraint(const Constraint &constraint) {
     return simplifyBindTupleOfFunctionParamsConstraint(
         constraint.getFirstType(), constraint.getSecondType(), /*flags*/ None,
         constraint.getLocator());
+
+  case ConstraintKind::PackElementOf:
+    return simplifyPackElementOfConstraint(
+        constraint.getFirstType(), constraint.getSecondType(), /*flags*/None,
+        constraint.getLocator());
   }
 
   llvm_unreachable("Unhandled ConstraintKind in switch.");

Diff for: lib/Sema/Constraint.cpp

@@ -79,6 +79,7 @@ Constraint::Constraint(ConstraintKind Kind, Type First, Type Second,
   case ConstraintKind::UnresolvedMemberChainBase:
   case ConstraintKind::PropertyWrapper:
   case ConstraintKind::BindTupleOfFunctionParams:
+  case ConstraintKind::PackElementOf:
     assert(!First.isNull());
     assert(!Second.isNull());
     break;
@@ -165,6 +166,7 @@ Constraint::Constraint(ConstraintKind Kind, Type First, Type Second, Type Third,
   case ConstraintKind::PropertyWrapper:
   case ConstraintKind::SyntacticElement:
   case ConstraintKind::BindTupleOfFunctionParams:
+  case ConstraintKind::PackElementOf:
     llvm_unreachable("Wrong constructor");
 
   case ConstraintKind::KeyPath:
@@ -310,6 +312,7 @@ Constraint *Constraint::clone(ConstraintSystem &cs) const {
   case ConstraintKind::UnresolvedMemberChainBase:
   case ConstraintKind::PropertyWrapper:
   case ConstraintKind::BindTupleOfFunctionParams:
+  case ConstraintKind::PackElementOf:
     return create(cs, getKind(), getFirstType(), getSecondType(), getLocator());
 
   case ConstraintKind::ApplicableFunction:
@@ -546,6 +549,10 @@ void Constraint::print(llvm::raw_ostream &Out, SourceManager *sm, unsigned inden
     Out << " bind tuple of function params to ";
     break;
 
+  case ConstraintKind::PackElementOf:
+    Out << " element of pack expansion pattern ";
+    break;
+
   case ConstraintKind::Disjunction:
     llvm_unreachable("disjunction handled above");
   case ConstraintKind::Conjunction:
@@ -710,6 +717,7 @@ gatherReferencedTypeVars(Constraint *constraint,
   case ConstraintKind::UnresolvedMemberChainBase:
   case ConstraintKind::PropertyWrapper:
   case ConstraintKind::BindTupleOfFunctionParams:
+  case ConstraintKind::PackElementOf:
     constraint->getFirstType()->getTypeVariables(typeVars);
     constraint->getSecondType()->getTypeVariables(typeVars);
     break;

Diff for: test/Constraints/pack-expansion-expressions.swift

@@ -20,14 +20,9 @@ func forward<U...>(_ u: U...) -> (U...) {
   return tuplify(u...)
 }
 
-// FIXME: This fails with opened element types. Add an ElementOf constraint to
-// handle cases where the pattern type is resolved after constraint generation
-// and mapped to pack archetypes in the solver.
-/*
 func forwardAndMap<U..., V...>(us u: U..., vs v: V...) -> ([(U, V)]...) {
   return tuplify([(u, v)]...)
 }
-*/
 
 func variadicMap<T..., Result...>(_ t: T..., transform: ((T) -> Result)...) -> (Result...) {
   return (transform(t)...)