[CSBindings] Instead of breaking a PackElementOf constraint down into Bind

constraints, vend potential bindings through PotentialBindings::infer.

This allows for bidirectional binding inference from the pack type to the
element type and vice versa.

Author: hborla

lib/Sema/CSBindings.cpp

@@ -15,6 +15,7 @@
 //===----------------------------------------------------------------------===//
 #include "swift/Sema/CSBindings.h"
 #include "TypeChecker.h"
+#include "swift/AST/GenericEnvironment.h"
 #include "swift/Sema/ConstraintGraph.h"
 #include "swift/Sema/ConstraintSystem.h"
 #include "llvm/ADT/SetVector.h"
@@ -1404,12 +1405,51 @@ void PotentialBindings::infer(Constraint *constraint) {
   case ConstraintKind::SyntacticElement:
   case ConstraintKind::Conjunction:
   case ConstraintKind::BindTupleOfFunctionParams:
-  case ConstraintKind::PackElementOf:
   case ConstraintKind::ShapeOf:
   case ConstraintKind::ExplicitGenericArguments:
     // Constraints from which we can't do anything.
     break;
 
+  case ConstraintKind::PackElementOf: {
+    auto elementType = CS.simplifyType(constraint->getFirstType());
+    auto packType = CS.simplifyType(constraint->getSecondType());
+
+    auto *loc = constraint->getLocator();
+    auto openedElement  =
+        loc->getLastElementAs<LocatorPathElt::OpenedPackElement>();
+
+    if (elementType->isTypeVariableOrMember() && packType->isTypeVariableOrMember())
+      break;
+
+    auto *elementVar = elementType->getAs<TypeVariableType>();
+    auto *packVar = packType->getAs<TypeVariableType>();
+
+    if (elementVar == TypeVar && !packVar) {
+      // Produce a potential binding to the opened element archetype corresponding
+      // to the pack type.
+      packType = packType->mapTypeOutOfContext();
+      if (auto *expansion = packType->getAs<PackExpansionType>())
+        packType = expansion->getPatternType();
+
+      auto *elementEnv = openedElement->getGenericEnvironment();
+      auto elementType = elementEnv->mapPackTypeIntoElementContext(packType);
+      addPotentialBinding({elementType, AllowedBindingKind::Exact, constraint});
+
+      break;
+    } else if (packVar == TypeVar && !elementVar) {
+      // Produce a potential binding to the pack archetype corresponding to
+      // the opened element type.
+      auto *packEnv = CS.DC->getGenericEnvironmentOfContext();
+      elementType = elementType->mapTypeOutOfContext();
+      auto patternType = packEnv->mapElementTypeIntoPackContext(elementType);
+      addPotentialBinding({patternType, AllowedBindingKind::Exact, constraint});
+
+      break;
+    }
+
+    break;
+  }
+
   // For now let's avoid inferring protocol requirements from
   // this constraint, but in the future we could do that to
   // to filter bindings.

lib/Sema/CSGen.cpp

@@ -2956,6 +2956,7 @@ namespace {
     }
 
     Type visitPackExpansionExpr(PackExpansionExpr *expr) {
+      auto *elementEnv = expr->getGenericEnvironment();
       for (auto *binding : expr->getBindings()) {
         auto type = visit(binding);
         CS.setType(binding, type);
@@ -2967,8 +2968,10 @@ namespace {
                                              TVO_CanBindToPack |
                                              TVO_CanBindToHole);
       auto elementResultType = CS.getType(expr->getPatternExpr());
+      auto *elementLoc = CS.getConstraintLocator(
+          expr, LocatorPathElt::OpenedPackElement(elementEnv));
       CS.addConstraint(ConstraintKind::PackElementOf, elementResultType,
-                       patternTy, CS.getConstraintLocator(expr));
+                       patternTy, elementLoc);
 
       auto *shapeLoc =
           CS.getConstraintLocator(expr, ConstraintLocator::PackShape);

lib/Sema/CSSimplify.cpp

@@ -8561,27 +8561,30 @@ ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
                                                   TypeMatchOptions flags,
                                                   ConstraintLocatorBuilder locator) {
   auto elementType = simplifyType(first, flags);
-  auto *loc = getConstraintLocator(locator);
+  auto packType = simplifyType(second, flags);
 
-  if (elementType->hasTypeVariable()) {
+  if (elementType->hasTypeVariable() || packType->hasTypeVariable()) {
     if (!flags.contains(TMF_GenerateConstraints))
       return SolutionKind::Unsolved;
 
+    auto *loc = getConstraintLocator(locator);
     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 *environment = DC->getGenericEnvironmentOfContext();
-  auto patternType = environment->mapElementTypeIntoPackContext(
-      elementType->mapTypeOutOfContext());
-  addConstraint(ConstraintKind::Bind, second, patternType, locator);
-
-  return SolutionKind::Solved;
+  // This constraint only exists to vend bindings.
+  auto *packEnv = DC->getGenericEnvironmentOfContext();
+  if (auto *expansion = packType->getAs<PackExpansionType>())
+    packType = expansion->getPatternType();
+  if (packType->isEqual(packEnv->mapElementTypeIntoPackContext
+                        (elementType->mapTypeOutOfContext()))) {
+    return SolutionKind::Solved;
+  } else {
+    return SolutionKind::Error;
+  }
 }
 
 static bool isForKeyPathSubscript(ConstraintSystem &cs,