Merge pull request swiftlang#24317 from benlangmuir/expand-stmt

[placeholder-expansion] Expand trailing closure in statements correctly

Author: shahmishal

test/SourceKit/CodeExpand/code-expand.swift

@@ -53,9 +53,9 @@ func f() {
 func f1() {
   bar(<#T##d: () -> ()##() -> ()#>)
 }
-// CHECK:   bar({
-// CHECK-NEXT:	<#code#>
-// CHECK-NEXT:	})
+// CHECK:      bar {
+// CHECK-NEXT: <#code#>
+// CHECK-NEXT: }
 
 func f1() {
   bar(<#T##d: () -> ()##() -> ()#>, <#T##d: () -> ()##() -> ()#>)
@@ -103,3 +103,95 @@ braced2(x: {<#T##() -> Void#>}, y: Int)
 // CHECK:   braced2(x: {
 // CHECK-NEXT:  <#code#>
 // CHECK-NEXT:  }, y: Int)
+
+func returnTrailing() -> Int {
+  return withtrail(<#T##() -> ()#>)
+// CHECK: return withtrail {
+// CHECK-NEXT: <#code#>
+}
+
+var yieldTrailing: Int {
+  _read {
+    yield withtrail(<#T##() -> ()#>)
+  // CHECK: yield withtrail {
+  // CHECK-NEXT: <#code#>
+  }
+}
+
+func caseTrailing() -> Int {
+  switch true {
+    case true: withtrail(<#T##() -> ()#>)
+// CHECK: case true: withtrail {
+// CHECK-NEXT: <#code#>
+    default: withtrail(<#T##() -> ()#>)
+// CHECK: default: withtrail {
+// CHECK-NEXT: <#code#>
+  }
+}
+
+func throwTrailing() -> Int {
+   throw withtrail(<#T##() -> ()#>)
+// CHECK: throw withtrail {
+// CHECK-NEXT: <#code#>
+}
+
+func singleExprTrailing1() -> Int {
+  withtrail(<#T##() -> ()#>)
+// CHECK: withtrail {
+// CHECK-NEXT: <#code#>
+}
+var singleExprTrailing2: Int {
+  withtrail(<#T##() -> ()#>)
+// CHECK: withtrail {
+// CHECK-NEXT: <#code#>
+}
+var singleExprTrailing3: Int {
+  get {
+    withtrail(<#T##() -> ()#>)
+// CHECK: withtrail {
+// CHECK-NEXT: <#code#>
+  }
+}
+
+closureTrailingMulti {
+  bah()
+  withtrail(<#T##() -> ()#>)
+// CHECK: bah()
+// CHECK-NEXT: withtrail {
+// CHECK-NEXT: <#code#>
+}
+
+closureIf {
+  if withtrail(<#T##() -> ()#>) {}
+// CHECK: if withtrail({
+// CHECK-NEXT: <#code#>
+}
+
+closureNonTrail {
+  nonTrail(<#T##() -> ()#>, 1)
+// CHECK: nonTrail({
+// CHECK-NEXT: <#code#>
+}
+
+singleExprClosureTrailing {
+  withtrail(<#T##() -> ()#>)
+// CHECK: withtrail {
+// CHECK-NEXT: <#code#>
+}
+
+singleExprClosureTrailingParens({
+  withtrail(<#T##() -> ()#>)
+// CHECK: withtrail {
+// CHECK-NEXT: <#code#>
+})
+
+singleExprClosureMultiArg(1) {
+  withtrail(<#T##() -> ()#>)
+// CHECK: withtrail {
+// CHECK-NEXT: <#code#>
+}
+singleExprClosureMultiArg(1) {
+  withtrail(<#T##() -> ()#>)
+// CHECK: withtrail {
+// CHECK-NEXT: <#code#>
+}

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -1516,8 +1516,26 @@ class PlaceholderExpansionScanner {
       bool walkToExprPre(Expr *E) override {
         auto SR = E->getSourceRange();
         if (SR.isValid() && SM.rangeContainsTokenLoc(SR, TargetLoc)) {
-          if (!checkCallExpr(E) && !EnclosingCallAndArg.first)
+          if (auto closure = dyn_cast<ClosureExpr>(E)) {
+            if (closure->hasSingleExpressionBody()) {
+              // Treat a single-expression body like a brace statement and reset
+              // the enclosing context. Note: when the placeholder is the whole
+              // body it is handled specially as wrapped in braces by
+              // shouldUseTrailingClosureInTuple().
+              auto SR = closure->getSingleExpressionBody()->getSourceRange();
+              if (SR.isValid() && SR.Start != TargetLoc &&
+                  SM.rangeContainsTokenLoc(SR, TargetLoc)) {
+                OuterStmt = nullptr;
+                OuterExpr = nullptr;
+                EnclosingCallAndArg = {nullptr, nullptr};
+                return true;
+              }
+            }
+          }
+
+          if (!checkCallExpr(E) && !EnclosingCallAndArg.first) {
             OuterExpr = E;
+          }
         }
         return true;
       }
@@ -1531,12 +1549,22 @@ class PlaceholderExpansionScanner {
       bool walkToStmtPre(Stmt *S) override {
         auto SR = S->getSourceRange();
         if (SR.isValid() && SM.rangeContainsTokenLoc(SR, TargetLoc)) {
-          if (!EnclosingCallAndArg.first) {
-            if (isa<BraceStmt>(S))
-              // In case OuterStmt is already set, we should clear it to nullptr.
-              OuterStmt = nullptr;
-            else
-              OuterStmt = S;
+          // A statement inside an expression - e.g. `foo({ if ... })` - resets
+          // the enclosing context.
+          OuterExpr = nullptr;
+          EnclosingCallAndArg = {nullptr, nullptr};
+
+          switch (S->getKind()) {
+          case StmtKind::Brace:
+          case StmtKind::Return:
+          case StmtKind::Yield:
+          case StmtKind::Throw:
+            // A trailing closure is allowed in these statements.
+            OuterStmt = nullptr;
+            break;
+          default:
+            OuterStmt = S;
+            break;
           }
         }
         return true;