[migrator] Fix missed dollar arg migration in closures relying on implicit destructuring of tuple type arg

lib/Migrator/TupleSplatMigratorPass.cpp

@@ -76,53 +76,94 @@ class ShorthandFinder: public ASTWalker {
 
 struct TupleSplatMigratorPass : public ASTMigratorPass,
   public SourceEntityWalker {
-
-  bool handleClosureShorthandMismatch(FunctionConversionExpr *FC) {
-    if (!SF->getASTContext().LangOpts.isSwiftVersion3() || !FC->isImplicit() ||
-        !isa<ClosureExpr>(FC->getSubExpr())) {
-      return false;
+
+  llvm::DenseSet<FunctionConversionExpr*> CallArgFuncConversions;
+
+  void blacklistFuncConversionArgs(CallExpr *CE) {
+    if (CE->isImplicit() || !SF->getASTContext().LangOpts.isSwiftVersion3())
+      return;
+
+    Expr *Arg = CE->getArg();
+    if (auto *Shuffle = dyn_cast<TupleShuffleExpr>(Arg))
+      Arg = Shuffle->getSubExpr();
+
+    if (auto *Paren = dyn_cast<ParenExpr>(Arg)) {
+      if (auto FC = dyn_cast_or_null<FunctionConversionExpr>(Paren->getSubExpr()))
+        CallArgFuncConversions.insert(FC);
+    } else if (auto *Tuple = dyn_cast<TupleExpr>(Arg)){
+      for (auto Elem : Tuple->getElements()) {
+        if (auto *FC = dyn_cast_or_null<FunctionConversionExpr>(Elem))
+          CallArgFuncConversions.insert(FC);
+      }
+    }
+  }
+
+  ClosureExpr *getShorthandClosure(Expr *E) {
+    if (auto *Closure = dyn_cast_or_null<ClosureExpr>(E)) {
+      if (Closure->hasAnonymousClosureVars())
+        return Closure;
     }
+    return nullptr;
+  }
+
+  bool handleClosureShorthandMismatch(const FunctionConversionExpr *FC) {
+    if (!SF->getASTContext().LangOpts.isSwiftVersion3() || !FC->isImplicit())
+      return false;
 
-    auto *Closure = cast<ClosureExpr>(FC->getSubExpr());
-    if (Closure->getInLoc().isValid())
+    ClosureExpr *Closure = getShorthandClosure(FC->getSubExpr());
+    if (!Closure)
       return false;
 
     FunctionType *FuncTy = FC->getType()->getAs<FunctionType>();
 
     unsigned NativeArity = FuncTy->getParams().size();
     unsigned ClosureArity = Closure->getParameters()->size();
-    if (NativeArity <= ClosureArity)
+    if (NativeArity == ClosureArity)
       return false;
 
-    ShorthandFinder Finder(Closure);
-
     if (ClosureArity == 1 && NativeArity > 1) {
       // Remove $0. from existing references or if it's only $0, replace it
       // with a tuple of the native arity, e.g. ($0, $1, $2)
-      Finder.forEachReference([this, NativeArity](Expr *Ref, ParamDecl *Def) {
-        if (auto *TE = dyn_cast<TupleElementExpr>(Ref)) {
-          SourceLoc Start = TE->getStartLoc();
-          SourceLoc End = TE->getLoc();
-          Editor.replace(CharSourceRange(SM, Start, End), "$");
-        } else {
-          std::string TupleText;
-          {
-            llvm::raw_string_ostream OS(TupleText);
-            for (size_t i = 1; i < NativeArity; ++i) {
-              OS << ", $" << i;
+      ShorthandFinder(Closure)
+        .forEachReference([this, NativeArity](Expr *Ref, ParamDecl *Def) {
+          if (auto *TE = dyn_cast<TupleElementExpr>(Ref)) {
+            SourceLoc Start = TE->getStartLoc();
+            SourceLoc End = TE->getLoc();
+            Editor.replace(CharSourceRange(SM, Start, End), "$");
+          } else {
+            std::string TupleText;
+            {
+              llvm::raw_string_ostream OS(TupleText);
+              for (size_t i = 1; i < NativeArity; ++i) {
+                OS << ", $" << i;
+              }
+              OS << ")";
             }
-            OS << ")";
+            Editor.insert(Ref->getStartLoc(), "(");
+            Editor.insertAfterToken(Ref->getEndLoc(), TupleText);
           }
-          Editor.insert(Ref->getStartLoc(), "(");
-          Editor.insertAfterToken(Ref->getEndLoc(), TupleText);
-        }
-      });
+        });
+      return true;
+    }
+
+    // This direction is only needed if not passed as a call argument. e.g.
+    // someFunc({ $0 > $1 }) // doesn't need migration
+    // let x: ((Int, Int)) -> Bool = { $0 > $1 } // needs migration
+    if (NativeArity == 1 && ClosureArity > 1 && !CallArgFuncConversions.count(FC)) {
+      // Prepend $0. to existing references
+      ShorthandFinder(Closure)
+        .forEachReference([this](Expr *Ref, ParamDecl *Def) {
+          if (auto *TE = dyn_cast<TupleElementExpr>(Ref))
+            Ref = TE->getBase();
+          SourceLoc AfterDollar = Ref->getStartLoc().getAdvancedLoc(1);
+          Editor.insert(AfterDollar, "0.");
+        });
       return true;
     }
     return false;
   }
 
-      /// Migrates code that compiles fine in Swift 3 but breaks in Swift 4 due to
+  /// Migrates code that compiles fine in Swift 3 but breaks in Swift 4 due to
   /// changes in how the typechecker handles tuple arguments.
   void handleTupleArgumentMismatches(const CallExpr *E) {
     if (!SF->getASTContext().LangOpts.isSwiftVersion3())
@@ -169,6 +210,7 @@ struct TupleSplatMigratorPass : public ASTMigratorPass,
     if (auto *FCE = dyn_cast<FunctionConversionExpr>(E)) {
       handleClosureShorthandMismatch(FCE);
     } else if (auto *CE = dyn_cast<CallExpr>(E)) {
+      blacklistFuncConversionArgs(CE);
       handleTupleArgumentMismatches(CE);
     }
     return true;

test/Migrator/tuple-arguments.swift

@@ -59,3 +59,7 @@ extension Dictionary {
 
 let dictionary: [String: String] = [:]
 _ = dictionary.first { (column, value) in true }!.value
+
+func doit(_ x: Int) -> Bool { return x > 0 }
+let _: ((String, Int)) -> [String:Bool] = { [$0: doit($1)] }
+func returnClosure() -> ((Int, Int)) -> Bool { return {$1 > $0} }

test/Migrator/tuple-arguments.swift.expected

@@ -59,3 +59,7 @@ extension Dictionary {
 
 let dictionary: [String: String] = [:]
 _ = dictionary.first { (column, value) in true }!.value
+
+func doit(_ x: Int) -> Bool { return x > 0 }
+let _: ((String, Int)) -> [String:Bool] = { [$0.0: doit($0.1)] }
+func returnClosure() -> ((Int, Int)) -> Bool { return {$0.1 > $0.0} }