[Migrator] Fix optional types not being updated correctly in some cases

If the replacement type was a function type or protocol-constrained type (e.g.  'Class & Proto) and the existing type had a postfix ? or !, we weren't wrapping the replacement in parens. This would result in an incorrect/ambiguous new type. This patch wraps the replacement in parens if it contains an & or ->, the existing type has a trailing ? or !.

Resolves rdar://problem/32082269.

Author: Nathan Hawes

lib/Migrator/APIDiffMigratorPass.cpp

@@ -37,19 +37,22 @@ namespace {
 
 struct FoundResult {
   SourceRange TokenRange;
-  bool Optional; // Range has a trailing ? or ! included
+  bool Optional; // Range includes a trailing ? or !, e.g. [SomeType!]
   bool Suffixable; // No need to wrap parens when adding optionality
+  bool Suffixed; // Range is followed by a trailing ? or !, e.g. [SomeType]!
   bool isValid() const { return TokenRange.isValid(); }
 };
 
 class ChildIndexFinder : public TypeReprVisitor<ChildIndexFinder, FoundResult> {
   ArrayRef<uint8_t> ChildIndices;
+  bool ParentIsOptional;
 
 public:
   ChildIndexFinder(ArrayRef<uint8_t> ChildIndices) :
     ChildIndices(ChildIndices) {}
 
   FoundResult findChild(AbstractFunctionDecl *Parent) {
+    ParentIsOptional = false;
     auto NextIndex = consumeNext();
     if (!NextIndex) {
       if (auto Func = dyn_cast<FuncDecl>(Parent))
@@ -60,9 +63,9 @@ class ChildIndexFinder : public TypeReprVisitor<ChildIndexFinder, FoundResult> {
         if (!Optional)
           End = Init->getNameLoc();
         return {SourceRange(Init->getNameLoc(), End), Optional,
-          /*suffixable=*/true};
+          /*suffixable=*/true, /*suffixed=*/false};
       }
-      return {SourceRange(), false, false};
+      return {SourceRange(), false, false, false};
     }
 
     for (auto *Params: Parent->getParameterLists()) {
@@ -100,7 +103,7 @@ class ChildIndexFinder : public TypeReprVisitor<ChildIndexFinder, FoundResult> {
 
   FoundResult findChild(TypeLoc Loc) {
     if (!Loc.hasLocation())
-      return {SourceRange(), false, false};
+      return {SourceRange(), false, false, false};
     return visit(Loc.getTypeRepr());
   }
 
@@ -110,12 +113,16 @@ class ChildIndexFinder : public TypeReprVisitor<ChildIndexFinder, FoundResult> {
   FoundResult handleParent(TypeRepr *Parent, const ArrayRef<T> Children,
                            bool Optional = false, bool Suffixable = true) {
     if (!hasNextIndex())
-      return {Parent->getSourceRange(), Optional, Suffixable};
+      return {
+        Parent->getSourceRange(),
+        Optional, Suffixable, /*Suffixed=*/ParentIsOptional
+      };
     auto NextIndex = consumeNext();
     if (isUserTypeAlias(Parent))
-      return {SourceRange(), false, false};
+      return {SourceRange(), false, false, false};
     assert(NextIndex < Children.size());
     TypeRepr *Child = Children[NextIndex];
+    ParentIsOptional = Optional;
     return visit(Child);
   }
 
@@ -137,7 +144,7 @@ class ChildIndexFinder : public TypeReprVisitor<ChildIndexFinder, FoundResult> {
   }
 
   FoundResult visitErrorTypeRepr(ErrorTypeRepr *T) {
-    return {SourceRange(), false, false};
+    return {SourceRange(), false, false, false};
   }
 
   FoundResult visitAttributedTypeRepr(AttributedTypeRepr *T) {
@@ -159,8 +166,10 @@ class ChildIndexFinder : public TypeReprVisitor<ChildIndexFinder, FoundResult> {
   FoundResult visitTupleTypeRepr(TupleTypeRepr *T) {
     // Single element TupleTypeReprs may be arbitrarily nested so don't count
     // as their own index level
-    if (T->getNumElements() == 1)
+    if (T->getNumElements() == 1) {
+      ParentIsOptional = false;
       return visit(T->getElement(0));
+    }
     return handleParent(T, T->getElements());
   }
 
@@ -425,6 +434,10 @@ struct APIDiffMigratorPass : public ASTMigratorPass, public SourceEntityWalker {
     }
   }
 
+  bool typeReplacementMayNeedParens(StringRef Replacement) const {
+    return Replacement.contains('&') || Replacement.contains("->");
+  }
+
   bool walkToDeclPre(Decl *D, CharSourceRange Range) override {
     if (D->isImplicit())
       return true;
@@ -465,6 +478,10 @@ struct APIDiffMigratorPass : public ASTMigratorPass, public SourceEntityWalker {
             break;
           case ide::api::NodeAnnotation::TypeRewritten:
             Editor.replace(Result.TokenRange, DiffItem->RightComment);
+            if (Result.Suffixed && typeReplacementMayNeedParens(DiffItem->RightComment)) {
+              Editor.insertBefore(Result.TokenRange.Start, "(");
+              Editor.insertAfterToken(Result.TokenRange.End, ")");
+            }
             break;
           default:
             break;

test/Migrator/API.json

@@ -207,12 +207,12 @@
   {
     "DiffItemKind": "CommonDiffItem",
     "NodeKind": "Function",
-    "NodeAnnotation": "UnwrapOptional",
-    "ChildIndex": "0",
+    "NodeAnnotation": "TypeRewritten",
+    "ChildIndex": "0:0",
     "LeftUsr": "s:6CitiesAAC7buderimABSgyF",
     "LeftComment": "",
     "RightUsr": "",
-    "RightComment": "",
+    "RightComment": "Cities & ExtraCities",
     "ModuleName": "Cities"
   },
   {
@@ -390,5 +390,27 @@
     "RightUsr": "",
     "RightComment": "setZooLocationNew(newX:newY:newZ:)",
     "ModuleName": "Cities"
+  },
+  {
+    "DiffItemKind": "CommonDiffItem",
+    "NodeKind": "Function",
+    "NodeAnnotation": "TypeRewritten",
+    "ChildIndex": "1:0",
+    "LeftUsr": "s:6CitiesAAC8maroochyySiSg1x_AD1ytF",
+    "LeftComment": "Int",
+    "RightUsr": "",
+    "RightComment": "(Int) -> Int",
+    "ModuleName": "Cities"
+  },
+  {
+    "DiffItemKind": "CommonDiffItem",
+    "NodeKind": "Function",
+    "NodeAnnotation": "TypeRewritten",
+    "ChildIndex": "2:0",
+    "LeftUsr": "s:6CitiesAAC8maroochyySiSg1x_AD1ytF",
+    "LeftComment": "Int",
+    "RightUsr": "",
+    "RightComment": "(Int) -> Int",
+    "ModuleName": "Cities"
   }
 ]

test/Migrator/Inputs/cities.swift

@@ -8,6 +8,7 @@ open class Cities {
   open func yandina(x: [[String : Cities]]!) {}
   open func buderim() -> Cities? { return Cities(x: 1) }
   open func noosa() -> [[String : Cities]?] { return [] }
+  open func maroochy(x: Int?, y: Int?) {}
 }
 
 public protocol ExtraCities {

test/Migrator/wrap_optional.swift

@@ -12,13 +12,14 @@ class MyCities : Cities {
   override func mooloolaba(x: Cities, y: Cities?) {}
   override func toowoomba(x: [Cities], y: [Cities]?) {}
   override func mareeba(x: [String : Cities?], y: [String : Cities]?) {}
+  override func maroochy(x: (Int)?, y: Int?) {}
 }
 
 class MySubCities : MyCities {}
 
 class MySubSubCities : MySubCities {
   override func yandina(x: [[String : Cities]]!) {}
-  override func buderim() -> Cities? { return Cities(x: 1) }
+  override func buderim() -> Cities? { return nil }
   override func noosa() -> [[String : Cities]?] { return [] }
 }
 

test/Migrator/wrap_optional.swift.expected

@@ -12,13 +12,14 @@ class MyCities : Cities {
   override func newMooloolaba(newX x: Cities?, newY y: Cities) {}
   override func toowoomba(x: [Cities?], y: [Cities?]) {}
   override func mareeba(x: [String? : Cities], y: [Int : Cities]) {}
+  override func maroochy(x: ((Int) -> Int)?, y: ((Int) -> Int)?) {}
 }
 
 class MySubCities : MyCities {}
 
 class MySubSubCities : MySubCities {
   override func yandina(x: [[Int : Cities]]?) {}
-  override func buderim() -> Cities { return Cities(x: 1) }
+  override func buderim() -> (Cities & ExtraCities)? { return nil }
   override func noosa() -> [[Int : Cities]] { return [] }
 }