migrator: handle function's return type changes from nonnull to nullable. rdar://47265255

Adding exclamation brings an optional type back to its original nonnull
version.

Author: nkcsgexi

lib/Migrator/APIDiffMigratorPass.cpp

@@ -1077,6 +1077,31 @@ struct APIDiffMigratorPass : public ASTMigratorPass, public SourceEntityWalker {
     return false;
   }
 
+  void handleResultTypeChange(ValueDecl *FD, Expr *Call) {
+    Optional<NodeAnnotation> ChangeKind;
+
+    // look for related change item for the function decl.
+    for (auto Item: getRelatedDiffItems(FD)) {
+      if (auto *CI = dyn_cast<CommonDiffItem>(Item)) {
+        // check if the function's return type has been changed from nonnull
+        // to nullable.
+        if (CI->DiffKind == NodeAnnotation::WrapOptional &&
+            CI->getChildIndices().size() == 1 &&
+            CI->getChildIndices().front() == 0) {
+          ChangeKind = NodeAnnotation::WrapOptional;
+          break;
+        }
+      }
+    }
+    if (!ChangeKind.hasValue())
+      return;
+    // If a function's return type has been changed from nonnull to nullable,
+    // append ! to the original call expression.
+    if (*ChangeKind == NodeAnnotation::WrapOptional) {
+      Editor.insertAfterToken(Call->getSourceRange().End, "!");
+    }
+  }
+
   bool walkToExprPre(Expr *E) override {
     if (E->getSourceRange().isInvalid())
       return false;
@@ -1100,6 +1125,7 @@ struct APIDiffMigratorPass : public ASTMigratorPass, public SourceEntityWalker {
           handleTypeHoist(FD, CE, Args);
           handleSpecialCases(FD, CE, Args);
           handleStringRepresentableArg(FD, Args, CE);
+          handleResultTypeChange(FD, CE);
         }
         break;
       }
@@ -1110,15 +1136,16 @@ struct APIDiffMigratorPass : public ASTMigratorPass, public SourceEntityWalker {
           handleFunctionCallToPropertyChange(FD, DSC->getFn(), Args);
           handleSpecialCases(FD, CE, Args);
           handleStringRepresentableArg(FD, Args, CE);
+          handleResultTypeChange(FD, CE);
         }
         break;
       }
       case ExprKind::ConstructorRefCall: {
         auto CCE = cast<ConstructorRefCallExpr>(Fn);
         if (auto FD = CCE->getFn()->getReferencedDecl().getDecl()) {
-          auto *CE = CCE->getFn();
           handleFuncRename(FD, CE, Args);
           handleStringRepresentableArg(FD, Args, CE);
+          handleResultTypeChange(FD, CE);
         }
         break;
       }

test/Migrator/Inputs/CallExpr.json

@@ -0,0 +1,24 @@
+[
+  {
+    "DiffItemKind": "CommonDiffItem",
+    "NodeKind": "Function",
+    "NodeAnnotation": "WrapOptional",
+    "ChildIndex": "0",
+    "LeftUsr": "s:6CitiesAAC1xABSi_tcfc",
+    "LeftComment": "",
+    "RightUsr": "",
+    "RightComment": "",
+    "ModuleName": "Cities"
+  },
+  {
+    "DiffItemKind": "CommonDiffItem",
+    "NodeKind": "Function",
+    "NodeAnnotation": "WrapOptional",
+    "ChildIndex": "0",
+    "LeftUsr": "s:6CitiesAAC5noosaSaySDySSABGSgGyF",
+    "LeftComment": "",
+    "RightUsr": "",
+    "RightComment": "",
+    "ModuleName": "Cities"
+  }
+]

test/Migrator/call_expr_result.swift

@@ -0,0 +1,13 @@
+// REQUIRES: objc_interop
+// RUN: %empty-directory(%t.mod)
+// RUN: %target-swift-frontend -emit-module -o %t.mod/Cities.swiftmodule %S/Inputs/Cities.swift -module-name Cities -parse-as-library
+// RUN: %empty-directory(%t) && %target-swift-frontend -c -update-code -disable-migrator-fixits -primary-file %s  -I %t.mod -api-diff-data-file %S/Inputs/CallExpr.json -emit-migrated-file-path %t/call_expr_result.swift.result -o /dev/null
+// RUN: diff -u %S/call_expr_result.swift.expected %t/call_expr_result.swift.result
+
+import Cities
+
+func foo() {
+  let c1 = Cities(x: 3)
+  _ = Cities.init(x: 3)
+  _ = c1.noosa()
+}

test/Migrator/call_expr_result.swift.expected

@@ -0,0 +1,13 @@
+// REQUIRES: objc_interop
+// RUN: %empty-directory(%t.mod)
+// RUN: %target-swift-frontend -emit-module -o %t.mod/Cities.swiftmodule %S/Inputs/Cities.swift -module-name Cities -parse-as-library
+// RUN: %empty-directory(%t) && %target-swift-frontend -c -update-code -disable-migrator-fixits -primary-file %s  -I %t.mod -api-diff-data-file %S/Inputs/CallExpr.json -emit-migrated-file-path %t/call_expr_result.swift.result -o /dev/null
+// RUN: diff -u %S/call_expr_result.swift.expected %t/call_expr_result.swift.result
+
+import Cities
+
+func foo() {
+  let c1 = Cities(x: 3)!
+  _ = Cities.init(x: 3)!
+  _ = c1.noosa()!
+}

test/Migrator/wrap_optional.swift.expected

@@ -7,7 +7,7 @@
 import Cities
 
 class MyCities : Cities {
-  override init?(x: Int?) { super.init(x: x) }
+  override init?(x: Int?) { super.init(x: x)! }
   override init?(y: Int) { super.init(y: y) }
   override func newMooloolaba(newX x: Cities?, newY y: Cities) {}
   override func toowoomba(x: [Cities?], y: [Cities?]) {}