Skip to content

Commit ff21dde

Browse files
atrickatrick
authored
Merge pull request #77675 from atrick/endapply_utils
[NFC] Utilities for lifetime dependent coroutines
2 parents 753cbcd + 442739a commit ff21dde

15 files changed

+164
-48
lines changed

‎SwiftCompilerSources/Sources/SIL/Builder.swift

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,18 @@ public struct Builder {
2929
private let notificationHandler: BridgedChangeNotificationHandler
3030
private let notifyNewInstruction: (Instruction) -> ()
3131

32+
/// Return 'nil' when inserting at the start of a function or in a global initializer.
33+
public var insertionBlock: BasicBlock? {
34+
switch insertAt {
35+
case let .before(inst):
36+
return inst.parentBlock
37+
case let .atEndOf(block):
38+
return block
39+
case .atStartOf, .staticInitializer:
40+
return nil
41+
}
42+
}
43+
3244
public var bridged: BridgedBuilder {
3345
switch insertAt {
3446
case .before(let inst):
@@ -482,6 +494,18 @@ public struct Builder {
482494
return notifyNew(endAccess.getAs(EndAccessInst.self))
483495
}
484496

497+
@discardableResult
498+
public func createEndApply(beginApply: BeginApplyInst) -> EndApplyInst {
499+
let endApply = bridged.createEndApply(beginApply.token.bridged)
500+
return notifyNew(endApply.getAs(EndApplyInst.self))
501+
}
502+
503+
@discardableResult
504+
public func createAbortApply(beginApply: BeginApplyInst) -> AbortApplyInst {
505+
let endApply = bridged.createAbortApply(beginApply.token.bridged)
506+
return notifyNew(endApply.getAs(AbortApplyInst.self))
507+
}
508+
485509
public func createConvertFunction(originalFunction: Value, resultType: Type, withoutActuallyEscaping: Bool) -> ConvertFunctionInst {
486510
let convertFunction = bridged.createConvertFunction(originalFunction.bridged, resultType.bridged, withoutActuallyEscaping)
487511
return notifyNew(convertFunction.getAs(ConvertFunctionInst.self))

‎SwiftCompilerSources/Sources/SIL/Instruction.swift

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1219,6 +1219,8 @@ final public class AllocExistentialBoxInst : SingleValueInstruction, Allocation
12191219
/// `end_borrow`).
12201220
public protocol ScopedInstruction {
12211221
var endOperands: LazyFilterSequence<UseList> { get }
1222+
1223+
var endInstructions: EndInstructions { get }
12221224
}
12231225

12241226
extension Instruction {
@@ -1330,7 +1332,7 @@ final public class BeginApplyInst : MultipleValueInstruction, FullApplySite {
13301332
}
13311333
}
13321334

1333-
final public class EndApplyInst : Instruction, UnaryInstruction {
1335+
final public class EndApplyInst : SingleValueInstruction, UnaryInstruction {
13341336
public var token: MultipleValueInstructionResult { operand.value as! MultipleValueInstructionResult }
13351337
public var beginApply: BeginApplyInst { token.parentInstruction as! BeginApplyInst }
13361338
}
@@ -1342,7 +1344,7 @@ final public class AbortApplyInst : Instruction, UnaryInstruction {
13421344

13431345
extension BeginApplyInst : ScopedInstruction {
13441346
public var endOperands: LazyFilterSequence<UseList> {
1345-
return token.uses.lazy.filter { _ in true }
1347+
return token.uses.lazy.filter { $0.endsLifetime }
13461348
}
13471349
}
13481350

‎SwiftCompilerSources/Sources/SIL/Utilities/AccessUtils.swift

Lines changed: 66 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -103,11 +103,33 @@ public enum AccessBase : CustomStringConvertible, Hashable {
103103
}
104104
case let sb as StoreBorrowInst:
105105
self = .storeBorrow(sb)
106+
case let p2a as PointerToAddressInst:
107+
if let global = p2a.resultOfGlobalAddressorCall {
108+
self = .global(global)
109+
} else {
110+
self = .pointer(p2a)
111+
}
106112
default:
107113
self = .unidentified
108114
}
109115
}
110116

117+
/// Return 'nil' for global varabiables and unidentified addresses.
118+
public var address: Value? {
119+
switch self {
120+
case .global, .unidentified: return nil
121+
case .box(let pbi): return pbi
122+
case .stack(let asi): return asi
123+
case .class(let rea): return rea
124+
case .tail(let rta): return rta
125+
case .argument(let arg): return arg
126+
case .yield(let result): return result
127+
case .storeBorrow(let sb): return sb
128+
case .pointer(let p): return p
129+
case .index(let ia): return ia
130+
}
131+
}
132+
111133
public var description: String {
112134
switch self {
113135
case .unidentified: return "?"
@@ -480,9 +502,40 @@ public enum EnclosingScope {
480502
case base(AccessBase)
481503
}
482504

505+
private struct EnclosingAccessWalker : AddressUseDefWalker {
506+
var enclosingScope: EnclosingScope?
507+
508+
mutating func walk(startAt address: Value, initialPath: UnusedWalkingPath = UnusedWalkingPath()) {
509+
if walkUp(address: address, path: UnusedWalkingPath()) == .abortWalk {
510+
assert(enclosingScope == nil, "shouldn't have set an enclosing scope in an aborted walk")
511+
}
512+
}
513+
514+
mutating func rootDef(address: Value, path: UnusedWalkingPath) -> WalkResult {
515+
assert(enclosingScope == nil, "rootDef should only called once")
516+
// Try identifying the address a pointer originates from
517+
if let p2ai = address as? PointerToAddressInst, let originatingAddr = p2ai.originatingAddress {
518+
return walkUp(address: originatingAddr, path: path)
519+
}
520+
enclosingScope = .base(AccessBase(baseAddress: address))
521+
return .continueWalk
522+
}
523+
524+
mutating func walkUp(address: Value, path: UnusedWalkingPath) -> WalkResult {
525+
if let ba = address as? BeginAccessInst {
526+
enclosingScope = .scope(ba)
527+
return .continueWalk
528+
}
529+
return walkUpDefault(address: address, path: path)
530+
}
531+
}
532+
483533
private struct AccessPathWalker : AddressUseDefWalker {
484534
var result = AccessPath.unidentified()
485-
var foundBeginAccess: BeginAccessInst?
535+
536+
// List of nested BeginAccessInst: inside-out order.
537+
var foundBeginAccesses = SingleInlineArray<BeginAccessInst>()
538+
486539
let enforceConstantProjectionPath: Bool
487540

488541
init(enforceConstantProjectionPath: Bool = false) {
@@ -528,18 +581,9 @@ private struct AccessPathWalker : AddressUseDefWalker {
528581
mutating func rootDef(address: Value, path: Path) -> WalkResult {
529582
assert(result.base == .unidentified, "rootDef should only called once")
530583
// Try identifying the address a pointer originates from
531-
if let p2ai = address as? PointerToAddressInst {
532-
if let originatingAddr = p2ai.originatingAddress {
533-
return walkUp(address: originatingAddr, path: path)
534-
} else if let global = p2ai.resultOfGlobalAddressorCall {
535-
self.result = AccessPath(base: .global(global), projectionPath: path.projectionPath)
536-
return .continueWalk
537-
} else {
538-
self.result = AccessPath(base: .pointer(p2ai), projectionPath: path.projectionPath)
539-
return .continueWalk
540-
}
584+
if let p2ai = address as? PointerToAddressInst, let originatingAddr = p2ai.originatingAddress {
585+
return walkUp(address: originatingAddr, path: path)
541586
}
542-
543587
let base = AccessBase(baseAddress: address)
544588
self.result = AccessPath(base: base, projectionPath: path.projectionPath)
545589
return .continueWalk
@@ -557,8 +601,8 @@ private struct AccessPathWalker : AddressUseDefWalker {
557601
// An `index_addr` instruction cannot be derived from an address
558602
// projection. Bail out
559603
return .abortWalk
560-
} else if let ba = address as? BeginAccessInst, foundBeginAccess == nil {
561-
foundBeginAccess = ba
604+
} else if let ba = address as? BeginAccessInst {
605+
foundBeginAccesses.push(ba)
562606
}
563607
return walkUpDefault(address: address, path: path.with(indexAddr: false))
564608
}
@@ -611,17 +655,20 @@ extension Value {
611655
public var accessPathWithScope: (AccessPath, scope: BeginAccessInst?) {
612656
var walker = AccessPathWalker()
613657
walker.walk(startAt: self)
614-
return (walker.result, walker.foundBeginAccess)
658+
return (walker.result, walker.foundBeginAccesses.first)
615659
}
616660

617661
/// Computes the enclosing access scope of this address value.
618662
public var enclosingAccessScope: EnclosingScope {
663+
var walker = EnclosingAccessWalker()
664+
walker.walk(startAt: self)
665+
return walker.enclosingScope ?? .base(.unidentified)
666+
}
667+
668+
public var accessBaseWithScopes: (AccessBase, SingleInlineArray<BeginAccessInst>) {
619669
var walker = AccessPathWalker()
620670
walker.walk(startAt: self)
621-
if let ba = walker.foundBeginAccess {
622-
return .scope(ba)
623-
}
624-
return .base(walker.result.base)
671+
return (walker.result.base, walker.foundBeginAccesses)
625672
}
626673

627674
/// The root definition of a reference, obtained by skipping ownership forwarding and ownership transition.

‎include/swift/SIL/SILBridging.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1195,6 +1195,10 @@ struct BridgedBuilder{
11951195

11961196
SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createEndAccess(BridgedValue value) const;
11971197

1198+
SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createEndApply(BridgedValue value) const;
1199+
1200+
SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createAbortApply(BridgedValue value) const;
1201+
11981202
SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createConvertFunction(BridgedValue originalFunction, BridgedType resultType, bool withoutActuallyEscaping) const;
11991203
SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createConvertEscapeToNoEscape(BridgedValue originalFunction, BridgedType resultType, bool isLifetimeGuaranteed) const;
12001204

‎include/swift/SIL/SILBridgingImpl.h

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2296,6 +2296,16 @@ BridgedInstruction BridgedBuilder::createEndAccess(BridgedValue value) const {
22962296
return {unbridged().createEndAccess(regularLoc(), value.getSILValue(), false)};
22972297
}
22982298

2299+
BridgedInstruction BridgedBuilder::createEndApply(BridgedValue value) const {
2300+
swift::ASTContext &ctxt = unbridged().getASTContext();
2301+
return {unbridged().createEndApply(regularLoc(), value.getSILValue(),
2302+
swift::SILType::getEmptyTupleType(ctxt))};
2303+
}
2304+
2305+
BridgedInstruction BridgedBuilder::createAbortApply(BridgedValue value) const {
2306+
return {unbridged().createAbortApply(regularLoc(), value.getSILValue())};
2307+
}
2308+
22992309
BridgedInstruction BridgedBuilder::createConvertFunction(BridgedValue originalFunction, BridgedType resultType, bool withoutActuallyEscaping) const {
23002310
return {unbridged().createConvertFunction(regularLoc(), originalFunction.getSILValue(), resultType.unbridged(), withoutActuallyEscaping)};
23012311
}

‎include/swift/SIL/SILInstruction.h

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3217,6 +3217,13 @@ class EndApplyInst;
32173217
class AbortApplyInst;
32183218
class EndBorrowInst;
32193219

3220+
struct EndApplyFilter {
3221+
std::optional<Operand*> operator()(Operand *use) const;
3222+
};
3223+
3224+
using EndApplyRange = OptionalTransformRange<ValueBase::use_range,
3225+
EndApplyFilter>;
3226+
32203227
/// BeginApplyInst - Represents the beginning of the full application of
32213228
/// a yield_once coroutine (up until the coroutine yields a value back).
32223229
class BeginApplyInst final
@@ -3267,6 +3274,8 @@ class BeginApplyInst final
32673274
&getAllResultsBuffer().drop_back(isCalleeAllocated() ? 1 : 0).back());
32683275
}
32693276

3277+
EndApplyRange getEndApplyUses() const;
3278+
32703279
MultipleValueInstructionResult *getCalleeAllocationResult() const {
32713280
if (!isCalleeAllocated()) {
32723281
return nullptr;
@@ -3336,6 +3345,25 @@ class EndApplyInst
33363345
}
33373346
};
33383347

3348+
inline std::optional<Operand*>
3349+
EndApplyFilter::operator()(Operand *use) const {
3350+
// An end_borrow ends the coroutine scope at a dead-end block without
3351+
// terminating the coroutine.
3352+
switch (use->getUser()->getKind()) {
3353+
case SILInstructionKind::EndApplyInst:
3354+
case SILInstructionKind::AbortApplyInst:
3355+
case SILInstructionKind::EndBorrowInst:
3356+
return use;
3357+
default:
3358+
return std::nullopt;
3359+
}
3360+
}
3361+
3362+
inline EndApplyRange BeginApplyInst::getEndApplyUses() const {
3363+
return makeOptionalTransformRange(
3364+
getTokenResult()->getUses(), EndApplyFilter());
3365+
}
3366+
33393367
//===----------------------------------------------------------------------===//
33403368
// Literal instructions.
33413369
//===----------------------------------------------------------------------===//
@@ -8802,6 +8830,11 @@ class MarkDependenceInst
88028830
uint8_t(MarkDependenceKind::NonEscaping);
88038831
}
88048832

8833+
void settleToEscaping() {
8834+
sharedUInt8().MarkDependenceInst.dependenceKind =
8835+
uint8_t(MarkDependenceKind::Escaping);
8836+
}
8837+
88058838
/// Visit the instructions that end the lifetime of an OSSA on-stack closure.
88068839
bool visitNonEscapingLifetimeEnds(
88078840
llvm::function_ref<bool (Operand*)> visitScopeEnd,

‎lib/SIL/IR/SILInstructions.cpp

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -712,8 +712,8 @@ void BeginApplyInst::getCoroutineEndPoints(
712712
SmallVectorImpl<EndApplyInst *> &endApplyInsts,
713713
SmallVectorImpl<AbortApplyInst *> &abortApplyInsts,
714714
SmallVectorImpl<EndBorrowInst *> *endBorrowInsts) const {
715-
for (auto *tokenUse : getTokenResult()->getUses()) {
716-
auto *user = tokenUse->getUser();
715+
for (auto *use : getEndApplyUses()) {
716+
auto *user = use->getUser();
717717
if (auto *end = dyn_cast<EndApplyInst>(user)) {
718718
endApplyInsts.push_back(end);
719719
continue;
@@ -733,19 +733,19 @@ void BeginApplyInst::getCoroutineEndPoints(
733733
SmallVectorImpl<Operand *> &endApplyInsts,
734734
SmallVectorImpl<Operand *> &abortApplyInsts,
735735
SmallVectorImpl<Operand *> *endBorrowInsts) const {
736-
for (auto *tokenUse : getTokenResult()->getUses()) {
737-
auto *user = tokenUse->getUser();
736+
for (auto *use : getEndApplyUses()) {
737+
auto *user = use->getUser();
738738
if (isa<EndApplyInst>(user)) {
739-
endApplyInsts.push_back(tokenUse);
739+
endApplyInsts.push_back(use);
740740
continue;
741741
}
742742
if (isa<AbortApplyInst>(user)) {
743-
abortApplyInsts.push_back(tokenUse);
743+
abortApplyInsts.push_back(use);
744744
continue;
745745
}
746746

747747
assert(isa<EndBorrowInst>(user));
748-
abortApplyInsts.push_back(tokenUse);
748+
abortApplyInsts.push_back(use);
749749
}
750750
}
751751

‎lib/SIL/Utils/OwnershipUtils.cpp

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -694,8 +694,7 @@ bool BorrowingOperand::visitScopeEndingUses(
694694
}
695695
case BorrowingOperandKind::BeginApply: {
696696
bool deadApply = true;
697-
auto *user = cast<BeginApplyInst>(op->getUser());
698-
for (auto *use : user->getTokenResult()->getUses()) {
697+
for (auto *use : cast<BeginApplyInst>(op->getUser())->getEndApplyUses()) {
699698
deadApply = false;
700699
if (!visitScopeEnd(use))
701700
return false;

‎lib/SIL/Verifier/SILOwnershipVerifier.cpp

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -583,8 +583,7 @@ bool SILValueOwnershipChecker::checkYieldWithoutLifetimeEndingUses(
583583
// If we have a guaranteed value, make sure that all uses are before our
584584
// end_yield.
585585
SmallVector<Operand *, 4> coroutineEndUses;
586-
for (auto *use : yield->getParent<BeginApplyInst>()->
587-
getTokenResult()->getUses()) {
586+
for (auto *use : yield->getParent<BeginApplyInst>()->getEndApplyUses()) {
588587
coroutineEndUses.push_back(use);
589588
}
590589

‎lib/SILOptimizer/Mandatory/AddressLowering.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2682,8 +2682,8 @@ void ApplyRewriter::convertBeginApplyWithOpaqueYield() {
26822682
SILValue load =
26832683
resultBuilder.emitLoadBorrowOperation(callLoc, &newResult);
26842684
oldResult.replaceAllUsesWith(load);
2685-
for (auto *user : origCall->getTokenResult()->getUsers()) {
2686-
pass.getBuilder(user->getIterator())
2685+
for (auto *use : origCall->getEndApplyUses()) {
2686+
pass.getBuilder(use->getUser()->getIterator())
26872687
.createEndBorrow(pass.genLoc(), load);
26882688
}
26892689
} else {

‎lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -453,7 +453,8 @@ static bool visitScopeEndsRequiringInit(
453453
// Check for yields from a modify coroutine.
454454
if (auto bai =
455455
dyn_cast_or_null<BeginApplyInst>(operand->getDefiningInstruction())) {
456-
for (auto *inst : bai->getTokenResult()->getUsers()) {
456+
for (auto *use : bai->getEndApplyUses()) {
457+
auto *inst = use->getUser();
457458
assert(isa<EndApplyInst>(inst) || isa<AbortApplyInst>(inst) ||
458459
isa<EndBorrowInst>(inst));
459460
visit(inst, ScopeRequiringFinalInit::Coroutine);

‎lib/SILOptimizer/Transforms/TempRValueElimination.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -210,8 +210,8 @@ collectLoads(Operand *addressUse, CopyAddrInst *originalCopy,
210210
// Register 'end_apply'/'abort_apply' as loads as well
211211
// 'checkNoSourceModification' should check instructions until
212212
// 'end_apply'/'abort_apply'.
213-
for (auto tokenUse : beginApply->getTokenResult()->getUses()) {
214-
SILInstruction *tokenUser = tokenUse->getUser();
213+
for (auto *tokenUse : beginApply->getEndApplyUses()) {
214+
auto *tokenUser = tokenUse->getUser();
215215
if (tokenUser->getParent() != block)
216216
return false;
217217
loadInsts.insert(tokenUser);

‎lib/SILOptimizer/Utils/Devirtualize.cpp

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -643,14 +643,11 @@ replaceBeginApplyInst(SILBuilder &builder, SILPassManager *pm, SILLocation loc,
643643
if (newArgBorrows.empty())
644644
return {newBAI, changedCFG};
645645

646-
SILValue token = newBAI->getTokenResult();
647-
648-
// The token will only be used by end_apply and abort_apply. Use that to
649-
// insert the end_borrows we need /after/ those uses.
650-
for (auto *use : token->getUses()) {
646+
// Insert the end_borrows after end_apply and abort_apply users.
647+
for (auto *use : newBAI->getEndApplyUses()) {
651648
SILBuilderWithScope borrowBuilder(
652-
&*std::next(use->getUser()->getIterator()),
653-
builder.getBuilderContext());
649+
&*std::next(use->getUser()->getIterator()),
650+
builder.getBuilderContext());
654651
for (SILValue borrow : newArgBorrows) {
655652
borrowBuilder.createEndBorrow(loc, borrow);
656653
}

0 commit comments

Comments
 (0)