[Migrator] Migrate references to shorthand closure params (e.g. $0, $1) where affected by SE110

In Swift 3:
 1) a closure referring to $1 or above type checked against a function type that takes a single tuple argument with that arity
 2) a closure referring only to $0 (or $0.1 etc) type checked against a function type that takes multiple arguments
but neither compiles in Swift 4.

This patch migrates shorthand references for
  1) by prefixing "0." in front of the existing references, e.g. "$1" to "$0.1"
  2) by removing the leading 0. if one existing or substituting a tuple of the correct arity if it doesn't, e.g. "$0.1" to "$1" and "$0" to "($0, $1, $2)"

Resolves rdar://problem/31969538

Author: Nathan Hawes

lib/Migrator/TupleSplatMigratorPass.cpp

@@ -22,9 +22,115 @@ using namespace swift;
 using namespace swift::migrator;
 
 namespace {
+  
+class ShorthandFinder: public ASTWalker {
+private:
+  llvm::DenseMap<ParamDecl*, std::vector<Expr*>> References;
+
+  std::pair<bool, Expr *> walkToExprPre(Expr *E) override {
+    Expr *ParentElementExpr = nullptr;
+    Expr *OrigE = E;
+
+    if (auto *TupleElem = dyn_cast<TupleElementExpr>(E)) {
+      ParentElementExpr = TupleElem;
+      E = TupleElem->getBase()->getSemanticsProvidingExpr();
+    }
+
+    if (auto *DeclRef = dyn_cast<DeclRefExpr>(E)) {
+      ParamDecl *Decl = dyn_cast<ParamDecl>(DeclRef->getDecl());
+      Expr *Reference = ParentElementExpr? ParentElementExpr : DeclRef;
+      if (References.count(Decl) && !Reference->isImplicit()) {
+        References[Decl].push_back(Reference);
+        return { false, OrigE };
+      }
+    }
+    return { true, OrigE };
+  }
+
+public:
+  ShorthandFinder(ClosureExpr *Expr) {
+    if (!Expr->hasAnonymousClosureVars())
+      return;
+    References.clear();
+    for(auto *Param: *Expr->getParameters()) {
+      References[Param] = {};
+    }
+    Expr->walk(*this);
+  }
+
+  void forEachReference(llvm::function_ref<void(Expr*, ParamDecl*)> Callback) {
+    for(auto Entry: References) {
+      for(auto *Expr : Entry.getSecond()) {
+        Callback(Expr, Entry.getFirst());
+      }
+    }
+  }
+};
 
 struct TupleSplatMigratorPass : public ASTMigratorPass,
   public SourceEntityWalker {
+    
+  bool handleClosureShorthandMismatch(FunctionConversionExpr *FC) {
+    if (!SF->getASTContext().LangOpts.isSwiftVersion3() || !FC->isImplicit() ||
+        !isa<ClosureExpr>(FC->getSubExpr())) {
+      return false;
+    }
+
+    auto *Closure = cast<ClosureExpr>(FC->getSubExpr());
+    if (Closure->getInLoc().isValid())
+      return false;
+
+    FunctionType *FuncTy = FC->getType()->getAs<FunctionType>();
+
+    unsigned NativeArity = 0;
+    if (isa<ParenType>(FuncTy->getInput().getPointer())) {
+      NativeArity = 1;
+    } else if (auto TT = FuncTy->getInput()->getAs<TupleType>()) {
+      NativeArity = TT->getNumElements();
+    }
+
+    unsigned ClosureArity = Closure->getParameters()->size();
+    if (NativeArity == ClosureArity)
+      return false;
+
+    ShorthandFinder Finder(Closure);
+    if (NativeArity == 1 && ClosureArity > 1) {
+      // Prepend $0. to existing references
+      Finder.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;
+    }
+
+    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 << ", ";
+              OS << "$" << i;
+            }
+            OS << ")";
+          }
+          Editor.insert(Ref->getStartLoc(), "(");
+          Editor.insertAfterToken(Ref->getEndLoc(), TupleText);
+        }
+      });
+      return true;
+    }
+    return false;
+  }
 
       /// Migrates code that compiles fine in Swift 3 but breaks in Swift 4 due to
   /// changes in how the typechecker handles tuple arguments.
@@ -86,7 +192,7 @@ struct TupleSplatMigratorPass : public ASTMigratorPass,
       auto parenT = dyn_cast<ParenType>(fnTy2->getInput().getPointer());
       if (!parenT)
         return false;
-      auto tupleInFn = dyn_cast<TupleType>(parenT->getUnderlyingType().getPointer());
+      auto tupleInFn = parenT->getAs<TupleType>();
       if (!tupleInFn)
         return false;
       if (!E->getArg())
@@ -197,7 +303,9 @@ struct TupleSplatMigratorPass : public ASTMigratorPass,
   }
 
   bool walkToExprPre(Expr *E) override {
-    if (auto *CE = dyn_cast<CallExpr>(E)) {
+    if (auto *FCE = dyn_cast<FunctionConversionExpr>(E)) {
+      handleClosureShorthandMismatch(FCE);
+    } else if (auto *CE = dyn_cast<CallExpr>(E)) {
       handleTupleArgumentMismatches(CE);
     }
     return true;

test/Migrator/tuple-arguments.swift

@@ -2,8 +2,6 @@
 // RUN: %target-swift-frontend -typecheck -update-code -primary-file %s -emit-migrated-file-path %t.result -disable-migrator-fixits -swift-version 3
 // RUN: diff -u %s.expected %t.result
 // RUN: %target-swift-frontend -typecheck %s.expected -swift-version 4
-// rdar://
-// XFAIL: *
 
 func test1(_: ()) {}
 test1(())
@@ -37,4 +35,15 @@ func toString(indexes: Int?...) -> String {
     if index != nil {}
     return ""
   })
+  let _ = indexes.reduce(0) { print($0); return $0.0 + ($0.1 ?? 0)}
+  let _ = indexes.reduce(0) { (true ? $0 : (1, 2)).0 + ($0.1 ?? 0) }
+  let _ = [(1, 2)].contains { $0 != $1 }
+  _ = ["Hello", "Foo"].sorted { print($0); return $0.0.characters.count > ($0).1.characters.count }
+  _ = ["Hello" : 2].map { ($0, ($1)) }
+}
+
+extension Dictionary {
+  public mutating func merge(with dictionary: Dictionary) {
+    dictionary.forEach { updateValue($1, forKey: $0) }
+  }
 }

test/Migrator/tuple-arguments.swift.expected

@@ -35,4 +35,15 @@ func toString(indexes: Int?...) -> String {
     if index != nil {}
     return ""
   })
+  let _ = indexes.reduce(0) { print(($0, $1)); return $0 + ($1 ?? 0)}
+  let _ = indexes.reduce(0) { (true ? ($0, $1) : (1, 2)).0 + ($1 ?? 0) }
+  let _ = [(1, 2)].contains { $0.0 != $0.1 }
+  _ = ["Hello", "Foo"].sorted { print(($0, $1)); return $0.characters.count > $1.characters.count }
+  _ = ["Hello" : 2].map { ($0.0, ($0.1)) }
+}
+
+extension Dictionary {
+  public mutating func merge(with dictionary: Dictionary) {
+    dictionary.forEach { updateValue($0.1, forKey: $0.0) }
+  }
 }