[placeholder-expansion] Expand trailing closure in statements correctly

In addition to brace statements (previously handled), several other
statements allow a trailing closure (return, yield, throw). Also fix the
handling of statements nested within expressions for closures - both
single-expression bodies and brace statements.

This also fixes a particular regression caused by single-expression
function bodies where we would fail to expand to a trailing closure when
a function body only contained a single expression.

rdar://50227500
This commit is contained in:
Ben Langmuir
2019-04-26 10:28:02 -07:00
parent f88be05394
commit 6377199e2f
2 changed files with 130 additions and 10 deletions

View File

@@ -1516,8 +1516,26 @@ private:
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 @@ private:
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;