[CS] Consolidate logic forming locators to callees

lib/Sema/CSApply.cpp

@@ -4479,8 +4479,7 @@ namespace {
     buildKeyPathPropertyComponent(const SelectedOverload &overload,
                                   SourceLoc componentLoc,
                                   ConstraintLocator *locator) {
-      if (overload.choice.isDecl()) {
-        auto property = overload.choice.getDecl();
+      if (auto *property = overload.choice.getDeclOrNull()) {
         // Key paths can only refer to properties currently.
         auto varDecl = cast<VarDecl>(property);
         // Key paths don't work with mutating-get properties.
@@ -4784,10 +4783,8 @@ static ConcreteDeclRef resolveLocatorToDecl(
     // Overloaded and unresolved cases: find the resolved overload.
     auto anchorLocator = cs.getConstraintLocator(anchor);
     if (auto selected = findOvlChoice(anchorLocator)) {
-      if (selected->choice.isDecl())
-        return getConcreteDeclRef(selected->choice.getDecl(),
-                                  selected->openedType,
-                                  anchorLocator);
+      if (auto *decl = selected->choice.getDeclOrNull())
+        return getConcreteDeclRef(decl, selected->openedType, anchorLocator);
     }
   }
 
@@ -4796,10 +4793,8 @@ static ConcreteDeclRef resolveLocatorToDecl(
     auto anchorLocator = cs.getConstraintLocator(anchor,
                            ConstraintLocator::UnresolvedMember);
     if (auto selected = findOvlChoice(anchorLocator)) {
-      if (selected->choice.isDecl())
-        return getConcreteDeclRef(selected->choice.getDecl(),
-                                  selected->openedType,
-                                  anchorLocator);
+      if (auto *decl = selected->choice.getDeclOrNull())
+        return getConcreteDeclRef(decl, selected->openedType, anchorLocator);
     }
   }
 
@@ -4808,10 +4803,8 @@ static ConcreteDeclRef resolveLocatorToDecl(
     auto anchorLocator = cs.getConstraintLocator(anchor,
                                                  ConstraintLocator::Member);
     if (auto selected = findOvlChoice(anchorLocator)) {
-      if (selected->choice.isDecl())
-        return getConcreteDeclRef(selected->choice.getDecl(),
-                                  selected->openedType,
-                                  anchorLocator);
+      if (auto *decl = selected->choice.getDeclOrNull())
+        return getConcreteDeclRef(decl, selected->openedType, anchorLocator);
     }
   }
 
@@ -4824,10 +4817,8 @@ static ConcreteDeclRef resolveLocatorToDecl(
     auto anchorLocator =
       cs.getConstraintLocator(anchor, ConstraintLocator::SubscriptMember);
     if (auto selected = findOvlChoice(anchorLocator)) {
-      if (selected->choice.isDecl())
-        return getConcreteDeclRef(selected->choice.getDecl(),
-                                  selected->openedType,
-                                  anchorLocator);
+      if (auto *decl = selected->choice.getDeclOrNull())
+        return getConcreteDeclRef(decl, selected->openedType, anchorLocator);
     }
   }
 

lib/Sema/CSDiag.cpp

@@ -867,9 +867,8 @@ bool FailureDiagnosis::diagnoseGeneralOverloadFailure(Constraint *constraint) {
   if (constraint->getKind() == ConstraintKind::Disjunction) {
     for (auto elt : constraint->getNestedConstraints()) {
       if (elt->getKind() != ConstraintKind::BindOverload) continue;
-      if (!elt->getOverloadChoice().isDecl()) continue;
-      auto candidate = elt->getOverloadChoice().getDecl();
-      diagnose(candidate, diag::found_candidate);
+      if (auto *candidate = elt->getOverloadChoice().getDeclOrNull())
+        diagnose(candidate, diag::found_candidate);
     }
   }
 
@@ -4663,7 +4662,11 @@ static bool isViableOverloadSet(const CalleeCandidateInfo &CCI,
   for (unsigned i = 0; i < CCI.size(); ++i) {
     auto &&cand = CCI[i];
     auto funcDecl = dyn_cast_or_null<AbstractFunctionDecl>(cand.getDecl());
-    if (!funcDecl)
+
+    // If we don't have a func decl or we haven't resolved its parameters,
+    // continue. The latter case can occur with `type(of:)`, which is introduced
+    // as a type variable.
+    if (!funcDecl || !cand.hasParameters())
       continue;
 
     auto params = cand.getParameters();

lib/Sema/CSDiagnostics.cpp

@@ -101,33 +101,9 @@ Expr *FailureDiagnostic::getArgumentExprFor(Expr *anchor) const {
   return nullptr;
 }
 
-// TODO: Replace duplications of this logic with calls to this.
-Optional<SelectedOverload> FailureDiagnostic::getChoiceFor(Expr *expr) {
+Optional<SelectedOverload> FailureDiagnostic::getChoiceFor(Expr *expr) const {
   auto &cs = getConstraintSystem();
-  ConstraintLocator *locator = nullptr;
-
-  if (auto *AE = dyn_cast<ApplyExpr>(expr)) {
-    if (auto *TE = dyn_cast<TypeExpr>(AE->getFn())) {
-      locator = cs.getConstraintLocator(AE,
-                                        {ConstraintLocator::ApplyFunction,
-                                         ConstraintLocator::ConstructorMember},
-                                        /*summaryFlags=*/0);
-    }
-    return getChoiceFor(AE->getFn());
-  } else if (auto *UDE = dyn_cast<UnresolvedDotExpr>(expr)) {
-    locator = cs.getConstraintLocator(UDE, ConstraintLocator::Member);
-  } else if (auto *UME = dyn_cast<UnresolvedMemberExpr>(expr)) {
-    locator = cs.getConstraintLocator(UME, ConstraintLocator::UnresolvedMember);
-  } else if (auto *SE = dyn_cast<SubscriptExpr>(expr)) {
-    locator = cs.getConstraintLocator(SE, ConstraintLocator::SubscriptMember);
-  } else {
-    locator = cs.getConstraintLocator(expr);
-  }
-
-  if (!locator)
-    return None;
-
-  return getOverloadChoiceIfAvailable(locator);
+  return getOverloadChoiceIfAvailable(cs.getCalleeLocator(expr));
 }
 
 Type RequirementFailure::getOwnerType() const {
@@ -179,10 +155,6 @@ ProtocolConformance *RequirementFailure::getConformanceForConditionalReq(
 
 ValueDecl *RequirementFailure::getDeclRef() const {
   auto &cs = getConstraintSystem();
-  auto &TC = getTypeChecker();
-
-  auto *anchor = getRawAnchor();
-  auto *locator = cs.getConstraintLocator(anchor);
 
   // Get a declaration associated with given type (if any).
   // This is used to retrieve affected declaration when
@@ -204,38 +176,7 @@ ValueDecl *RequirementFailure::getDeclRef() const {
   if (isFromContextualType())
     return getAffectedDeclFromType(cs.getContextualType());
 
-  if (auto *AE = dyn_cast<CallExpr>(anchor)) {
-    // NOTE: In valid code, the function can only be a TypeExpr
-    assert(isa<TypeExpr>(AE->getFn()) ||
-           isa<OverloadedDeclRefExpr>(AE->getFn()));
-    ConstraintLocatorBuilder ctor(locator);
-    locator = cs.getConstraintLocator(
-        ctor.withPathElement(PathEltKind::ApplyFunction)
-            .withPathElement(PathEltKind::ConstructorMember));
-  } else if (auto *UDE = dyn_cast<UnresolvedDotExpr>(anchor)) {
-    ConstraintLocatorBuilder member(locator);
-
-    if (TC.getSelfForInitDelegationInConstructor(getDC(), UDE)) {
-      member = member.withPathElement(PathEltKind::ConstructorMember);
-    } else {
-      member = member.withPathElement(PathEltKind::Member);
-    }
-
-    locator = cs.getConstraintLocator(member);
-  } else if (auto *UME = dyn_cast<UnresolvedMemberExpr>(anchor)) {
-    locator = cs.getConstraintLocator(locator, PathEltKind::UnresolvedMember);
-  } else if (isa<SubscriptExpr>(anchor)) {
-    ConstraintLocatorBuilder subscript(locator);
-    locator = cs.getConstraintLocator(
-        subscript.withPathElement(PathEltKind::SubscriptMember));
-  } else if (isa<MemberRefExpr>(anchor)) {
-    ConstraintLocatorBuilder memberRef(locator);
-    locator =
-        cs.getConstraintLocator(memberRef.withPathElement(PathEltKind::Member));
-  }
-
-  auto overload = getOverloadChoiceIfAvailable(locator);
-  if (overload)
+  if (auto overload = getChoiceFor(getRawAnchor()))
     return overload->choice.getDecl();
 
   return getAffectedDeclFromType(getOwnerType());
@@ -568,32 +509,7 @@ bool NoEscapeFuncToTypeConversionFailure::diagnoseParameterUse() const {
 
 Type NoEscapeFuncToTypeConversionFailure::getParameterTypeFor(
     Expr *expr, unsigned paramIdx) const {
-  auto &cs = getConstraintSystem();
-  ConstraintLocator *locator = nullptr;
-
-  if (auto *call = dyn_cast<CallExpr>(expr)) {
-    auto *fnExpr = call->getFn();
-    if (isa<UnresolvedDotExpr>(fnExpr)) {
-      locator = cs.getConstraintLocator(fnExpr, ConstraintLocator::Member);
-    } else if (isa<UnresolvedMemberExpr>(fnExpr)) {
-      locator =
-          cs.getConstraintLocator(fnExpr, ConstraintLocator::UnresolvedMember);
-    } else if (auto *TE = dyn_cast<TypeExpr>(fnExpr)) {
-      locator = cs.getConstraintLocator(call,
-                                        {ConstraintLocator::ApplyFunction,
-                                         ConstraintLocator::ConstructorMember},
-                                        /*summaryFlags=*/0);
-    } else {
-      locator = cs.getConstraintLocator(fnExpr);
-    }
-  } else if (auto *SE = dyn_cast<SubscriptExpr>(expr)) {
-    locator = cs.getConstraintLocator(SE, ConstraintLocator::SubscriptMember);
-  }
-
-  if (!locator)
-    return Type();
-
-  auto choice = getOverloadChoiceIfAvailable(locator);
+  auto choice = getChoiceFor(expr);
   if (!choice)
     return Type();
 

lib/Sema/CSDiagnostics.h

@@ -155,7 +155,7 @@ class FailureDiagnostic {
   /// reference or subscript, nullptr otherwise.
   Expr *getArgumentExprFor(Expr *anchor) const;
 
-  Optional<SelectedOverload> getChoiceFor(Expr *);
+  Optional<SelectedOverload> getChoiceFor(Expr *) const;
 
 private:
   /// Compute anchor expression associated with current diagnostic.

lib/Sema/CSGen.cpp

@@ -622,9 +622,9 @@ namespace {
     unsigned numFavoredConstraints = 0;
     Constraint *firstFavored = nullptr;
     for (auto constraint : disjunction->getNestedConstraints()) {
-      if (!constraint->getOverloadChoice().isDecl())
+      auto *decl = constraint->getOverloadChoice().getDeclOrNull();
+      if (!decl)
         continue;
-      auto decl = constraint->getOverloadChoice().getDecl();
 
       if (mustConsider && mustConsider(decl)) {
         // Roll back any constraints we favored.

lib/Sema/CSSimplify.cpp

@@ -859,8 +859,7 @@ getCalleeDeclAndArgs(ConstraintSystem &cs,
                            hasTrailingClosure);
 
   // If there's a declaration, return it.
-  if (choice->isDecl()) {
-    auto decl = choice->getDecl();
+  if (auto *decl = choice->getDeclOrNull()) {
     bool hasCurriedSelf = false;
     if (decl->getDeclContext()->isTypeContext()) {
       if (auto function = dyn_cast<AbstractFunctionDecl>(decl)) {
@@ -4475,9 +4474,7 @@ fixMemberRef(ConstraintSystem &cs, Type baseTy,
              Optional<MemberLookupResult::UnviableReason> reason = None) {
   // Not all of the choices handled here are going
   // to refer to a declaration.
-  if (choice.isDecl()) {
-    auto *decl = choice.getDecl();
-
+  if (auto *decl = choice.getDeclOrNull()) {
     if (auto *CD = dyn_cast<ConstructorDecl>(decl)) {
       if (auto *fix = validateInitializerRef(cs, CD, locator))
         return fix;

lib/Sema/CSSolver.cpp

@@ -2092,11 +2092,9 @@ void ConstraintSystem::partitionDisjunction(
     forEachChoice(Choices, [&](unsigned index, Constraint *constraint) -> bool {
       if (constraint->getKind() != ConstraintKind::BindOverload)
         return false;
-      if (!constraint->getOverloadChoice().isDecl())
-        return false;
 
-      auto *decl = constraint->getOverloadChoice().getDecl();
-      auto *funcDecl = dyn_cast<FuncDecl>(decl);
+      auto *decl = constraint->getOverloadChoice().getDeclOrNull();
+      auto *funcDecl = dyn_cast_or_null<FuncDecl>(decl);
       if (!funcDecl)
         return false;
 

lib/Sema/ConstraintSystem.cpp

@@ -411,6 +411,44 @@ ConstraintLocator *ConstraintSystem::getConstraintLocator(
   return getConstraintLocator(anchor, path, builder.getSummaryFlags());
 }
 
+ConstraintLocator *ConstraintSystem::getCalleeLocator(Expr *expr) {
+  if (auto *applyExpr = dyn_cast<ApplyExpr>(expr)) {
+    auto *fnExpr = applyExpr->getFn();
+    // For an apply of a metatype, we have a short-form constructor. Unlike
+    // other locators to callees, these are anchored on the apply expression
+    // rather than the function expr.
+    if (simplifyType(getType(fnExpr))->is<AnyMetatypeType>()) {
+      auto *fnLocator =
+          getConstraintLocator(applyExpr, ConstraintLocator::ApplyFunction);
+      return getConstraintLocator(fnLocator,
+                                  ConstraintLocator::ConstructorMember);
+    }
+    // Otherwise fall through and look for locators anchored on the fn expr.
+    expr = fnExpr;
+  }
+
+  auto *locator = getConstraintLocator(expr);
+  if (auto *ude = dyn_cast<UnresolvedDotExpr>(expr)) {
+    if (TC.getSelfForInitDelegationInConstructor(DC, ude)) {
+      return getConstraintLocator(locator,
+                                  ConstraintLocator::ConstructorMember);
+    } else {
+      return getConstraintLocator(locator, ConstraintLocator::Member);
+    }
+  }
+
+  if (isa<UnresolvedMemberExpr>(expr))
+    return getConstraintLocator(locator, ConstraintLocator::UnresolvedMember);
+
+  if (isa<SubscriptExpr>(expr))
+    return getConstraintLocator(locator, ConstraintLocator::SubscriptMember);
+
+  if (isa<MemberRefExpr>(expr))
+    return getConstraintLocator(locator, ConstraintLocator::Member);
+
+  return locator;
+}
+
 Type ConstraintSystem::openUnboundGenericType(UnboundGenericType *unbound,
                                               ConstraintLocatorBuilder locator,
                                               OpenedTypeMap &replacements) {
@@ -2075,8 +2113,7 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
   }
   assert(!refType->hasTypeParameter() && "Cannot have a dependent type here");
 
-  if (choice.isDecl()) {
-    auto decl = choice.getDecl();
+  if (auto *decl = choice.getDeclOrNull()) {
     // If we're binding to an init member, the 'throws' need to line up between
     // the bound and reference types.
     if (auto CD = dyn_cast<ConstructorDecl>(decl)) {
@@ -2398,16 +2435,11 @@ bool ConstraintSystem::diagnoseAmbiguityWithFixes(
   if (solutions.empty())
     return false;
 
-  auto getOverloadDecl = [&](SelectedOverload &overload) -> ValueDecl * {
-    auto &choice = overload.choice;
-    return choice.isDecl() ? choice.getDecl() : nullptr;
-  };
-
   // Problems related to fixes forming ambiguous solution set
   // could only be diagnosed (at the moment), if all of the fixes
-  // are attached to the same anchor, which means they fix
-  // different overloads of the same declaration.
-  Expr *commonAnchor = nullptr;
+  // have the same callee locator, which means they fix different
+  // overloads of the same declaration.
+  ConstraintLocator *commonCalleeLocator = nullptr;
   SmallPtrSet<ValueDecl *, 4> distinctChoices;
   SmallVector<std::pair<const Solution *, const ConstraintFix *>, 4>
       viableSolutions;
@@ -2421,21 +2453,17 @@ bool ConstraintSystem::diagnoseAmbiguityWithFixes(
       return false;
 
     const auto *fix = fixes.front();
-    if (commonAnchor && commonAnchor != fix->getAnchor())
+    auto *calleeLocator = getCalleeLocator(fix->getAnchor());
+    if (commonCalleeLocator && commonCalleeLocator != calleeLocator)
       return false;
 
-    commonAnchor = fix->getAnchor();
+    commonCalleeLocator = calleeLocator;
 
-    SmallVector<SelectedOverload, 2> overloads;
-    solution.getOverloadChoices(commonAnchor, overloads);
-    // There is unfortunately no way, at the moment, to figure out
-    // what declaration the fix is attached to, so we have to make
-    // sure that there is only one declaration associated with common
-    // anchor to be sure that the right problem is being diagnosed.
-    if (overloads.size() != 1)
+    auto overload = solution.getOverloadChoiceIfAvailable(calleeLocator);
+    if (!overload)
       return false;
 
-    auto *decl = getOverloadDecl(overloads.front());
+    auto *decl = overload->choice.getDeclOrNull();
     if (!decl)
       return false;
 
@@ -2458,6 +2486,7 @@ bool ConstraintSystem::diagnoseAmbiguityWithFixes(
     DiagnosticTransaction transaction(TC.Diags);
 
     const auto *fix = viableSolutions.front().second;
+    auto *commonAnchor = commonCalleeLocator->getAnchor();
     if (fix->getKind() == FixKind::UseSubscriptOperator) {
       auto *UDE = cast<UnresolvedDotExpr>(commonAnchor);
       TC.diagnose(commonAnchor->getLoc(),