[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});