[SourceKit] Handle macro arguments in placeholder expansion

Allow macro argument placeholders to expand into trailing closures.

rdar://150550747

Author: hamishknight

test/SourceKit/CodeExpand/code-expand.swift

@@ -25,6 +25,10 @@ await foo(x: <#T##() -> Void#>)
 // CHECK-NEXT: <#code#>
 // CHECK-NEXT: }
 
+foo(bar(<#T##() -> Void#>))
+// CHECK:      foo(bar({
+// CHECK-NEXT: <#code#>
+// CHECK-NEXT: }))
 
 anArr.indexOfObjectPassingTest(<#T##predicate: ((AnyObject!, Int, UnsafePointer<ObjCBool>) -> Bool)?##((AnyObject!, Int, UnsafePointer<ObjCBool>) -> Bool)?#>)
 // CHECK:      anArr.indexOfObjectPassingTest { <#AnyObject!#>, <#Int#>, <#UnsafePointer<ObjCBool>#> in
@@ -273,3 +277,41 @@ expandClosureWithInternalParameterNames {
 // CHECK: withtrail { a, b in
 // CHECK-NEXT: <#code#>
 }
+
+// CHECK-LABEL: func expandMacro()
+func expandMacro() {
+  #foo(<#T##() -> Int#>)
+  // CHECK:      #foo {
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }
+
+  #foo(bar: <#T##() -> ()#>)
+  // CHECK:      #foo {
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }
+
+  #foo(bar: <#T##() -> Int#>, baz: <#T##() -> ()#>)
+  // CHECK:      #foo {
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: } baz: {
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: struct ExpandDeclMacro
+struct ExpandDeclMacro {
+  #foo(<#T##() -> ()#>)
+  // CHECK:      #foo {
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }
+
+  #foo(bar(<#T##() -> ()#>))
+  // CHECK:      #foo(bar({
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }))
+
+  #foo(#bar(<#T##() -> ()#>))
+  // CHECK:      #foo(#bar({
+  // CHECK-NEXT:   <#code#>
+  // CHECK-NEXT: }))
+}

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -1679,34 +1679,65 @@ class PlaceholderExpansionScanner {
     public:
       const SourceManager &SM;
       SourceLoc TargetLoc;
-      std::pair<Expr *, ArgumentList *> EnclosingCallAndArg;
+      std::pair<ASTNode, ArgumentList *> EnclosingCallAndArg;
       Expr *OuterExpr;
       Stmt *OuterStmt;
-      explicit CallExprFinder(const SourceManager &SM)
-        :SM(SM) { }
-
-      bool checkCallExpr(Expr *E) {
-        ArgumentList *Args = nullptr;
-        if (auto *CE = dyn_cast<CallExpr>(E)) {
-          // Call expression can have argument.
-          Args = CE->getArgs();
+      Decl *OuterDecl;
+
+      explicit CallExprFinder(const SourceManager &SM) : SM(SM) {}
+
+      void checkArgumentList(ArgumentList *Args) {
+        ASTNode Enclosing;
+        if (auto *E = Parent.getAsExpr()) {
+          // Must have a function call or macro parent.
+          if (!isa<CallExpr, MacroExpansionExpr>(E) ||
+              E->getArgs() != Args) {
+            return;
+          }
+          // If the outer expression is the call itself, disregard it.
+          if (OuterExpr == E)
+            OuterExpr = nullptr;
+
+          Enclosing = E;
         }
-        if (!Args)
-          return false;
-        if (EnclosingCallAndArg.first)
-          OuterExpr = EnclosingCallAndArg.first;
-        EnclosingCallAndArg = {E, Args};
-        return true;
+        if (auto *D = Parent.getAsDecl()) {
+          // Must have a macro decl parent.
+          auto *MED = dyn_cast<MacroExpansionDecl>(D);
+          if (!MED || MED->getArgs() != Args)
+            return;
+          Enclosing = D;
+        }
+        if (!Enclosing)
+          return;
+
+        // If we have an outer enclosing call, update the outer nodes if needed.
+        if (auto EnclosingCall = EnclosingCallAndArg.first) {
+          if (!OuterExpr)
+            OuterExpr = EnclosingCall.dyn_cast<Expr *>();
+          if (!OuterDecl)
+            OuterDecl = EnclosingCall.dyn_cast<Decl *>();
+        }
+
+        EnclosingCallAndArg = {Enclosing, Args};
       }
 
       MacroWalking getMacroWalkingBehavior() const override {
         return MacroWalking::Arguments;
       }
 
+      PreWalkResult<ArgumentList *>
+      walkToArgumentListPre(ArgumentList *ArgList) override {
+        auto SR = ArgList->getSourceRange();
+        if (SR.isValid() && SM.rangeContainsTokenLoc(SR, TargetLoc))
+          checkArgumentList(ArgList);
+
+        return Action::Continue(ArgList);
+      }
+
       PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
         auto SR = E->getSourceRange();
         if (SR.isValid() && SM.rangeContainsTokenLoc(SR, TargetLoc) &&
-            !checkCallExpr(E) && !EnclosingCallAndArg.first) {
+            !OuterExpr && !EnclosingCallAndArg.first) {
           if (!isa<TryExpr>(E) && !isa<AwaitExpr>(E) && !isa<UnsafeExpr>(E) &&
               !isa<PrefixUnaryExpr>(E)) {
             // We don't want to expand to trailing closures if the call is
@@ -1782,6 +1813,7 @@ class PlaceholderExpansionScanner {
         EnclosingCallAndArg = {nullptr, nullptr};
         OuterExpr = nullptr;
         OuterStmt = nullptr;
+        OuterDecl = nullptr;
         TargetLoc = SL;
         SF.walk(*this);
         return EnclosingCallAndArg.second;
@@ -1791,14 +1823,9 @@ class PlaceholderExpansionScanner {
     CallExprFinder CEFinder(SM);
     auto *Args = CEFinder.findEnclosingCallArg(SF, SL);
 
-    if (!Args)
-      return std::make_pair(Args, false);
-    if (CEFinder.OuterExpr)
-      return std::make_pair(Args, false);
-    if (CEFinder.OuterStmt)
-      return std::make_pair(Args, false);
-
-    return std::make_pair(Args, true);
+    auto HasOuterNode =
+        CEFinder.OuterExpr || CEFinder.OuterStmt || CEFinder.OuterDecl;
+    return std::make_pair(Args, /*useTrailingClosure*/ Args && !HasOuterNode);
   }
 
   struct ParamClosureInfo {