Skip to content

Commit ddd154b

Browse files
jckarterjckarter
authored
Merge pull request #61213 from jckarter/key-path-literal-as-function-ast-optimization
Sema: Simplify AST representation of key path literals as functions.
2 parents ef6bf93 + 5755387 commit ddd154b

File tree

6 files changed

+111
-67
lines changed

6 files changed

+111
-67
lines changed

‎include/swift/AST/Expr.h

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4047,14 +4047,6 @@ class ClosureExpr : public AbstractClosureExpr {
40474047
/// func f(x : @autoclosure () -> Int)
40484048
/// f(42) // AutoclosureExpr convert from Int to ()->Int
40494049
/// \endcode
4050-
///
4051-
/// They are also created when key path expressions are converted to function
4052-
/// type, in which case, a pair of nested implicit closures are formed:
4053-
/// \code
4054-
/// { $kp$ in { $0[keyPath: $kp$] } }( \(E) )
4055-
/// \endcode
4056-
/// This is to ensure side effects of the key path expression (mainly indices in
4057-
/// subscripts) are only evaluated once.
40584050
class AutoClosureExpr : public AbstractClosureExpr {
40594051
BraceStmt *Body;
40604052

@@ -5397,6 +5389,10 @@ class KeyPathExpr : public Expr {
53975389
/// a contextual root type.
53985390
bool HasLeadingDot = false;
53995391

5392+
/// When we parse a key path literal, we claim a closure discriminator for it, since it may be used as
5393+
/// a closure value in function type context.
5394+
unsigned ClosureDiscriminator;
5395+
54005396
public:
54015397
/// A single stored component, which will be one of:
54025398
/// - an unresolved DeclNameRef, which has to be type-checked
@@ -5728,11 +5724,12 @@ class KeyPathExpr : public Expr {
57285724

57295725
KeyPathExpr(SourceLoc startLoc, Expr *parsedRoot, Expr *parsedPath,
57305726
SourceLoc endLoc, bool hasLeadingDot, bool isObjC,
5731-
bool isImplicit);
5727+
bool isImplicit, unsigned closureDiscriminator);
57325728

57335729
/// Create a key path with unresolved root and path expressions.
57345730
KeyPathExpr(SourceLoc backslashLoc, Expr *parsedRoot, Expr *parsedPath,
5735-
bool hasLeadingDot, bool isImplicit);
5731+
bool hasLeadingDot, bool isImplicit,
5732+
unsigned closureDiscriminator);
57365733

57375734
/// Create a key path with components.
57385735
KeyPathExpr(ASTContext &ctx, SourceLoc startLoc,
@@ -5742,8 +5739,9 @@ class KeyPathExpr : public Expr {
57425739
public:
57435740
/// Create a new parsed Swift key path expression.
57445741
static KeyPathExpr *createParsed(ASTContext &ctx, SourceLoc backslashLoc,
5745-
Expr *parsedRoot, Expr *parsedPath,
5746-
bool hasLeadingDot);
5742+
Expr *parsedRoot, Expr *parsedPath,
5743+
bool hasLeadingDot,
5744+
unsigned closureDiscriminator = AbstractClosureExpr::InvalidDiscriminator);
57475745

57485746
/// Create a new parsed #keyPath expression.
57495747
static KeyPathExpr *createParsedPoundKeyPath(ASTContext &ctx,
@@ -5838,6 +5836,9 @@ class KeyPathExpr : public Expr {
58385836
/// True if this key path expression has a leading dot.
58395837
bool expectsContextualRoot() const { return HasLeadingDot; }
58405838

5839+
/// Return the discriminator to use if this key path becomes a closure.
5840+
unsigned getClosureDiscriminator() const { return ClosureDiscriminator; }
5841+
58415842
static bool classof(const Expr *E) {
58425843
return E->getKind() == ExprKind::KeyPath;
58435844
}

‎lib/AST/Expr.cpp

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2252,21 +2252,24 @@ OpenedArchetypeType *OpenExistentialExpr::getOpenedArchetype() const {
22522252

22532253
KeyPathExpr::KeyPathExpr(SourceLoc startLoc, Expr *parsedRoot,
22542254
Expr *parsedPath, SourceLoc endLoc, bool hasLeadingDot,
2255-
bool isObjC, bool isImplicit)
2255+
bool isObjC, bool isImplicit,
2256+
unsigned closureDiscriminator)
22562257
: Expr(ExprKind::KeyPath, isImplicit), StartLoc(startLoc), EndLoc(endLoc),
22572258
ParsedRoot(parsedRoot), ParsedPath(parsedPath),
2258-
HasLeadingDot(hasLeadingDot) {
2259+
HasLeadingDot(hasLeadingDot), ClosureDiscriminator(closureDiscriminator) {
22592260
assert(!(isObjC && (parsedRoot || parsedPath)) &&
22602261
"Obj-C key paths should only have components");
22612262
Bits.KeyPathExpr.IsObjC = isObjC;
22622263
}
22632264

22642265
KeyPathExpr::KeyPathExpr(SourceLoc backslashLoc, Expr *parsedRoot,
2265-
Expr *parsedPath, bool hasLeadingDot, bool isImplicit)
2266+
Expr *parsedPath, bool hasLeadingDot, bool isImplicit,
2267+
unsigned closureDiscriminator)
22662268
: KeyPathExpr(backslashLoc, parsedRoot, parsedPath,
22672269
parsedPath ? parsedPath->getEndLoc()
22682270
: parsedRoot->getEndLoc(),
2269-
hasLeadingDot, /*isObjC*/ false, isImplicit) {
2271+
hasLeadingDot, /*isObjC*/ false, isImplicit,
2272+
closureDiscriminator) {
22702273
assert((parsedRoot || parsedPath) &&
22712274
"Key path must have either root or path");
22722275
}
@@ -2275,7 +2278,8 @@ KeyPathExpr::KeyPathExpr(ASTContext &ctx, SourceLoc startLoc,
22752278
ArrayRef<Component> components, SourceLoc endLoc,
22762279
bool isObjC, bool isImplicit)
22772280
: KeyPathExpr(startLoc, /*parsedRoot*/ nullptr, /*parsedPath*/ nullptr,
2278-
endLoc, /*hasLeadingDot*/ false, isObjC, isImplicit) {
2281+
endLoc, /*hasLeadingDot*/ false, isObjC, isImplicit,
2282+
/*closure discriminator*/ AbstractClosureExpr::InvalidDiscriminator) {
22792283
assert(!components.empty());
22802284
Components = ctx.AllocateCopy(components);
22812285
}
@@ -2289,9 +2293,11 @@ KeyPathExpr *KeyPathExpr::createParsedPoundKeyPath(
22892293

22902294
KeyPathExpr *KeyPathExpr::createParsed(ASTContext &ctx, SourceLoc backslashLoc,
22912295
Expr *parsedRoot, Expr *parsedPath,
2292-
bool hasLeadingDot) {
2296+
bool hasLeadingDot,
2297+
unsigned closureDiscriminator) {
22932298
return new (ctx) KeyPathExpr(backslashLoc, parsedRoot, parsedPath,
2294-
hasLeadingDot, /*isImplicit*/ false);
2299+
hasLeadingDot, /*isImplicit*/ false,
2300+
closureDiscriminator);
22952301
}
22962302

22972303
KeyPathExpr *KeyPathExpr::createImplicit(ASTContext &ctx,
@@ -2307,7 +2313,8 @@ KeyPathExpr *KeyPathExpr::createImplicit(ASTContext &ctx,
23072313
Expr *parsedRoot, Expr *parsedPath,
23082314
bool hasLeadingDot) {
23092315
return new (ctx) KeyPathExpr(backslashLoc, parsedRoot, parsedPath,
2310-
hasLeadingDot, /*isImplicit*/ true);
2316+
hasLeadingDot, /*isImplicit*/ true,
2317+
/*closureDiscriminator*/ AbstractClosureExpr::InvalidDiscriminator);
23112318
}
23122319

23132320
void

‎lib/Parse/ParseExpr.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -690,7 +690,8 @@ ParserResult<Expr> Parser::parseExprKeyPath() {
690690

691691
auto *keypath = KeyPathExpr::createParsed(
692692
Context, backslashLoc, rootResult.getPtrOrNull(),
693-
pathResult.getPtrOrNull(), hasLeadingDot);
693+
pathResult.getPtrOrNull(), hasLeadingDot,
694+
CurLocalContext->claimNextClosureDiscriminator());
694695
return makeParserResult(parseStatus, keypath);
695696
}
696697

‎lib/SILGen/SILGenExpr.cpp

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1748,6 +1748,11 @@ static bool canPeepholeLiteralClosureConversion(Type literalType,
17481748

17491749
if (!literalFnType || !convertedFnType)
17501750
return false;
1751+
1752+
// Is it an identity conversion?
1753+
if (literalFnType->isEqual(convertedFnType)) {
1754+
return true;
1755+
}
17511756

17521757
// Does the conversion only add `throws`?
17531758
if (literalFnType->isEqual(convertedFnType->getWithoutThrowing())) {

‎lib/SILOptimizer/IPO/CapturePropagation.cpp

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -329,6 +329,22 @@ void CapturePropagation::rewritePartialApply(PartialApplyInst *OrigPAI,
329329
auto *T2TF = Builder.createThinToThickFunction(OrigPAI->getLoc(), FuncRef,
330330
OrigPAI->getType());
331331
OrigPAI->replaceAllUsesWith(T2TF);
332+
333+
// Bypass any mark_dependence on the captures we specialized away.
334+
//
335+
// TODO: If we start to specialize away key path literals with operands
336+
// (subscripts etc.), then a dependence of the new partial_apply on those
337+
// operands may still exist. However, we should still leave the key path
338+
// itself out of the dependency chain, and introduce dependencies on those
339+
// operands instead, so that the key path object itself can be made dead.
340+
for (auto user : T2TF->getUsersOfType<MarkDependenceInst>()) {
341+
if (auto depUser = user->getBase()->getSingleUserOfType<PartialApplyInst>()){
342+
if (depUser == OrigPAI) {
343+
user->replaceAllUsesWith(T2TF);
344+
}
345+
}
346+
}
347+
332348
// Remove any dealloc_stack users.
333349
SmallVector<Operand*, 16> Uses(T2TF->getUses());
334350
for (auto *Use : Uses)

‎lib/Sema/CSApply.cpp

Lines changed: 60 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -4971,53 +4971,72 @@ namespace {
49714971
cs.cacheType(E);
49724972

49734973
// To ensure side effects of the key path expression (mainly indices in
4974-
// subscripts) are only evaluated once, we construct an outer closure,
4975-
// which is immediately evaluated, and an inner closure, which it returns.
4976-
// The result looks like this:
4974+
// subscripts) are only evaluated once, we use a capture list to evaluate
4975+
// the key path immediately and capture it in the function value created.
4976+
// The result looks like:
49774977
//
4978-
// return "{ $kp$ in { $0[keyPath: $kp$] } }( \(E) )"
4978+
// return "{ [$kp$ = \(E)] in $0[keyPath: $kp$] }"
49794979

49804980
auto &ctx = cs.getASTContext();
4981-
auto discriminator = AutoClosureExpr::InvalidDiscriminator;
4981+
auto discriminator = E->getClosureDiscriminator();
49824982

4983-
// The inner closure.
4984-
//
4985-
// let closure = "{ $0[keyPath: $kp$] }"
4986-
//
4987-
// FIXME: Verify ExtInfo state is correct, not working by accident.
4983+
auto contextInfo = exprType->castTo<AnyFunctionType>()->getExtInfo();
49884984
FunctionType::ExtInfo closureInfo;
4985+
closureInfo = closureInfo.withNoEscape(contextInfo.isNoEscape());
49894986
auto closureTy =
49904987
FunctionType::get({FunctionType::Param(baseTy)}, leafTy, closureInfo);
4991-
auto closure = new (ctx)
4992-
AutoClosureExpr(/*set body later*/nullptr, leafTy,
4993-
discriminator, dc);
4988+
4989+
ClosureExpr *closure = new (ctx)
4990+
ClosureExpr(DeclAttributes(),
4991+
SourceRange(),
4992+
/*captured self*/ nullptr,
4993+
/*params*/ nullptr,
4994+
SourceLoc(),
4995+
SourceLoc(),
4996+
SourceLoc(),
4997+
SourceLoc(),
4998+
/*explicit result type*/ nullptr,
4999+
discriminator,
5000+
dc);
5001+
49945002
auto param = new (ctx) ParamDecl(
49955003
SourceLoc(),
49965004
/*argument label*/ SourceLoc(), Identifier(),
49975005
/*parameter name*/ SourceLoc(), ctx.getIdentifier("$0"), closure);
49985006
param->setInterfaceType(baseTy->mapTypeOutOfContext());
49995007
param->setSpecifier(ParamSpecifier::Default);
50005008
param->setImplicit();
5009+
5010+
auto params = ParameterList::create(ctx, SourceLoc(),
5011+
param, SourceLoc());
50015012

5002-
// The outer closure.
5003-
//
5004-
// let outerClosure = "{ $kp$ in \(closure) }"
5005-
//
5006-
// FIXME: Verify ExtInfo state is correct, not working by accident.
5007-
FunctionType::ExtInfo outerClosureInfo;
5008-
auto outerClosureTy = FunctionType::get({FunctionType::Param(keyPathTy)},
5009-
closureTy, outerClosureInfo);
5010-
auto outerClosure = new (ctx)
5011-
AutoClosureExpr(/*set body later*/nullptr, closureTy,
5012-
discriminator, dc);
5013-
auto outerParam =
5014-
new (ctx) ParamDecl(SourceLoc(),
5015-
/*argument label*/ SourceLoc(), Identifier(),
5016-
/*parameter name*/ SourceLoc(),
5017-
ctx.getIdentifier("$kp$"), outerClosure);
5018-
outerParam->setInterfaceType(keyPathTy->mapTypeOutOfContext());
5019-
outerParam->setSpecifier(ParamSpecifier::Default);
5013+
closure->setParameterList(params);
5014+
5015+
// The capture list.
5016+
VarDecl *outerParam = new (ctx) VarDecl(/*static*/ false,
5017+
VarDecl::Introducer::Let,
5018+
SourceLoc(),
5019+
ctx.getIdentifier("$kp$"),
5020+
dc);
50205021
outerParam->setImplicit();
5022+
outerParam->setInterfaceType(keyPathTy->mapTypeOutOfContext());
5023+
5024+
NamedPattern *outerParamPat = new (ctx) NamedPattern(outerParam);
5025+
outerParamPat->setImplicit();
5026+
outerParamPat->setType(keyPathTy);
5027+
5028+
auto outerParamPBE = PatternBindingEntry(outerParamPat,
5029+
SourceLoc(), E, dc);
5030+
solution.setExprTypes(E);
5031+
auto outerParamDecl = PatternBindingDecl::create(ctx, SourceLoc(),
5032+
{}, SourceLoc(),
5033+
outerParamPBE,
5034+
dc);
5035+
outerParamDecl->setImplicit();
5036+
auto outerParamCapture = CaptureListEntry(outerParamDecl);
5037+
auto captureExpr = CaptureListExpr::create(ctx, outerParamCapture,
5038+
closure);
5039+
captureExpr->setImplicit();
50215040

50225041
// let paramRef = "$0"
50235042
auto *paramRef = new (ctx)
@@ -5037,26 +5056,21 @@ namespace {
50375056
E->getStartLoc(), outerParamRef, E->getEndLoc(),
50385057
leafTy, /*implicit=*/true);
50395058
cs.cacheType(application);
5059+
5060+
auto returnStmt = new (ctx) ReturnStmt(SourceLoc(), application);
5061+
auto bodyStmt = BraceStmt::create(ctx,
5062+
SourceLoc(), ASTNode(returnStmt), SourceLoc());
50405063

50415064
// Finish up the inner closure.
50425065
closure->setParameterList(ParameterList::create(ctx, {param}));
5043-
closure->setBody(application);
5066+
closure->setBody(bodyStmt, /*singleExpr*/ true);
50445067
closure->setType(closureTy);
50455068
cs.cacheType(closure);
5046-
5047-
// Finish up the outer closure.
5048-
outerClosure->setParameterList(ParameterList::create(ctx, {outerParam}));
5049-
outerClosure->setBody(closure);
5050-
outerClosure->setType(outerClosureTy);
5051-
cs.cacheType(outerClosure);
5052-
5053-
// let outerApply = "\( outerClosure )( \(E) )"
5054-
auto *argList = ArgumentList::forImplicitUnlabeled(ctx, {E});
5055-
auto outerApply = CallExpr::createImplicit(ctx, outerClosure, argList);
5056-
outerApply->setType(closureTy);
5057-
cs.cacheExprTypes(outerApply);
5058-
5059-
return coerceToType(outerApply, exprType, cs.getConstraintLocator(E));
5069+
5070+
captureExpr->setType(closureTy);
5071+
cs.cacheType(captureExpr);
5072+
5073+
return coerceToType(captureExpr, exprType, cs.getConstraintLocator(E));
50605074
}
50615075

50625076
void buildKeyPathOptionalForceComponent(

0 commit comments

Comments
 (0)