Merge pull request #12982 from nkcsgexi/code-expand-fix

Author: swift-ci

test/SourceKit/CodeExpand/code-expand.swift

@@ -87,4 +87,7 @@ if true {
 // CHECK:   withtrail {
 // CHECK-NEXT: <#code#>
 }
-}
+}
+
+foo(.foo(<#T##block: () -> Void##() -> Void#>))
+// CHECK: foo(.foo({

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -1537,27 +1537,39 @@ class PlaceholderExpansionScanner {
   /// For example, if the CallExpr is enclosed in another expression or statement
   /// such as "outer(inner(<#closure#>))", or "if inner(<#closure#>)", then trailing
   /// closure should not be applied to the inner call.
-  std::pair<CallExpr *, bool> enclosingCallExpr(SourceFile &SF, SourceLoc SL) {
+  std::pair<Expr*, bool> enclosingCallExprArg(SourceFile &SF, SourceLoc SL) {
 
     class CallExprFinder : public SourceEntityWalker {
     public:
       const SourceManager &SM;
       SourceLoc TargetLoc;
-      CallExpr *EnclosingCall;
+      std::pair<Expr *, Expr*> EnclosingCallAndArg;
       Expr *OuterExpr;
       Stmt *OuterStmt;
       explicit CallExprFinder(const SourceManager &SM)
         :SM(SM) { }
 
+      bool checkCallExpr(Expr *E) {
+        Expr* Arg = nullptr;
+        if (auto *CE = dyn_cast<CallExpr>(E)) {
+          // Call expression can have argument.
+          Arg = CE->getArg();
+        } else if (auto UME = dyn_cast<UnresolvedMemberExpr>(E)) {
+          // Unresolved member can have argument too.
+          Arg = UME->getArgument();
+        }
+        if (!Arg)
+          return false;
+        if (EnclosingCallAndArg.first)
+          OuterExpr = EnclosingCallAndArg.first;
+        EnclosingCallAndArg = {E, Arg};
+        return true;
+      }
+
       bool walkToExprPre(Expr *E) override {
         auto SR = E->getSourceRange();
         if (SR.isValid() && SM.rangeContainsTokenLoc(SR, TargetLoc)) {
-          if (auto *CE = dyn_cast<CallExpr>(E)) {
-            if (EnclosingCall)
-              OuterExpr = EnclosingCall;
-            EnclosingCall = CE;
-          }
-          else if (!EnclosingCall)
+          if (!checkCallExpr(E) && !EnclosingCallAndArg.first)
             OuterExpr = E;
         }
         return true;
@@ -1572,7 +1584,7 @@ class PlaceholderExpansionScanner {
       bool walkToStmtPre(Stmt *S) override {
         auto SR = S->getSourceRange();
         if (SR.isValid() && SM.rangeContainsTokenLoc(SR, TargetLoc)) {
-          if (!EnclosingCall) {
+          if (!EnclosingCallAndArg.first) {
             if (isa<BraceStmt>(S))
               // In case OuterStmt is already set, we should clear it to nullptr.
               OuterStmt = nullptr;
@@ -1583,18 +1595,18 @@ class PlaceholderExpansionScanner {
         return true;
       }
 
-      CallExpr *findEnclosingCall(SourceFile &SF, SourceLoc SL) {
-        EnclosingCall = nullptr;
+      Expr *findEnclosingCallArg(SourceFile &SF, SourceLoc SL) {
+        EnclosingCallAndArg = {nullptr, nullptr};
         OuterExpr = nullptr;
         OuterStmt = nullptr;
         TargetLoc = SL;
         walk(SF);
-        return EnclosingCall;
+        return EnclosingCallAndArg.second;
       }
     };
 
     CallExprFinder CEFinder(SM);
-    auto *CE = CEFinder.findEnclosingCall(SF, SL);
+    auto *CE = CEFinder.findEnclosingCallArg(SF, SL);
 
     if (!CE)
       return std::make_pair(CE, false);
@@ -1647,8 +1659,8 @@ class PlaceholderExpansionScanner {
     // and if the call parens can be removed in that case.
     // We'll first find the enclosing CallExpr, and then do further analysis.
     bool UseTrailingClosure = false;
-    std::pair<CallExpr*, bool> ECE = enclosingCallExpr(SF, PlaceholderStartLoc);
-    Expr *Args = ECE.first ? ECE.first->getArg() : nullptr;
+    auto ECE = enclosingCallExprArg(SF, PlaceholderStartLoc);
+    Expr *Args = ECE.first;
     if (Args && ECE.second) {
       if (isa<ParenExpr>(Args)) {
         UseTrailingClosure = true;