[CSSimplify] Propagate contextual types to pack expansion variables

If there are explicit generic arguments that fully resolves the
pack expansion, let's bind opened pack expansion to its contextual
type early (while resolving pack expansion variable), doing so
helps with performance and diagnostics.

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -3892,11 +3892,11 @@ class ConstraintSystem {
   /// variable representing a pack expansion type, let's resolve the expansion.
   ///
   /// \param typeVar The type variable representing pack expansion type.
-  /// \param locator The locator associated with contextual type.
+  /// \param contextualType The contextual type this pack expansion variable
+  /// would be bound/equated to.
   ///
   /// \returns `true` if pack expansion has been resolved, `false` otherwise.
-  bool resolvePackExpansion(TypeVariableType *typeVar,
-                            ConstraintLocatorBuilder locator);
+  bool resolvePackExpansion(TypeVariableType *typeVar, Type contextualType);
 
   /// Assign a fixed type to the given type variable.
   ///

lib/Sema/CSSimplify.cpp

@@ -4196,15 +4196,15 @@ static bool isBindable(TypeVariableType *typeVar, Type type) {
 
 ConstraintSystem::TypeMatchResult
 ConstraintSystem::matchTypesBindTypeVar(
-    TypeVariableType *typeVar, Type type, ConstraintKind kind,
+    TypeVariableType *typeVar, Type origType, ConstraintKind kind,
     TypeMatchOptions flags, ConstraintLocatorBuilder locator,
     llvm::function_ref<TypeMatchResult()> formUnsolvedResult) {
   assert(typeVar->is<TypeVariableType>() && "Expected a type variable!");
-  assert(!type->is<TypeVariableType>() && "Expected a non-type variable!");
+  assert(!origType->is<TypeVariableType>() && "Expected a non-type variable!");
 
   // Simplify the right-hand type and perform the "occurs" check.
   typeVar = getRepresentative(typeVar);
-  type = simplifyType(type, flags);
+  auto type = simplifyType(origType, flags);
   if (!isBindable(typeVar, type)) {
     if (shouldAttemptFixes()) {
       // If type variable is allowed to be a hole and it can't be bound to
@@ -4306,7 +4306,7 @@ ConstraintSystem::matchTypesBindTypeVar(
     if (!flags.contains(TMF_BindingTypeVariable))
       return formUnsolvedResult();
 
-    return resolvePackExpansion(typeVar, locator)
+    return resolvePackExpansion(typeVar, origType)
                ? getTypeMatchSuccess()
                : getTypeMatchFailure(locator);
   }
@@ -11354,17 +11354,30 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
 }
 
 bool ConstraintSystem::resolvePackExpansion(TypeVariableType *typeVar,
-                                            ConstraintLocatorBuilder locator) {
+                                            Type contextualType) {
+  auto *locator = typeVar->getImpl().getLocator();
+
   Type openedExpansionType;
-  if (auto last = locator.last()) {
-    auto expansionElt = last->castTo<LocatorPathElt::PackExpansionType>();
-    openedExpansionType = expansionElt.getOpenedType();
+  if (auto expansionElt =
+          locator->getLastElementAs<LocatorPathElt::PackExpansionType>()) {
+    openedExpansionType = expansionElt->getOpenedType();
   }
 
   if (!openedExpansionType)
     return false;
 
-  assignFixedType(typeVar, openedExpansionType, getConstraintLocator(locator));
+  assignFixedType(typeVar, openedExpansionType, locator);
+
+  // We have a fully resolved contextual pack expansion type, let's
+  // apply it right away.
+  if (!contextualType->isEqual(openedExpansionType)) {
+    assert(contextualType->is<PackExpansionType>() &&
+           !contextualType->hasTypeVariable());
+    auto result = matchTypes(openedExpansionType, contextualType,
+                             ConstraintKind::Equal, {}, locator);
+    return !result.isFailure();
+  }
+
   return true;
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -7235,21 +7235,6 @@ TypeVarBindingProducer::TypeVarBindingProducer(BindingSet &bindings)
     return;
   }
 
-  // Pack expansion type variable can only ever have one binding
-  // which is handled by \c resolvePackExpansion.
-  //
-  // There is no need to iterate over other bindings here because
-  // there is no use for contextual types (unlike closures that can
-  // propagate contextual information into the body).
-  if (TypeVar->getImpl().isPackExpansion()) {
-    for (const auto &entry : bindings.Defaults) {
-      auto *constraint = entry.second;
-      Bindings.push_back(getDefaultBinding(constraint));
-    }
-
-    return;
-  }
-
   // A binding to `Any` which should always be considered as a last resort.
   Optional<Binding> Any;
 
@@ -7266,6 +7251,40 @@ TypeVarBindingProducer::TypeVarBindingProducer(BindingSet &bindings)
     }
   };
 
+  if (TypeVar->getImpl().isPackExpansion()) {
+    SmallVector<Binding> viableBindings;
+
+    // Collect possible contextual types (keep in mind that pack
+    // expansion type variable gets bound to its "opened" type
+    // regardless). To be viable the binding has to come from `bind`
+    // or `equal` constraint (i.e. same-type constraint or explicit
+    // generic argument) and be fully resolved.
+    llvm::copy_if(bindings.Bindings, std::back_inserter(viableBindings),
+                  [&](const Binding &binding) {
+                    auto *source = binding.getSource();
+                    if (source->getKind() == ConstraintKind::Bind ||
+                        source->getKind() == ConstraintKind::Equal) {
+                      auto type = binding.BindingType;
+                      return type->is<PackExpansionType>() &&
+                             !type->hasTypeVariable();
+                    }
+                    return false;
+                  });
+
+    // If there is a single fully resolved contextual type, let's
+    // use it as a binding to help with performance and diagnostics.
+    if (viableBindings.size() == 1) {
+      addBinding(viableBindings.front());
+    } else {
+      for (const auto &entry : bindings.Defaults) {
+        auto *constraint = entry.second;
+        Bindings.push_back(getDefaultBinding(constraint));
+      }
+    }
+
+    return;
+  }
+
   for (const auto &binding : bindings.Bindings) {
     addBinding(binding);
   }