[SR-15281] [Sema] Couple of contextual mismatch and runtime cast diagnostic fixes

include/swift/Sema/ConstraintSystem.h

@@ -2180,6 +2180,10 @@ class SolutionApplicationTarget {
 using RewriteTargetFn = std::function<
     Optional<SolutionApplicationTarget> (SolutionApplicationTarget)>;
 
+/// Represents a conversion restriction between two types.
+using ConversionRestriction =
+    std::tuple<TypeBase *, TypeBase *, ConversionRestrictionKind>;
+
 enum class ConstraintSystemPhase {
   ConstraintGeneration,
   Solving,
@@ -2364,8 +2368,7 @@ class ConstraintSystem {
   /// there are multiple ways in which one type could convert to another, e.g.,
   /// given class types A and B, the solver might choose either a superclass
   /// conversion or a user-defined conversion.
-  std::vector<std::tuple<Type, Type, ConversionRestrictionKind>>
-      ConstraintRestrictions;
+  std::vector<ConversionRestriction> ConstraintRestrictions;
 
   /// The set of fixes applied to make the solution work.
   llvm::SmallVector<ConstraintFix *, 4> Fixes;

lib/Sema/CSDiagnostics.cpp

@@ -2953,6 +2953,20 @@ bool ContextualFailure::tryTypeCoercionFixIt(
   if (!toType->hasTypeRepr())
     return false;
 
+  // If object of the optional type is a subtype of the specified contextual
+  // type, let's suggest a force unwrap "!". Otherwise fallback to potential
+  // coercion or force cast.
+  if (!bothOptional && fromType->getOptionalObjectType()) {
+    if (TypeChecker::isSubtypeOf(fromType->lookThroughAllOptionalTypes(),
+                                 toType, getDC())) {
+      diagnostic.fixItInsert(
+          Lexer::getLocForEndOfToken(getASTContext().SourceMgr,
+                                     getSourceRange().End),
+          "!");
+      return true;
+    }
+  }
+
   CheckedCastKind Kind =
       TypeChecker::typeCheckCheckedCast(fromType, toType,
                                         CheckedCastContextKind::None, getDC(),

lib/Sema/CSSimplify.cpp

@@ -3902,9 +3902,8 @@ bool ConstraintSystem::repairFailures(
       // If the result type of the coercion has an value to optional conversion
       // we can instead suggest the conditional downcast as it is safer in
       // situations like conditional binding.
-      auto useConditionalCast = llvm::any_of(
-          ConstraintRestrictions,
-          [&](std::tuple<Type, Type, ConversionRestrictionKind> restriction) {
+      auto useConditionalCast =
+          llvm::any_of(ConstraintRestrictions, [&](auto &restriction) {
             ConversionRestrictionKind restrictionKind;
             Type type1, type2;
             std::tie(type1, type2, restrictionKind) = restriction;
@@ -6709,7 +6708,9 @@ static bool isCastToExpressibleByNilLiteral(ConstraintSystem &cs, Type fromType,
 static ConstraintFix *maybeWarnAboutExtraneousCast(
     ConstraintSystem &cs, Type origFromType, Type origToType, Type fromType,
     Type toType, SmallVector<Type, 4> fromOptionals,
-    SmallVector<Type, 4> toOptionals, ConstraintSystem::TypeMatchOptions flags,
+    SmallVector<Type, 4> toOptionals,
+    const std::vector<ConversionRestriction> &constraintRestrictions,
+    ConstraintSystem::TypeMatchOptions flags,
     ConstraintLocatorBuilder locator) {
 
   auto last = locator.last();
@@ -6731,6 +6732,18 @@ static ConstraintFix *maybeWarnAboutExtraneousCast(
   // "from" could be less optional than "to" e.g. `0 as Any?`, so
   // we need to store the difference as a signed integer.
   int extraOptionals = fromOptionals.size() - toOptionals.size();
+
+  // "from" expression could be a type variable with value-to-optional
+  // restrictions that we have to account for optionality mismatch.
+  const auto subExprType = cs.getType(castExpr->getSubExpr());
+  if (llvm::is_contained(
+          constraintRestrictions,
+          std::make_tuple(fromType.getPointer(), subExprType.getPointer(),
+                          ConversionRestrictionKind::ValueToOptional))) {
+    extraOptionals++;
+    origFromType = OptionalType::get(origFromType);
+  }
+
   // Removing the optionality from to type when the force cast expr is an IUO.
   const auto *const TR = castExpr->getCastTypeRepr();
   if (isExpr<ForcedCheckedCastExpr>(anchor) && TR &&
@@ -6886,7 +6899,7 @@ ConstraintSystem::simplifyCheckedCastConstraint(
 
     if (auto *fix = maybeWarnAboutExtraneousCast(
             *this, origFromType, origToType, fromType, toType, fromOptionals,
-            toOptionals, flags, locator)) {
+            toOptionals, ConstraintRestrictions, flags, locator)) {
       (void)recordFix(fix);
     }
   };
@@ -6954,7 +6967,7 @@ ConstraintSystem::simplifyCheckedCastConstraint(
     // succeed or fail.
     if (auto *fix = maybeWarnAboutExtraneousCast(
             *this, origFromType, origToType, fromType, toType, fromOptionals,
-            toOptionals, flags, locator)) {
+            toOptionals, ConstraintRestrictions, flags, locator)) {
       (void)recordFix(fix);
     }
 
@@ -11382,7 +11395,8 @@ ConstraintSystem::simplifyRestrictedConstraint(
       addFixConstraint(fix, matchKind, type1, type2, locator);
     }
 
-    ConstraintRestrictions.push_back(std::make_tuple(type1, type2, restriction));
+    ConstraintRestrictions.push_back(
+        std::make_tuple(type1.getPointer(), type2.getPointer(), restriction));
     return SolutionKind::Solved;
   }
   case SolutionKind::Unsolved:

lib/Sema/CSSolver.cpp

@@ -231,9 +231,10 @@ void ConstraintSystem::applySolution(const Solution &solution) {
   // Register constraint restrictions.
   // FIXME: Copy these directly into some kind of partial solution?
   for (auto restriction : solution.ConstraintRestrictions) {
-    ConstraintRestrictions.push_back(
-        std::make_tuple(restriction.first.first, restriction.first.second,
-                        restriction.second));
+    auto &types = restriction.first;
+    ConstraintRestrictions.push_back(std::make_tuple(types.first.getPointer(),
+                                                     types.second.getPointer(),
+                                                     restriction.second));
   }
 
   // Register the solution's disjunction choices.

test/Constraints/casts.swift

@@ -613,3 +613,41 @@ func decodeStringOrIntDictionary<T: FixedWidthInteger>() -> [Int: T] {
     fatalError()
   }
 }
+
+
+// SR-15281
+struct SR15281_A { }
+struct SR15281_B {
+  init(a: SR15281_A) { }
+}
+
+struct SR15281_S {
+  var a: SR15281_A? = SR15281_A()
+
+  var b: SR15281_B {
+    a.flatMap(SR15281_B.init(a:)) // expected-error{{cannot convert return expression of type 'SR15281_B?' to return type 'SR15281_B'}} {{34-34=!}}
+  }
+
+  var b1: SR15281_B {
+    a.flatMap(SR15281_B.init(a:)) as! SR15281_B 
+    // expected-warning@-1 {{forced cast from 'SR15281_B?' to 'SR15281_B' only unwraps optionals; did you mean to use '!'?}} {{34-34=!}} {{34-48=}}
+  }
+}
+
+class SR15281_AC {}
+class SR15281_BC {
+  init(a: SR15281_AC) { }
+}
+class SR15281_CC: SR15281_BC {}
+
+struct SR15281_SC {
+  var a: SR15281_AC? = SR15281_AC()
+
+  var b: SR15281_BC {
+    a.flatMap(SR15281_BC.init(a:)) // expected-error{{cannot convert return expression of type 'SR15281_BC?' to return type 'SR15281_BC'}} {{35-35=!}}
+  }
+
+  var c: SR15281_BC {
+    a.flatMap(SR15281_CC.init(a:)) // expected-error{{cannot convert return expression of type 'SR15281_CC?' to return type 'SR15281_BC'}} {{35-35=!}}
+  }
+}