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

When we’re trying to find the overload set corresponding to a particular
type variable, look through “optional object of” constraints that represent
the use of ? binding or ! forcing. This allows us to find overload sets
when referring to, e.g., @objc optional protocol requirements.

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,32 +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;
-      });
+  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;
+          }
 
-  if (disjunctions.empty())
+          // 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;
+        });
+
+    // If we found a disjunction, return it.
+    if (!disjunctions.empty())
+      return disjunctions[0];
+
+    // If we found an "optional object of" constraint, follow it.
+    if (optionalObjectTypeVar && !getFixedType(optionalObjectTypeVar)) {
+      if (numOptionalUnwraps)
+        ++*numOptionalUnwraps;
+
+      tyvar = optionalObjectTypeVar;
+      rep = getRepresentative(tyvar);
+      continue;
+    }
+
+    // There is nowhere else to look.
     return nullptr;
+  }
 
-  return disjunctions[0];
+  return nullptr;
 }
 
 /// solely resolved by an overload set.

lib/Sema/ConstraintSystem.h

@@ -3166,7 +3166,14 @@ 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);
+  ///
+  /// \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:
   /// Given a type variable that might represent an overload set, retrieve

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)
+}