IUO: Put the pieces in place to handle coercions and casts to IUO types.

For casts like 'x as T!' or 'x as! T!', generate a disjunction that
attempts to bind both the optional and non-optional T, preferring the
optional branch of the disjunction.

Note that this should effectively be NFC at the moment because we're
still generating the IUO type for T! rather than a plain optional, so
we first attempt the IUO type (as opposed to an Optional<T>) which
goes through existing logic that handles the IUOs as part of the type
system.

The new rewriting logic will actually do something once we switch over
to generating Optional<T> when T! is uttered.

Author: rudkx

lib/Sema/CSApply.cpp

@@ -3170,7 +3170,7 @@ namespace {
           cast->setImplicit();
 
         // Type-check this conditional case.
-        Expr *result = visitConditionalCheckedCastExpr(cast, true);
+        Expr *result = handleConditionalCheckedCastExpr(cast, true);
         if (!result)
           return nullptr;
 
@@ -3409,7 +3409,28 @@ namespace {
       return addFinalOptionalInjections(result);
     }
 
+    bool hasForcedOptionalResult(ExplicitCastExpr *expr) {
+      auto *TR = expr->getCastTypeLoc().getTypeRepr();
+      if (TR && TR->getKind() == TypeReprKind::ImplicitlyUnwrappedOptional) {
+        auto *locator = cs.getConstraintLocator(
+            expr, ConstraintLocator::ImplicitlyUnwrappedCoercionResult);
+        return solution.getDisjunctionChoice(locator);
+      }
+      return false;
+    }
+
     Expr *visitCoerceExpr(CoerceExpr *expr) {
+      // If we need to insert a force-unwrap for coercions of the form
+      // 'as T!', do so now.
+      if (hasForcedOptionalResult(expr)) {
+        auto *coerced = visitCoerceExpr(expr, None);
+        if (!coerced)
+          return nullptr;
+
+        return coerceImplicitlyUnwrappedOptionalToValue(
+            coerced, cs.getType(coerced)->getOptionalObjectType());
+      }
+
       return visitCoerceExpr(expr, None);
     }
 
@@ -3482,7 +3503,24 @@ namespace {
       return expr;
     }
 
+    // Rewrite ForcedCheckedCastExpr based on what the solver computed.
     Expr *visitForcedCheckedCastExpr(ForcedCheckedCastExpr *expr) {
+      // If we need to insert a force-unwrap for coercions of the form
+      // 'as! T!', do so now.
+      if (hasForcedOptionalResult(expr)) {
+        auto *coerced = handleForcedCheckedCastExpr(expr);
+        if (!coerced)
+          return nullptr;
+
+        return coerceImplicitlyUnwrappedOptionalToValue(
+            coerced, cs.getType(coerced)->getOptionalObjectType());
+      }
+
+      return handleForcedCheckedCastExpr(expr);
+    }
+
+    // Most of the logic for dealing with ForcedCheckedCastExpr.
+    Expr *handleForcedCheckedCastExpr(ForcedCheckedCastExpr *expr) {
       // Simplify the type we're casting to.
       auto toType = simplifyType(expr->getCastTypeLoc().getType());
       expr->getCastTypeLoc().setType(toType, /*validated=*/true);
@@ -3546,8 +3584,23 @@ namespace {
                                            OptionalBindingsCastKind::Forced);
     }
 
-    Expr *visitConditionalCheckedCastExpr(ConditionalCheckedCastExpr *expr,
-                                          bool isInsideIsExpr = false) {
+    Expr *visitConditionalCheckedCastExpr(ConditionalCheckedCastExpr *expr) {
+      // If we need to insert a force-unwrap for coercions of the form
+      // 'as! T!', do so now.
+      if (hasForcedOptionalResult(expr)) {
+        auto *coerced = handleConditionalCheckedCastExpr(expr);
+        if (!coerced)
+          return nullptr;
+
+        return coerceImplicitlyUnwrappedOptionalToValue(
+            coerced, cs.getType(coerced)->getOptionalObjectType());
+      }
+
+      return handleConditionalCheckedCastExpr(expr);
+    }
+
+    Expr *handleConditionalCheckedCastExpr(ConditionalCheckedCastExpr *expr,
+                                           bool isInsideIsExpr = false) {
       // Simplify the type we're casting to.
       auto toType = simplifyType(expr->getCastTypeLoc().getType());
       checkForImportedUsedConformances(toType);

lib/Sema/CSGen.cpp

@@ -2494,6 +2494,18 @@ namespace {
       llvm_unreachable("Already type-checked");
     }
 
+    Type
+    createTypeVariableAndDisjunctionForIUOCoercion(Type toType,
+                                                   ConstraintLocator *locator) {
+      auto implicitUnwrapLocator = CS.getConstraintLocator(
+          locator, ConstraintLocator::ImplicitlyUnwrappedCoercionResult);
+      auto typeVar =
+          CS.createTypeVariable(implicitUnwrapLocator, /*options=*/0);
+      CS.buildDisjunctionForImplicitlyUnwrappedOptional(typeVar, toType,
+                                                        implicitUnwrapLocator);
+      return typeVar;
+    }
+
     Type visitForcedCheckedCastExpr(ForcedCheckedCastExpr *expr) {
       auto &tc = CS.getTypeChecker();
       auto fromExpr = expr->getSubExpr();
@@ -2517,6 +2529,12 @@ namespace {
       // The source type can be checked-cast to the destination type.
       CS.addConstraint(ConstraintKind::CheckedCast, fromType, toType, locator);
 
+      // If the result type was declared IUO, add a disjunction for
+      // bindings for the result of the coercion.
+      auto *TR = expr->getCastTypeLoc().getTypeRepr();
+      if (TR && TR->getKind() == TypeReprKind::ImplicitlyUnwrappedOptional)
+        return createTypeVariableAndDisjunctionForIUOCoercion(toType, locator);
+
       return toType;
     }
 
@@ -2537,8 +2555,17 @@ namespace {
       auto fromType = CS.getType(expr->getSubExpr());
       auto locator = CS.getConstraintLocator(expr);
 
-      CS.addExplicitConversionConstraint(fromType, toType, /*allowFixes=*/true,
-                                         locator);
+      // Add a conversion constraint for the direct conversion between
+      // types.
+      CS.addExplicitConversionConstraint(fromType, toType,
+                                         /*allowFixes=*/true, locator);
+
+      // If the result type was declared IUO, add a disjunction for
+      // bindings for the result of the coercion.
+      auto *TR = expr->getCastTypeLoc().getTypeRepr();
+      if (TR && TR->getKind() == TypeReprKind::ImplicitlyUnwrappedOptional)
+        return createTypeVariableAndDisjunctionForIUOCoercion(toType, locator);
+
       return toType;
     }
 
@@ -2561,7 +2588,16 @@ namespace {
 
       auto fromType = CS.getType(fromExpr);
       auto locator = CS.getConstraintLocator(expr);
+
       CS.addConstraint(ConstraintKind::CheckedCast, fromType, toType, locator);
+
+      // If the result type was declared IUO, add a disjunction for
+      // bindings for the result of the coercion.
+      auto *TR = expr->getCastTypeLoc().getTypeRepr();
+      if (TR && TR->getKind() == TypeReprKind::ImplicitlyUnwrappedOptional)
+        return createTypeVariableAndDisjunctionForIUOCoercion(
+            OptionalType::get(toType), locator);
+
       return OptionalType::get(toType);
     }