Merge pull request #23121 from DougGregor/solver-look-through-optional-binding-overloads

 [Constraint solver] Look through optional binding for overload sets.

Author: DougGregor

lib/Sema/CSSimplify.cpp

@@ -4952,7 +4952,9 @@ Type ConstraintSystem::simplifyAppliedOverloads(
   fnTypeVar = getRepresentative(fnTypeVar);
 
   // Dig out the disjunction that describes this overload.
-  auto disjunction = getUnboundBindOverloadDisjunction(fnTypeVar);
+  unsigned numOptionalUnwraps = 0;
+  auto disjunction =
+      getUnboundBindOverloadDisjunction(fnTypeVar, &numOptionalUnwraps);
   if (!disjunction) return fnType;
 
   /// The common result type amongst all function overloads.
@@ -5013,6 +5015,13 @@ Type ConstraintSystem::simplifyAppliedOverloads(
           return true;
         }
 
+        // Account for any optional unwrapping/binding
+        for (unsigned i : range(numOptionalUnwraps)) {
+          (void)i;
+          if (Type objectType = choiceType->getOptionalObjectType())
+            choiceType = objectType;
+        }
+
         // If we have a function type, we can compute a common result type.
         updateCommonResultType(choiceType);
         return true;

lib/Sema/CSSolver.cpp

@@ -1650,64 +1650,67 @@ bool ConstraintSystem::haveTypeInformationForAllArguments(
 }
 
 Constraint *ConstraintSystem::getUnboundBindOverloadDisjunction(
-                                                  TypeVariableType *tyvar) {
+    TypeVariableType *tyvar, unsigned *numOptionalUnwraps) {
+  if (numOptionalUnwraps)
+    *numOptionalUnwraps = 0;
+
   auto *rep = getRepresentative(tyvar);
   assert(!getFixedType(rep));
 
-  llvm::SetVector<Constraint *> disjunctions;
-  getConstraintGraph().gatherConstraints(
-      rep, disjunctions, ConstraintGraph::GatheringKind::EquivalenceClass,
-      [this, rep](Constraint *match) {
-        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;
-      });
-
-  if (disjunctions.empty())
-    return nullptr;
+  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;
+    llvm::SetVector<Constraint *> disjunctions;
+    getConstraintGraph().gatherConstraints(
+        rep, disjunctions, 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) {
+            auto rhsTypeVar = match->getSecondType()->getAs<TypeVariableType>();
+            if (rhsTypeVar && getRepresentative(rhsTypeVar) == rep) {
+              optionalObjectTypeVar =
+                  match->getFirstType()->getAs<TypeVariableType>();
+            }
+            return false;
+          }
 
-  return disjunctions[0];
-}
+          // We only care about disjunctions of overload bindings.
+          if (match->getKind() != ConstraintKind::Disjunction ||
+              match->getNestedConstraints().front()->getKind() !=
+                     ConstraintKind::BindOverload)
+            return false;
 
-/// solely resolved by an overload set.
-SmallVector<OverloadChoice, 2> ConstraintSystem::getUnboundBindOverloads(
-                                                     TypeVariableType *tyvar) {
-  // Always work on the representation.
-  tyvar = getRepresentative(tyvar);
+          auto lhsTypeVar =
+              match->getNestedConstraints().front()->getFirstType()
+                ->getAs<TypeVariableType>();
+          if (!lhsTypeVar)
+            return false;
 
-  SmallVector<OverloadChoice, 2> choices;
+          return getRepresentative(lhsTypeVar) == rep;
+        });
 
-  auto disjunction = getUnboundBindOverloadDisjunction(tyvar);
-  if (!disjunction) return choices;
+    // If we found a disjunction, return it.
+    if (!disjunctions.empty())
+      return disjunctions[0];
 
-  for (auto constraint : disjunction->getNestedConstraints()) {
-    // We must have bind-overload constraints.
-    if (constraint->getKind() != ConstraintKind::BindOverload) {
-      choices.clear();
-      return choices;
-    }
+    // If we found an "optional object of" constraint, follow it.
+    if (optionalObjectTypeVar && !getFixedType(optionalObjectTypeVar)) {
+      if (numOptionalUnwraps)
+        ++*numOptionalUnwraps;
 
-    // We must be binding the type variable (or a type variable equivalent to
-    // it).
-    auto boundTypeVar = constraint->getFirstType()->getAs<TypeVariableType>();
-    if (!boundTypeVar || getRepresentative(boundTypeVar) != tyvar) {
-      choices.clear();
-      return choices;
+      tyvar = optionalObjectTypeVar;
+      rep = getRepresentative(tyvar);
+      continue;
     }
 
-    choices.push_back(constraint->getOverloadChoice());
+    // There is nowhere else to look.
+    return nullptr;
   }
 
-  return choices;
+  return nullptr;
 }
 
 // Find a disjunction associated with an ApplicableFunction constraint
@@ -1941,7 +1944,7 @@ static Constraint *tryOptimizeGenericDisjunction(
       return type->isAny();
     });
 
-    // If function declaration references `Any` or `Any?` type
+    // If function declaration references `Any` or an optional type,
     // let's not attempt it, because it's unclear
     // without solving which overload is going to be better.
     return !hasAnyOrOptional;

lib/Sema/ConstraintSystem.h

@@ -3166,17 +3166,16 @@ class ConstraintSystem {
   // bind overloads associated with it. This may return null in cases where
   // the disjunction has either not been created or binds the type variable
   // in some manner other than by binding overloads.
-  Constraint *getUnboundBindOverloadDisjunction(TypeVariableType *tyvar);
-
-private:
-  /// Given a type variable that might represent an overload set, retrieve
   ///
-  /// \returns the set of overload choices to which this type variable
-  /// could be bound, or an empty vector if the type variable is not
-  /// solely resolved by an overload set.
-  SmallVector<OverloadChoice, 2> getUnboundBindOverloads(
-                                                  TypeVariableType *tyvar);
+  /// \param numOptionalUnwraps If non-null, this will receive the number
+  /// of "optional object of" constraints that this function looked through
+  /// to uncover the disjunction. The actual overloads will have this number
+  /// of optionals wrapping the type.
+  Constraint *getUnboundBindOverloadDisjunction(
+    TypeVariableType *tyvar,
+    unsigned *numOptionalUnwraps = nullptr);
 
+private:
   /// Solve the system of constraints after it has already been
   /// simplified.
   ///

test/Constraints/common_type_objc.swift

@@ -0,0 +1,22 @@
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -typecheck -verify %s -debug-constraints 2>%t.err
+// RUN: %FileCheck %s < %t.err
+
+// REQUIRES: objc_interop
+
+import Foundation
+
+@objc protocol P {
+  func foo(_ i: Int) -> Double
+  func foo(_ d: Double) -> Double
+
+  @objc optional func opt(_ i: Int) -> Int
+  @objc optional func opt(_ d: Double) -> Int
+}
+
+func testOptional(obj: P) {
+  // CHECK: common result type for {{.*}} is Int
+  _ = obj.opt?(1)
+
+  // CHECK: common result type for {{.*}} is Int
+  _ = obj.opt!(1)
+}