[Sema] Adjust getCalleeLocator to look into coercion operand

Author: LucianoPAlmeida

include/swift/Sema/ConstraintSystem.h

@@ -3249,12 +3249,6 @@ class ConstraintSystem {
   ConstraintLocator *
   getConstraintLocator(const ConstraintLocatorBuilder &builder);
 
-  /// Retrieves a locator that identifies where an overload ambiguity may be
-  /// present. In most cases this is a no-op, but may simplify locator e.g
-  /// coercions for ambiguities in the expression operand.
-  ConstraintLocator *
-  getConstraintLocatorForAmbiguity(ConstraintLocator *locator);
-
   /// Compute a constraint locator for an implicit value-to-value
   /// conversion rooted at the given location.
   ConstraintLocator *

lib/Sema/CSApply.cpp

@@ -4275,15 +4275,8 @@ namespace {
     bool hasForcedOptionalResult(ExplicitCastExpr *expr) {
       const auto *const TR = expr->getCastTypeRepr();
       if (TR && TR->getKind() == TypeReprKind::ImplicitlyUnwrappedOptional) {
-        ConstraintLocator *locator;
-        if (isExpr<CoerceExpr>(expr)) {
-          locator = cs.getConstraintLocator(
-              expr, {LocatorPathElt::CoercionOperand(),
-                     ConstraintLocator::ImplicitlyUnwrappedDisjunctionChoice});
-        } else {
-          locator = cs.getConstraintLocator(
-              expr, ConstraintLocator::ImplicitlyUnwrappedDisjunctionChoice);
-        }
+        auto *locator = cs.getConstraintLocator(
+            expr, ConstraintLocator::ImplicitlyUnwrappedDisjunctionChoice);
         return solution.getDisjunctionChoice(locator);
       }
       return false;
@@ -4357,7 +4350,7 @@ namespace {
       // get it from the solution to determine whether we've picked a coercion
       // or a bridging conversion.
       auto *locator =
-          cs.getConstraintLocator(expr, LocatorPathElt::CoercionOperand());
+          cs.getConstraintLocator(expr, ConstraintLocator::CoercionOperand);
       auto choice = solution.getDisjunctionChoice(locator);
 
       // Handle the coercion/bridging of the underlying subexpression, where

lib/Sema/CSDiagnostics.cpp

@@ -1225,21 +1225,28 @@ ASTNode InvalidCoercionFailure::getAnchor() const {
   return anchor;
 }
 
+SourceLoc InvalidCoercionFailure::getLoc() const {
+  if (getLocator()->isForCoercion()) {
+    auto *CE = castToExpr<CoerceExpr>(getRawAnchor());
+    return CE->getAsLoc();
+  }
+  return FailureDiagnostic::getLoc();
+}
+
 bool InvalidCoercionFailure::diagnoseAsError() {
   auto fromType = getFromType();
   auto toType = getToType();
-  auto *CE = getAsExpr<CoerceExpr>(getRawAnchor());
 
   emitDiagnostic(diag::cannot_coerce_to_type, fromType, toType);
 
   if (UseConditionalCast) {
     emitDiagnostic(diag::missing_optional_downcast)
-        .highlight(CE->getSourceRange())
-        .fixItReplace(CE->getAsLoc(), "as?");
+        .highlight(getSourceRange())
+        .fixItReplace(getLoc(), "as?");
   } else {
     emitDiagnostic(diag::missing_forced_downcast)
-        .highlight(CE->getSourceRange())
-        .fixItReplace(CE->getAsLoc(), "as!");
+        .highlight(getSourceRange())
+        .fixItReplace(getLoc(), "as!");
   }
 
   return true;
@@ -2662,7 +2669,7 @@ bool ContextualFailure::diagnoseAsError() {
 }
 
 bool ContextualFailure::diagnoseAsNote() {
-  auto *locator = getAmbiguityLocator();
+  auto *locator = getLocator();
 
   auto overload = getCalleeOverloadChoiceIfAvailable(locator);
   if (!(overload && overload->choice.isDecl()))
@@ -4931,7 +4938,7 @@ bool MissingArgumentsFailure::diagnoseAsError() {
 }
 
 bool MissingArgumentsFailure::diagnoseAsNote() {
-  auto *locator = getAmbiguityLocator();
+  auto *locator = getLocator();
   if (auto overload = getCalleeOverloadChoiceIfAvailable(locator)) {
     auto *fn = resolveType(overload->adjustedOpenedType)->getAs<AnyFunctionType>();
     auto loc = overload->choice.getDecl()->getLoc();
@@ -5733,16 +5740,16 @@ bool ExtraneousArgumentsFailure::diagnoseAsError() {
 }
 
 bool ExtraneousArgumentsFailure::diagnoseAsNote() {
-  auto overload = getCalleeOverloadChoiceIfAvailable(getAmbiguityLocator());
+  auto overload = getCalleeOverloadChoiceIfAvailable(getLocator());
   if (!(overload && overload->choice.isDecl()))
     return false;
 
   auto *decl = overload->choice.getDecl();
   auto numArgs = getTotalNumArguments();
   if (isExpr<ClosureExpr>(getAnchor())) {
     emitDiagnosticAt(decl, diag::candidate_with_extraneous_args_closure,
-                    ContextualType, ContextualType->getNumParams(), numArgs,
-                    (numArgs == 1));
+                     ContextualType, ContextualType->getNumParams(), numArgs,
+                     (numArgs == 1));
   } else {
     emitDiagnosticAt(decl, diag::candidate_with_extraneous_args,
                      overload->adjustedOpenedType, numArgs, ContextualType,

lib/Sema/CSDiagnostics.h

@@ -90,11 +90,6 @@ class FailureDiagnostic {
 
   ConstraintLocator *getLocator() const { return Locator; }
 
-  ConstraintLocator *getAmbiguityLocator() const {
-    auto &cs = getConstraintSystem();
-    return cs.getConstraintLocatorForAmbiguity(Locator);
-  }
-
   Type getType(ASTNode node, bool wantRValue = true) const;
 
   /// Get type associated with a given ASTNode without resolving it,
@@ -2182,6 +2177,8 @@ class InvalidCoercionFailure final : public ContextualFailure {
 
   ASTNode getAnchor() const override;
 
+  SourceLoc getLoc() const override;
+
   bool diagnoseAsError() override;
 };
 

lib/Sema/CSGen.cpp

@@ -3307,7 +3307,7 @@ namespace {
 
       auto fromType = CS.getType(expr->getSubExpr());
       auto locator =
-          CS.getConstraintLocator(expr, LocatorPathElt::CoercionOperand());
+          CS.getConstraintLocator(expr, ConstraintLocator::CoercionOperand);
 
       // Literal initialization (e.g. `UInt32(0)`) doesn't require
       // a conversion because the literal is supposed to assume the
@@ -3327,8 +3327,12 @@ namespace {
 
       // If the result type was declared IUO, add a disjunction for
       // bindings for the result of the coercion.
-      if (repr && repr->getKind() == TypeReprKind::ImplicitlyUnwrappedOptional)
-        return createTypeVariableAndDisjunctionForIUOCoercion(toType, locator);
+      if (repr &&
+          repr->getKind() == TypeReprKind::ImplicitlyUnwrappedOptional) {
+        auto *locatorIUO = CS.getConstraintLocator(expr);
+        return createTypeVariableAndDisjunctionForIUOCoercion(toType,
+                                                              locatorIUO);
+      }
 
       return toType;
     }

lib/Sema/ConstraintSystem.cpp

@@ -422,14 +422,6 @@ ConstraintLocator *ConstraintSystem::getConstraintLocator(
   return getConstraintLocator(anchor, newPath);
 }
 
-ConstraintLocator *
-ConstraintSystem::getConstraintLocatorForAmbiguity(ConstraintLocator *locator) {
-  if (locator->findLast<LocatorPathElt::CoercionOperand>()) {
-    return getConstraintLocator(simplifyLocatorToAnchor(locator));
-  }
-  return locator;
-}
-
 ConstraintLocator *ConstraintSystem::getImplicitValueConversionLocator(
     ConstraintLocatorBuilder root, ConversionRestrictionKind restriction) {
   SmallVector<LocatorPathElt, 4> path;
@@ -630,6 +622,14 @@ ConstraintLocator *ConstraintSystem::getCalleeLocator(
   if (isExpr<ObjectLiteralExpr>(anchor))
     return getConstraintLocator(anchor, ConstraintLocator::ConstructorMember);
 
+  if (locator->isFirstElement<LocatorPathElt::CoercionOperand>()) {
+    auto *CE = castToExpr<CoerceExpr>(anchor);
+    locator = getConstraintLocator(CE->getSubExpr()->getValueProvidingExpr(),
+                                   path.drop_front());
+    return getCalleeLocator(locator, lookThroughApply, getType, simplifyType,
+                            getOverloadFor);
+  }
+
   return getConstraintLocator(anchor);
 }
 
@@ -4980,7 +4980,7 @@ bool ConstraintSystem::diagnoseAmbiguityWithFixes(
     const auto &solution = *entry.first;
     const auto *fix = entry.second;
 
-    auto *locator = getConstraintLocatorForAmbiguity(fix->getLocator());
+    auto *locator = fix->getLocator();
 
     if (locator->isForContextualType()) {
       contextualFixes.push_back({&solution, fix});