[Diagnostics] Port diagnostic for CTP_YieldByReference

lib/Sema/CSDiag.cpp

@@ -1683,24 +1683,6 @@ bool FailureDiagnosis::diagnoseContextualConversionError(
   if (isUnresolvedOrTypeVarType(exprType) || exprType->isEqual(contextualType))
     return false;
 
-  if (CTP == CTP_YieldByReference) {
-    if (auto contextualLV = contextualType->getAs<LValueType>())
-      contextualType = contextualLV->getObjectType();
-    if (auto exprLV = exprType->getAs<LValueType>()) {
-      diagnose(expr->getLoc(), diag::cannot_yield_wrong_type_by_reference,
-               exprLV->getObjectType(), contextualType);
-    } else if (exprType->isEqual(contextualType)) {
-      diagnose(expr->getLoc(), diag::cannot_yield_rvalue_by_reference_same_type,
-               exprType);
-    } else {
-      diagnose(expr->getLoc(), diag::cannot_yield_rvalue_by_reference, exprType,
-               contextualType);
-    }
-    return true;
-  }
-
-  exprType = exprType->getRValueType();
-
   // Don't attempt fixits if we have an unsolved type variable, since
   // the recovery path's recursion into the type checker via typeCheckCast()
   // will confuse matters.

lib/Sema/CSDiagnostics.cpp

@@ -1885,6 +1885,9 @@ bool ContextualFailure::diagnoseAsError() {
     if (diagnoseThrowsTypeMismatch())
       return true;
 
+    if (diagnoseYieldByReferenceMismatch())
+      return true;
+
     auto contextualType = getToType();
     if (auto msg = getDiagnosticFor(CTP, contextualType->isExistentialType())) {
       diagnostic = *msg;
@@ -2238,6 +2241,27 @@ bool ContextualFailure::diagnoseThrowsTypeMismatch() const {
   return true;
 }
 
+bool ContextualFailure::diagnoseYieldByReferenceMismatch() const {
+  if (CTP != CTP_YieldByReference)
+    return false;
+
+  auto *anchor = getAnchor();
+  auto exprType = getType(anchor);
+  auto contextualType = getToType();
+
+  if (auto exprLV = exprType->getAs<LValueType>()) {
+    emitDiagnostic(anchor->getLoc(), diag::cannot_yield_wrong_type_by_reference,
+                   exprLV->getObjectType(), contextualType);
+  } else if (exprType->isEqual(contextualType)) {
+    emitDiagnostic(anchor->getLoc(),
+                   diag::cannot_yield_rvalue_by_reference_same_type, exprType);
+  } else {
+    emitDiagnostic(anchor->getLoc(), diag::cannot_yield_rvalue_by_reference,
+                   exprType, contextualType);
+  }
+  return true;
+}
+
 bool ContextualFailure::tryRawRepresentableFixIts(
     InFlightDiagnostic &diagnostic,
     KnownProtocolKind rawRepresentableProtocol) const {

lib/Sema/CSDiagnostics.h

@@ -633,6 +633,10 @@ class ContextualFailure : public FailureDiagnostic {
   /// something with doesn't conform to `Error`.
   bool diagnoseThrowsTypeMismatch() const;
 
+  /// Produce a specialized diagnostic if this is an attempt to `yield`
+  /// something of incorrect type.
+  bool diagnoseYieldByReferenceMismatch() const;
+
   /// Attempt to attach any relevant fix-its to already produced diagnostic.
   void tryFixIts(InFlightDiagnostic &diagnostic) const;
 

lib/Sema/CSSimplify.cpp

@@ -2028,15 +2028,36 @@ ConstraintSystem::matchTypesBindTypeVar(
     type = type->getRValueType();
   }
 
+  // Attempt to fix situations where type variable can't be bound
+  // to a particular type e.g. `l-value` or `inout`.
+  auto fixReferenceMismatch = [&](TypeVariableType *typeVar,
+                                  Type type) -> bool {
+    if (auto last = locator.last()) {
+      if (last->is<LocatorPathElt::ContextualType>()) {
+        auto *fix = IgnoreContextualType::create(*this, typeVar, type,
+                                                 getConstraintLocator(locator));
+        return !recordFix(fix);
+      }
+    }
+
+    return false;
+  };
+
   // If the left-hand type variable cannot bind to an lvalue,
   // but we still have an lvalue, fail.
   if (!typeVar->getImpl().canBindToLValue() && type->hasLValueType()) {
-      return getTypeMatchFailure(locator);
+    if (shouldAttemptFixes() && fixReferenceMismatch(typeVar, type))
+      return getTypeMatchSuccess();
+
+    return getTypeMatchFailure(locator);
   }
 
   // If the left-hand type variable cannot bind to an inout,
   // but we still have an inout, fail.
   if (!typeVar->getImpl().canBindToInOut() && type->is<InOutType>()) {
+    if (shouldAttemptFixes() && fixReferenceMismatch(typeVar, type))
+      return getTypeMatchSuccess();
+
     return getTypeMatchFailure(locator);
   }
 
@@ -2399,9 +2420,31 @@ bool ConstraintSystem::repairFailures(
                         });
   };
 
-  auto &elt = path.back();
+  auto elt = path.back();
   switch (elt.getKind()) {
-  case ConstraintLocator::LValueConversion:
+  case ConstraintLocator::LValueConversion: {
+    auto CTP = getContextualTypePurpose();
+    // Special case for `CTP_CallArgument` set by CSDiag
+    // while type-checking each argument because we yet
+    // to cover argument-to-parameter conversions in the
+    // new framework.
+    if (CTP != CTP_CallArgument) {
+      // Ignore l-value conversion element since it has already
+      // played its role.
+      path.pop_back();
+      // If this is a contextual mismatch between l-value types e.g.
+      // `@lvalue String vs. @lvalue Int`, let's pretend that it's okay.
+      if (!path.empty() && path.back().is<LocatorPathElt::ContextualType>()) {
+        auto *locator = getConstraintLocator(anchor, path.back());
+        conversionsOrFixes.push_back(
+            IgnoreContextualType::create(*this, lhs, rhs, locator));
+        break;
+      }
+    }
+
+    LLVM_FALLTHROUGH;
+  }
+
   case ConstraintLocator::ApplyArgToParam: {
     auto loc = getConstraintLocator(locator);
     if (repairByInsertingExplicitCall(lhs, rhs))

test/stmt/yield.swift

@@ -25,7 +25,7 @@ struct YieldVariables {
     }
     _modify {
       var x = 0
-      yield &x // expected-error {{cannot yield immutable value of type 'Int' as an inout yield of type 'String'}}
+      yield &x // expected-error {{cannot yield reference to storage of type 'Int' as an inout yield of type 'String'}}
     }
   }