[CSSimplify] Merge binding inference logic with PackElementOf constraint simplification

It's only safe to infer element type from `PackElementOf` constraint
when pattern type is fully resolved (because it can have pack element
archetypes which should be mapped out of context), the same applies
to the pattern type inference as well. Since constraints are re-activated
every time a referenced type variable is bound, simplication logic
can act as inference source by decaying into `Equal` constraints
where pattern/element type are resolved via surrounding information
first.

Author: xedin

lib/Sema/CSBindings.cpp

@@ -1479,56 +1479,10 @@ void PotentialBindings::infer(Constraint *constraint) {
   case ConstraintKind::BindTupleOfFunctionParams:
   case ConstraintKind::ShapeOf:
   case ConstraintKind::ExplicitGenericArguments:
+  case ConstraintKind::PackElementOf:
     // Constraints from which we can't do anything.
     break;
 
-  case ConstraintKind::PackElementOf: {
-    auto elementType = CS.simplifyType(constraint->getFirstType());
-    auto packType = CS.simplifyType(constraint->getSecondType());
-
-    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.
-      auto shapeClass = packType->getReducedShape();
-      packType = packType->mapTypeOutOfContext();
-      auto *elementEnv = CS.getPackElementEnvironment(constraint->getLocator(),
-                                                      shapeClass);
-
-      // Without an opened element environment, we cannot derive the
-      // element binding.
-      if (!elementEnv)
-        break;
-
-      auto elementType = elementEnv->mapPackTypeIntoElementContext(packType);
-      assert(!elementType->is<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.
-      Type patternType;
-      auto *packEnv = CS.DC->getGenericEnvironmentOfContext();
-      if (!elementType->hasElementArchetype()) {
-        patternType = elementType;
-      } else {
-        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/CSSimplify.cpp

@@ -8998,67 +8998,78 @@ ConstraintSystem::simplifyBindTupleOfFunctionParamsConstraint(
   return SolutionKind::Solved;
 }
 
-static Type lookThroughSingletonPackExpansion(Type ty) {
-  if (auto pack = ty->getAs<PackType>()) {
-    if (pack->getNumElements() == 1) {
-      if (auto expansion = pack->getElementType(0)->getAs<PackExpansionType>()) {
-        auto countType = expansion->getCountType();
-        if (countType->isEqual(expansion->getPatternType()))
-          return countType;
-      }
-    }
-  }
-  return ty;
-}
-
 ConstraintSystem::SolutionKind
 ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
                                                   TypeMatchOptions flags,
                                                   ConstraintLocatorBuilder locator) {
   auto elementType = simplifyType(first, flags);
-  auto packType = simplifyType(second, flags);
+  auto patternType = simplifyType(second, flags);
 
-  if (elementType->hasTypeVariable() || packType->hasTypeVariable()) {
+  auto formUnsolved = [&]() {
     if (!flags.contains(TMF_GenerateConstraints))
       return SolutionKind::Unsolved;
 
-    auto *loc = getConstraintLocator(locator);
     addUnsolvedConstraint(
-        Constraint::create(*this, ConstraintKind::PackElementOf,
-                           first, second, loc));
+        Constraint::create(*this, ConstraintKind::PackElementOf, first, second,
+                           getConstraintLocator(locator)));
 
     return SolutionKind::Solved;
+  };
+
+  // If neither side is fully resolved yet, there is nothing we can do.
+  if (elementType->hasTypeVariable() && patternType->hasTypeVariable())
+    return formUnsolved();
+
+  if (shouldAttemptFixes()) {
+    if (elementType->isPlaceholder() || patternType->isPlaceholder())
+      return SolutionKind::Solved;
   }
 
-  // FIXME: I'm not sure this is actually necessary; I may only be seeing
-  // this because of something I've screwed up in element generic
-  // environments.
-  elementType = lookThroughSingletonPackExpansion(elementType);
+  // Let's try to resolve element type based on the pattern type.
+  if (!patternType->hasTypeVariable()) {
+    auto *loc = getConstraintLocator(locator);
+    auto shapeClass = patternType->getReducedShape();
+    patternType = patternType->mapTypeOutOfContext();
+    auto *elementEnv = getPackElementEnvironment(loc, shapeClass);
 
-  // This constraint only exists to vend bindings.
-  auto *packEnv = DC->getGenericEnvironmentOfContext();
+    // Without an opened element environment, we cannot derive the
+    // element binding.
+    if (!elementEnv) {
+      if (!shouldAttemptFixes())
+        return SolutionKind::Error;
 
-  // Map element archetypes to the pack context to check for equality.
-  if (elementType->hasElementArchetype()) {
-    auto mappedPack = packEnv->mapElementTypeIntoPackContext(elementType);
-    return (packType->isEqual(mappedPack) ?
-            SolutionKind::Solved : SolutionKind::Error);
-  }
+      // `each` was applied to a concrete type.
+      if (!shapeClass->is<PackArchetypeType>()) {
+        if (recordFix(AllowInvalidPackElement::create(*this, patternType, loc)))
+          return SolutionKind::Error;
+      }
 
-  // Pack expansions can have concrete pattern types. In this case, the pack
-  // type and element type will be equal.
-  if (packType->isEqual(elementType)) {
-    return SolutionKind::Solved;
-  }
+      // Only other posibility is that there is a shape mismatch between
+      // elements of the pack expansion pattern which is detected separately.
 
-  if (shouldAttemptFixes()) {
-    auto *loc = getConstraintLocator(locator);
-    if (elementType->isPlaceholder() ||
-        !recordFix(AllowInvalidPackElement::create(*this, packType, loc)))
+      recordAnyTypeVarAsPotentialHole(elementType);
       return SolutionKind::Solved;
+    }
+
+    auto expectedElementTy =
+        elementEnv->mapPackTypeIntoElementContext(patternType);
+    assert(!expectedElementTy->is<PackType>());
+
+    addConstraint(ConstraintKind::Equal, elementType, expectedElementTy,
+                  locator);
+    return SolutionKind::Solved;
   }
 
-  return SolutionKind::Error;
+  // Otherwise we are inferred or checking pattern type.
+
+  auto *packEnv = DC->getGenericEnvironmentOfContext();
+
+  // Map element archetypes to the pack context to check for equality.
+  if (elementType->hasElementArchetype())
+    elementType = packEnv->mapElementTypeIntoPackContext(elementType);
+
+  addConstraint(ConstraintKind::Equal, elementType, patternType, locator);
+  return SolutionKind::Solved;
 }
 
 static bool isForKeyPathSubscript(ConstraintSystem &cs,