[ConstraintSystem] Implicitly force results of @optional protocol requirements.

We were not handling IUO results of @optional protocol methods
properly, sometimes forcing the @optional requirement rather than the
result of the call.

Fixes rdar://problem/37240984.

Author: rudkx

lib/Sema/CSApply.cpp

@@ -1065,7 +1065,8 @@ namespace {
         // We also need to handle the implicitly unwrap of the result
         // of the called function if that's the type checking solution
         // we ended up with.
-        return forceUnwrapIfExpected(ref, member, memberLocator);
+        return forceUnwrapIfExpected(ref, member, memberLocator,
+                                     member->getAttrs().hasAttribute<OptionalAttr>());
       }
 
       // For types and properties, build member references.
@@ -2396,9 +2397,22 @@ namespace {
       return expr;
     }
 
-    Expr *forceUnwrapResult(Expr *expr) {
+    // Add a forced unwrap of an expression which either has type Optional<T>
+    // or is a function that returns an Optional<T>. The latter turns into a
+    // conversion expression that we will hoist above the ApplyExpr
+    // that needs to be forced during the process of rewriting the expression.
+    //
+    // forForcedOptional is used to indicate that we will further need
+    // to hoist this result above an explicit force of an optional that is
+    // in place for something like an @optional protocol member from
+    // Objective C that we might otherwise mistake for the thing we mean to
+    // force here.
+    Expr *forceUnwrapResult(Expr *expr, bool forForcedOptional =false) {
       auto ty = simplifyType(cs.getType(expr));
 
+      if (forForcedOptional)
+        ty = ty->getOptionalObjectType();
+
       if (auto *fnTy = ty->getAs<AnyFunctionType>()) {
         auto underlyingType = cs.replaceFinalResultTypeWithUnderlying(fnTy);
 
@@ -2422,12 +2436,13 @@ namespace {
     }
 
     Expr *forceUnwrapIfExpected(Expr *expr, Decl *decl,
-                                ConstraintLocatorBuilder locator) {
+                                ConstraintLocatorBuilder locator,
+                                bool forForcedOptional =false) {
       if (!shouldForceUnwrapResult(decl, locator))
         return expr;
 
       // Force the expression if required for the solution.
-      return forceUnwrapResult(expr);
+      return forceUnwrapResult(expr, forForcedOptional);
     }
 
     Expr *visitDeclRefExpr(DeclRefExpr *expr) {
@@ -3774,9 +3789,29 @@ namespace {
     }
 
     Expr *visitForceValueExpr(ForceValueExpr *expr) {
+      // Check to see if we are forcing an
+      // ImplicitlyUnwrappedFunctionConversionExpr.  This can happen
+      // in cases where we had a ForceValueExpr of an optional for a
+      // declaration for a function whose result type we need to
+      // implicitly force after applying. We need to hoist the function
+      // conversion above the ForceValueExpr, so that we may ultimately
+      // hoist it above the ApplyExpr where we will eventually rewrite the
+      // function conversion into a force of the result.
+      Expr *replacement = expr;
+      if (auto fnConv =
+          dyn_cast<ImplicitlyUnwrappedFunctionConversionExpr>(expr->getSubExpr())) {
+        auto fnConvSubExpr = fnConv->getSubExpr();
+        auto fnConvSubObjTy =
+          cs.getType(fnConvSubExpr)->getOptionalObjectType();
+        cs.setType(expr, fnConvSubObjTy);
+        expr->setSubExpr(fnConvSubExpr);
+        fnConv->setSubExpr(expr);
+        replacement = fnConv;
+      }
+
       Type valueType = simplifyType(cs.getType(expr));
       cs.setType(expr, valueType);
-      
+
       // Coerce the object type, if necessary.
       auto subExpr = expr->getSubExpr();
       if (auto objectTy = cs.getType(subExpr)->getOptionalObjectType()) {
@@ -3789,7 +3824,7 @@ namespace {
         }
       }
 
-      return expr;
+      return replacement;
     }
 
     Expr *visitOpenExistentialExpr(OpenExistentialExpr *expr) {

lib/Sema/ConstraintSystem.cpp

@@ -1565,7 +1565,7 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
   Type openedFullType;
 
   bool isDynamicResult = choice.getKind() == OverloadChoiceKind::DeclViaDynamic;
-  bool createdDynamicResultDisjunction = false;
+  bool bindConstraintCreated = false;
 
   switch (auto kind = choice.getKind()) {
   case OverloadChoiceKind::Decl:
@@ -1622,6 +1622,19 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
       //
       // Subscript declarations are handled within
       // getTypeOfMemberReference(); their result types are optional.
+
+      // Deal with values declared as implicitly unwrapped, or
+      // functions with return types that are implicitly unwrapped.
+      if (choice.isImplicitlyUnwrappedValueOrReturnValue()) {
+        // Build the disjunction to attempt binding both T? and T (or
+        // function returning T? and function returning T).
+        Type ty = createTypeVariable(locator, TVO_CanBindToInOut);
+        buildDisjunctionForImplicitlyUnwrappedOptional(ty, refType,
+                                                       locator);
+        addConstraint(ConstraintKind::Bind, boundType,
+                      OptionalType::get(ty->getRValueType()), locator);
+        bindConstraintCreated = true;
+      }
       refType = OptionalType::get(refType->getRValueType());
     }
     // For a non-subscript declaration found via dynamic lookup, strip
@@ -1696,7 +1709,7 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
         refType = OptionalType::get(refType->getRValueType());
       }
 
-      createdDynamicResultDisjunction = true;
+      bindConstraintCreated = true;
     }
 
     // If the declaration is unavailable, note that in the score.
@@ -1785,8 +1798,8 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
                                               openedFullType,
                                               refType};
 
-  // We created appropriate disjunctions for dynamic result above.
-  if (!createdDynamicResultDisjunction) {
+  // In some cases we already created the appropriate bind constraints.
+  if (!bindConstraintCreated) {
     if (choice.isImplicitlyUnwrappedValueOrReturnValue()) {
       // Build the disjunction to attempt binding both T? and T (or
       // function returning T? and function returning T).

test/Constraints/iuo_objc.swift

@@ -0,0 +1,31 @@
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -typecheck -verify %s
+// REQUIRES: objc_interop
+
+import Foundation
+
+func iuo_error(prop: IUOProperty) {
+  let _: Coat? = prop.iuo.optional()
+  // expected-error@-1 {{value of optional type '(() -> Coat?)?' not unwrapped; did you mean to use '!' or '?'?}}
+  let _: Coat? = prop.iuo.optional()!
+  // expected-error@-1 {{cannot invoke 'optional' with no arguments}}
+  let _: Coat? = prop.iuo.optional!()
+  let _: Coat? = prop.iuo.optional!()!
+  let _: Coat? = prop.iuo!.optional()
+  // expected-error@-1 {{value of optional type '(() -> Coat?)?' not unwrapped; did you mean to use '!' or '?'?}}
+  let _: Coat? = prop.iuo!.optional()!
+  // expected-error@-1 {{cannot invoke 'optional' with no arguments}}
+  let _: Coat? = prop.iuo!.optional!()
+  let _: Coat? = prop.iuo!.optional!()!
+  let _: Coat = prop.iuo.optional()
+  // expected-error@-1 {{value of optional type '(() -> Coat)?' not unwrapped; did you mean to use '!' or '?'?}}
+  let _: Coat = prop.iuo.optional()!
+  // expected-error@-1 {{cannot invoke 'optional' with no arguments}}
+  let _: Coat = prop.iuo.optional!()
+  let _: Coat = prop.iuo.optional!()!
+  let _: Coat = prop.iuo!.optional()
+  // expected-error@-1 {{value of optional type '(() -> Coat)?' not unwrapped; did you mean to use '!' or '?'?}}
+  let _: Coat = prop.iuo!.optional()!
+  // expected-error@-1 {{cannot invoke 'optional' with no arguments}}
+  let _: Coat = prop.iuo!.optional!()
+  let _: Coat = prop.iuo!.optional!()!
+}

test/Inputs/clang-importer-sdk/usr/include/Foundation.h

@@ -868,6 +868,7 @@ typedef struct NonNilableReferences {
 @protocol NSProtocolWithOptionalRequirement
 @optional
 -(void)optionalRequirement;
+-(DummyClass *)optionalRequirementMethodWithIUOResult;
 @end
 
 @interface NSClassWithMethodFromNSProtocolWithOptionalRequirement
@@ -1144,3 +1145,15 @@ void install_global_event_handler(_Nullable event_handler handler);
 
 __nullable id returnNullableId(void);
 void takeNullableId(__nullable id);
+
+@interface I
+@end
+
+@protocol OptionalMethods
+@optional
+- (Coat *)optional;
+@end
+
+@interface IUOProperty
+@property (readonly) id<OptionalMethods> iuo;
+@end