[CS] Adjust applied overload simplification

Currently `simplifyAppliedOverloads` depends on
the order in which constraints are simplified,
specifically that a lookup constraint for a
function gets simplified before the applicable
function constraint. This happens to work out
just fine today with the order in which we
re-activate constraints, but I'm planning on
changing that order.

This commit changes the logic such that it it's no
longer affected by the order in which constraints
are simplified. We'll now run it when either an
applicable function constraint is added, or a new
bind overload disjunction is added. This also
means we no longer need to run it potentially
multiple times when simplifying the applicable fn.

Author: hamishknight

lib/Sema/CSSimplify.cpp

@@ -7810,21 +7810,10 @@ ConstraintSystem::simplifyKeyPathApplicationConstraint(
   return unsolved();
 }
 
-Type ConstraintSystem::simplifyAppliedOverloads(
-                                TypeVariableType *fnTypeVar,
-                                const FunctionType *argFnType,
-                                ConstraintLocatorBuilder locator) {
-  Type fnType(fnTypeVar);
-
-  // Always work on the representation.
-  fnTypeVar = getRepresentative(fnTypeVar);
-
-  // Dig out the disjunction that describes this overload.
-  unsigned numOptionalUnwraps = 0;
-  auto disjunction =
-      getUnboundBindOverloadDisjunction(fnTypeVar, &numOptionalUnwraps);
-  if (!disjunction) return fnType;
-
+bool ConstraintSystem::simplifyAppliedOverloadsImpl(
+    Constraint *disjunction, TypeVariableType *fnTypeVar,
+    const FunctionType *argFnType, unsigned numOptionalUnwraps,
+    ConstraintLocatorBuilder locator) {
   if (shouldAttemptFixes()) {
     auto arguments = argFnType->getParams();
     bool allHoles =
@@ -7854,7 +7843,7 @@ Type ConstraintSystem::simplifyAppliedOverloads(
     // regardless of problems related to missing or extraneous labels
     // and/or arguments.
     if (solverState)
-      return fnTypeVar;
+      return false;
   }
 
   /// The common result type amongst all function overloads.
@@ -7941,22 +7930,16 @@ Type ConstraintSystem::simplifyAppliedOverloads(
       argumentInfo.reset();
       goto retry_after_fail;
     }
-
-    return Type();
-
+    return true;
   case SolutionKind::Solved:
-    // We should now have a type for the one remaining overload.
-    fnType = getFixedTypeRecursive(fnType, /*wantRValue=*/true);
-    break;
-
   case SolutionKind::Unsolved:
     break;
   }
 
   // If there was a constraint that we couldn't reason about, don't use the
   // results of any common-type computations.
   if (hasUnhandledConstraints)
-    return fnType;
+    return false;
 
   // If we have a common result type, bind the expected result type to it.
   if (commonResultType && !commonResultType->is<ErrorType>()) {
@@ -7969,16 +7952,67 @@ Type ConstraintSystem::simplifyAppliedOverloads(
         << ")\n";
     }
 
-    // FIXME: Could also rewrite fnType to include this result type.
-    // Introduction of `Bind` constraint here could result in the disconnect
+    // Introduction of a `Bind` constraint here could result in the disconnect
     // in the constraint system with unintended consequences because e.g.
     // in case of key path application it could disconnect one of the
     // components like subscript from the rest of the context.
     addConstraint(ConstraintKind::Equal, argFnType->getResult(),
                   commonResultType, locator);
   }
+  return false;
+}
+
+bool ConstraintSystem::simplifyAppliedOverloads(
+    Constraint *disjunction, ConstraintLocatorBuilder locator) {
+  auto choices = disjunction->getNestedConstraints();
+  assert(choices.size() >= 2);
+  assert(choices.front()->getKind() == ConstraintKind::BindOverload);
+
+  // If we've already bound the overload type var, bail.
+  auto *typeVar = choices.front()->getFirstType()->getAs<TypeVariableType>();
+  if (!typeVar || getFixedType(typeVar))
+    return false;
+
+  // Try to find an applicable fn constraint that applies the overload choice.
+  auto result = findConstraintThroughOptionals(
+      typeVar, OptionalWrappingDirection::Unwrap,
+      [&](Constraint *match, TypeVariableType *currentRep) {
+        // Check to see if we have an applicable fn with a type var RHS that
+        // matches the disjunction.
+        if (match->getKind() != ConstraintKind::ApplicableFunction)
+          return false;
 
-  return fnType;
+        auto *rhsTyVar = match->getSecondType()->getAs<TypeVariableType>();
+        return rhsTyVar && currentRep == getRepresentative(rhsTyVar);
+      });
+
+  if (!result)
+    return false;
+
+  auto *applicableFn = result->first;
+  auto *fnTypeVar = applicableFn->getSecondType()->castTo<TypeVariableType>();
+  auto argFnType = applicableFn->getFirstType()->castTo<FunctionType>();
+  return simplifyAppliedOverloadsImpl(disjunction, fnTypeVar, argFnType,
+                                      /*numOptionalUnwraps*/ result->second,
+                                      locator);
+}
+
+bool ConstraintSystem::simplifyAppliedOverloads(
+    Type fnType, const FunctionType *argFnType,
+    ConstraintLocatorBuilder locator) {
+  // If we've already bound the function type, bail.
+  auto *fnTypeVar = fnType->getAs<TypeVariableType>();
+  if (!fnTypeVar || getFixedType(fnTypeVar))
+    return false;
+
+  // Try to find a corresponding bind overload disjunction.
+  unsigned numOptionalUnwraps = 0;
+  auto *disjunction =
+      getUnboundBindOverloadDisjunction(fnTypeVar, &numOptionalUnwraps);
+  if (!disjunction)
+    return false;
+  return simplifyAppliedOverloadsImpl(disjunction, fnTypeVar, argFnType,
+                                      numOptionalUnwraps, locator);
 }
 
 ConstraintSystem::SolutionKind
@@ -8055,16 +8089,6 @@ ConstraintSystem::simplifyApplicableFnConstraint(
 
   };
 
-  // If the right-hand side is a type variable,
-  // try to simplify the overload set.
-  if (auto typeVar = desugar2->getAs<TypeVariableType>()) {
-    Type newType2 = simplifyAppliedOverloads(typeVar, func1, locator);
-    if (!newType2)
-      return SolutionKind::Error;
-
-    desugar2 = newType2->getDesugaredType();
-  }
-
   // If right-hand side is a type variable, the constraint is unsolved.
   if (desugar2->isTypeVariableOrMember()) {
     return formUnsolved();
@@ -9408,9 +9432,14 @@ ConstraintSystem::addConstraintImpl(ConstraintKind kind, Type first,
   case ConstraintKind::BridgingConversion:
     return simplifyBridgingConstraint(first, second, subflags, locator);
 
-  case ConstraintKind::ApplicableFunction:
+  case ConstraintKind::ApplicableFunction: {
+    // First try to simplify the overload set for the function being applied.
+    if (simplifyAppliedOverloads(second, first->castTo<FunctionType>(),
+                                 locator)) {
+      return SolutionKind::Error;
+    }
     return simplifyApplicableFnConstraint(first, second, subflags, locator);
-
+  }
   case ConstraintKind::DynamicCallableApplicableFunction:
     return simplifyDynamicCallableApplicableFnConstraint(first, second,
                                                          subflags, locator);

lib/Sema/CSSolver.cpp

@@ -1790,67 +1790,92 @@ bool ConstraintSystem::haveTypeInformationForAllArguments(
                       });
 }
 
-Constraint *ConstraintSystem::getUnboundBindOverloadDisjunction(
-    TypeVariableType *tyvar, unsigned *numOptionalUnwraps) {
-  if (numOptionalUnwraps)
-    *numOptionalUnwraps = 0;
-
-  auto *rep = getRepresentative(tyvar);
-  assert(!getFixedType(rep));
+Optional<std::pair<Constraint *, unsigned>>
+ConstraintSystem::findConstraintThroughOptionals(
+    TypeVariableType *typeVar, OptionalWrappingDirection optionalDirection,
+    llvm::function_ref<bool(Constraint *, TypeVariableType *)> predicate) {
+  unsigned numOptionals = 0;
+  auto *rep = getRepresentative(typeVar);
 
   SmallPtrSet<TypeVariableType *, 4> visitedVars;
   while (visitedVars.insert(rep).second) {
     // Look for a disjunction that binds this type variable to an overload set.
     TypeVariableType *optionalObjectTypeVar = nullptr;
-    auto disjunctions = getConstraintGraph().gatherConstraints(
+    auto constraints = getConstraintGraph().gatherConstraints(
         rep, ConstraintGraph::GatheringKind::EquivalenceClass,
-        [this, rep, &optionalObjectTypeVar](Constraint *match) {
-          // If we have an "optional object of" constraint where the right-hand
-          // side is this type variable, we may need to follow that type
-          // variable to find the disjunction.
-          if (match->getKind() == ConstraintKind::OptionalObject) {
+        [&](Constraint *match) {
+          // If we have an "optional object of" constraint, we may need to
+          // look through it to find the constraint we're looking for.
+          if (match->getKind() != ConstraintKind::OptionalObject)
+            return predicate(match, rep);
+
+          switch (optionalDirection) {
+          case OptionalWrappingDirection::Promote: {
+            // We want to go from T to T?, so check if we're on the RHS, and
+            // move over to the LHS if we can.
             auto rhsTypeVar = match->getSecondType()->getAs<TypeVariableType>();
             if (rhsTypeVar && getRepresentative(rhsTypeVar) == rep) {
               optionalObjectTypeVar =
                   match->getFirstType()->getAs<TypeVariableType>();
             }
-            return false;
+            break;
           }
-
-          // We only care about disjunctions of overload bindings.
-          if (match->getKind() != ConstraintKind::Disjunction ||
-              match->getNestedConstraints().front()->getKind() !=
-                     ConstraintKind::BindOverload)
-            return false;
-
-          auto lhsTypeVar =
-              match->getNestedConstraints().front()->getFirstType()
-                ->getAs<TypeVariableType>();
-          if (!lhsTypeVar)
-            return false;
-
-          return getRepresentative(lhsTypeVar) == rep;
+          case OptionalWrappingDirection::Unwrap: {
+            // We want to go from T? to T, so check if we're on the LHS, and
+            // move over to the RHS if we can.
+            auto lhsTypeVar = match->getFirstType()->getAs<TypeVariableType>();
+            if (lhsTypeVar && getRepresentative(lhsTypeVar) == rep) {
+              optionalObjectTypeVar =
+                  match->getSecondType()->getAs<TypeVariableType>();
+            }
+            break;
+          }
+          }
+          // Don't include the optional constraint in the results.
+          return false;
         });
 
-    // If we found a disjunction, return it.
-    if (!disjunctions.empty())
-      return disjunctions[0];
+    // If we found a result, return it.
+    if (!constraints.empty())
+      return std::make_pair(constraints[0], numOptionals);
 
     // If we found an "optional object of" constraint, follow it.
     if (optionalObjectTypeVar && !getFixedType(optionalObjectTypeVar)) {
-      if (numOptionalUnwraps)
-        ++*numOptionalUnwraps;
-
-      tyvar = optionalObjectTypeVar;
-      rep = getRepresentative(tyvar);
+      numOptionals += 1;
+      rep = getRepresentative(optionalObjectTypeVar);
       continue;
     }
 
-    // There is nowhere else to look.
-    return nullptr;
+    // Otherwise we're done.
+    return None;
   }
+  return None;
+}
 
-  return nullptr;
+Constraint *ConstraintSystem::getUnboundBindOverloadDisjunction(
+    TypeVariableType *tyvar, unsigned *numOptionalUnwraps) {
+  assert(!getFixedType(tyvar));
+  auto result = findConstraintThroughOptionals(
+      tyvar, OptionalWrappingDirection::Promote,
+      [&](Constraint *match, TypeVariableType *currentRep) {
+        // Check to see if we have a bind overload disjunction that binds the
+        // type var we need.
+        if (match->getKind() != ConstraintKind::Disjunction ||
+            match->getNestedConstraints().front()->getKind() !=
+                ConstraintKind::BindOverload)
+          return false;
+
+        auto lhsTy = match->getNestedConstraints().front()->getFirstType();
+        auto *lhsTyVar = lhsTy->getAs<TypeVariableType>();
+        return lhsTyVar && currentRep == getRepresentative(lhsTyVar);
+      });
+  if (!result)
+    return nullptr;
+
+  if (numOptionalUnwraps)
+    *numOptionalUnwraps = result->second;
+
+  return result->first;
 }
 
 // Find a disjunction associated with an ApplicableFunction constraint

lib/Sema/ConstraintSystem.cpp

@@ -1694,7 +1694,14 @@ void ConstraintSystem::addOverloadSet(ArrayRef<Constraint *> choices,
     return;
   }
 
-  addDisjunctionConstraint(choices, locator, ForgetChoice);
+  auto *disjunction =
+      Constraint::createDisjunction(*this, choices, locator, ForgetChoice);
+  addUnsolvedConstraint(disjunction);
+  if (simplifyAppliedOverloads(disjunction, locator)) {
+    retireConstraint(disjunction);
+    if (!failedConstraint)
+      failedConstraint = disjunction;
+  }
 }
 
 /// If we're resolving an overload set with a decl that has special type

lib/Sema/ConstraintSystem.h

@@ -3357,22 +3357,77 @@ class ConstraintSystem {
       llvm::function_ref<void(unsigned int, Type, ConstraintLocator *)>
           verifyThatArgumentIsHashable);
 
-public:
+  /// Describes a direction of optional wrapping, either increasing optionality
+  /// or decreasing optionality.
+  enum class OptionalWrappingDirection {
+    /// Unwrap an optional type T? to T.
+    Unwrap,
+
+    /// Promote a type T to optional type T?.
+    Promote
+  };
+
+  /// Attempts to find a constraint that involves \p typeVar and satisfies
+  /// \p predicate, looking through optional object constraints if necessary. If
+  /// multiple candidates are found, returns the first one.
+  ///
+  /// \param optionalDirection The direction to travel through optional object
+  /// constraints, either increasing or decreasing optionality.
+  ///
+  /// \param predicate Checks whether a given constraint is the one being
+  /// searched for. The type variable passed is the current representative
+  /// after looking through the optional object constraints.
+  ///
+  /// \returns The constraint found along with the number of optional object
+  /// constraints looked through, or \c None if no constraint was found.
+  Optional<std::pair<Constraint *, unsigned>> findConstraintThroughOptionals(
+      TypeVariableType *typeVar, OptionalWrappingDirection optionalDirection,
+      llvm::function_ref<bool(Constraint *, TypeVariableType *)> predicate);
+
   /// Attempt to simplify the set of overloads corresponding to a given
   /// function application constraint.
   ///
+  /// \param disjunction The disjunction for the set of overloads.
+  ///
   /// \param fnTypeVar The type variable that describes the set of
   /// overloads for the function.
   ///
   /// \param argFnType The call signature, which includes the call arguments
   /// (as the function parameters) and the expected result type of the
   /// call.
   ///
-  /// \returns \c fnType, or some simplified form of it if this function
-  /// was able to find a single overload or derive some common structure
-  /// among the overloads.
-  Type simplifyAppliedOverloads(TypeVariableType *fnTypeVar,
-                                const FunctionType *argFnType,
+  /// \param numOptionalUnwraps The number of unwraps required to get the
+  /// underlying function from the overload choice.
+  ///
+  /// \returns \c true if an error was encountered, \c false otherwise.
+  bool simplifyAppliedOverloadsImpl(Constraint *disjunction,
+                                    TypeVariableType *fnTypeVar,
+                                    const FunctionType *argFnType,
+                                    unsigned numOptionalUnwraps,
+                                    ConstraintLocatorBuilder locator);
+
+public:
+  /// Attempt to simplify the set of overloads corresponding to a given
+  /// bind overload disjunction.
+  ///
+  /// \param disjunction The disjunction for the set of overloads.
+  ///
+  /// \returns \c true if an error was encountered, \c false otherwise.
+  bool simplifyAppliedOverloads(Constraint *disjunction,
+                                ConstraintLocatorBuilder locator);
+
+  /// Attempt to simplify the set of overloads corresponding to a given
+  /// function application constraint.
+  ///
+  /// \param fnType The type that describes the set of overloads for the
+  /// function.
+  ///
+  /// \param argFnType The call signature, which includes the call arguments
+  /// (as the function parameters) and the expected result type of the
+  /// call.
+  ///
+  /// \returns \c true if an error was encountered, \c false otherwise.
+  bool simplifyAppliedOverloads(Type fnType, const FunctionType *argFnType,
                                 ConstraintLocatorBuilder locator);
 
   /// Retrieve the type that will be used when matching the given overload.