Skip to content

Commit 80e1ae3

Browse files
Merge pull request #73274 from nate-chandler/lifetime-completion/20240425/1
[LifetimeCompletion] Complete coro token lifetimes with end_borrow.
2 parents f3224ca + 29d046f commit 80e1ae3

File tree

12 files changed

+202
-15
lines changed

12 files changed

+202
-15
lines changed

include/swift/SIL/OwnershipUtils.h

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -490,6 +490,7 @@ class BorrowedValueKind {
490490
BeginBorrow,
491491
SILFunctionArgument,
492492
Phi,
493+
BeginApplyToken,
493494
};
494495

495496
private:
@@ -520,6 +521,13 @@ class BorrowedValueKind {
520521
}
521522
return Kind::Phi;
522523
}
524+
case ValueKind::MultipleValueInstructionResult:
525+
if (!isaResultOf<BeginApplyInst>(value))
526+
return Kind::Invalid;
527+
if (value->isBeginApplyToken()) {
528+
return Kind::BeginApplyToken;
529+
}
530+
return Kind::Invalid;
523531
}
524532
}
525533

@@ -540,6 +548,7 @@ class BorrowedValueKind {
540548
case BorrowedValueKind::BeginBorrow:
541549
case BorrowedValueKind::LoadBorrow:
542550
case BorrowedValueKind::Phi:
551+
case BorrowedValueKind::BeginApplyToken:
543552
return true;
544553
case BorrowedValueKind::SILFunctionArgument:
545554
return false;

include/swift/SIL/SILInstruction.h

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3163,6 +3163,7 @@ class PartialApplyInst final
31633163

31643164
class EndApplyInst;
31653165
class AbortApplyInst;
3166+
class EndBorrowInst;
31663167

31673168
/// BeginApplyInst - Represents the beginning of the full application of
31683169
/// a yield_once coroutine (up until the coroutine yields a value back).
@@ -3216,10 +3217,13 @@ class BeginApplyInst final
32163217

32173218
void getCoroutineEndPoints(
32183219
SmallVectorImpl<EndApplyInst *> &endApplyInsts,
3219-
SmallVectorImpl<AbortApplyInst *> &abortApplyInsts) const;
3220+
SmallVectorImpl<AbortApplyInst *> &abortApplyInsts,
3221+
SmallVectorImpl<EndBorrowInst *> *endBorrowInsts = nullptr) const;
32203222

3221-
void getCoroutineEndPoints(SmallVectorImpl<Operand *> &endApplyInsts,
3222-
SmallVectorImpl<Operand *> &abortApplyInsts) const;
3223+
void getCoroutineEndPoints(
3224+
SmallVectorImpl<Operand *> &endApplyInsts,
3225+
SmallVectorImpl<Operand *> &abortApplyInsts,
3226+
SmallVectorImpl<Operand *> *endBorrowInsts = nullptr) const;
32233227
};
32243228

32253229
/// AbortApplyInst - Unwind the full application of a yield_once coroutine.
@@ -4539,8 +4543,6 @@ class StoreInst
45394543
}
45404544
};
45414545

4542-
class EndBorrowInst;
4543-
45444546
/// Represents a load of a borrowed value. Must be paired with an end_borrow
45454547
/// instruction in its use-def list.
45464548
class LoadBorrowInst :

include/swift/SIL/SILValue.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -590,6 +590,8 @@ class ValueBase : public SILNode, public SILAllocated<ValueBase> {
590590

591591
bool isGuaranteedForwarding() const;
592592

593+
bool isBeginApplyToken() const;
594+
593595
/// Unsafely eliminate moveonly from this value's type. Returns true if the
594596
/// value's underlying type was move only and thus was changed. Returns false
595597
/// otherwise.

lib/SIL/IR/SILInstructions.cpp

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -693,29 +693,41 @@ BeginApplyInst *BeginApplyInst::create(
693693

694694
void BeginApplyInst::getCoroutineEndPoints(
695695
SmallVectorImpl<EndApplyInst *> &endApplyInsts,
696-
SmallVectorImpl<AbortApplyInst *> &abortApplyInsts) const {
696+
SmallVectorImpl<AbortApplyInst *> &abortApplyInsts,
697+
SmallVectorImpl<EndBorrowInst *> *endBorrowInsts) const {
697698
for (auto *tokenUse : getTokenResult()->getUses()) {
698699
auto *user = tokenUse->getUser();
699700
if (auto *end = dyn_cast<EndApplyInst>(user)) {
700701
endApplyInsts.push_back(end);
701702
continue;
702703
}
703-
704-
abortApplyInsts.push_back(cast<AbortApplyInst>(user));
704+
if (auto *abort = dyn_cast<AbortApplyInst>(user)) {
705+
abortApplyInsts.push_back(abort);
706+
continue;
707+
}
708+
auto *end = cast<EndBorrowInst>(user);
709+
if (endBorrowInsts) {
710+
endBorrowInsts->push_back(end);
711+
}
705712
}
706713
}
707714

708715
void BeginApplyInst::getCoroutineEndPoints(
709716
SmallVectorImpl<Operand *> &endApplyInsts,
710-
SmallVectorImpl<Operand *> &abortApplyInsts) const {
717+
SmallVectorImpl<Operand *> &abortApplyInsts,
718+
SmallVectorImpl<Operand *> *endBorrowInsts) const {
711719
for (auto *tokenUse : getTokenResult()->getUses()) {
712720
auto *user = tokenUse->getUser();
713721
if (isa<EndApplyInst>(user)) {
714722
endApplyInsts.push_back(tokenUse);
715723
continue;
716724
}
725+
if (isa<AbortApplyInst>(user)) {
726+
abortApplyInsts.push_back(tokenUse);
727+
continue;
728+
}
717729

718-
assert(isa<AbortApplyInst>(user));
730+
assert(isa<EndBorrowInst>(user));
719731
abortApplyInsts.push_back(tokenUse);
720732
}
721733
}

lib/SIL/IR/SILValue.cpp

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,13 @@ bool ValueBase::isGuaranteedForwarding() const {
191191
return phi->isGuaranteedForwarding();
192192
}
193193

194+
bool ValueBase::isBeginApplyToken() const {
195+
auto *result = isaResultOf<BeginApplyInst>(this);
196+
if (!result)
197+
return false;
198+
return result->isBeginApplyToken();
199+
}
200+
194201
bool ValueBase::hasDebugTrace() const {
195202
for (auto *op : getUses()) {
196203
if (auto *debugValue = dyn_cast<DebugValueInst>(op->getUser())) {

lib/SIL/Utils/OwnershipUtils.cpp

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -913,6 +913,9 @@ void BorrowedValueKind::print(llvm::raw_ostream &os) const {
913913
case BorrowedValueKind::Phi:
914914
os << "Phi";
915915
return;
916+
case BorrowedValueKind::BeginApplyToken:
917+
os << "BeginApply";
918+
return;
916919
}
917920
llvm_unreachable("Covered switch isn't covered?!");
918921
}
@@ -945,6 +948,7 @@ bool BorrowedValue::visitLocalScopeEndingUses(
945948
case BorrowedValueKind::LoadBorrow:
946949
case BorrowedValueKind::BeginBorrow:
947950
case BorrowedValueKind::Phi:
951+
case BorrowedValueKind::BeginApplyToken:
948952
for (auto *use : lookThroughBorrowedFromUser(value)->getUses()) {
949953
if (use->isLifetimeEnding()) {
950954
if (!visitor(use))
@@ -1819,6 +1823,7 @@ class FindEnclosingDefs {
18191823

18201824
case BorrowedValueKind::LoadBorrow:
18211825
case BorrowedValueKind::SILFunctionArgument:
1826+
case BorrowedValueKind::BeginApplyToken:
18221827
// There is no enclosing def on this path.
18231828
return true;
18241829
}
@@ -2037,6 +2042,7 @@ class FindEnclosingDefs {
20372042
}
20382043
case BorrowedValueKind::LoadBorrow:
20392044
case BorrowedValueKind::SILFunctionArgument:
2045+
case BorrowedValueKind::BeginApplyToken:
20402046
// There is no enclosing def on this path.
20412047
break;
20422048
}

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -447,7 +447,8 @@ static bool visitScopeEndsRequiringInit(
447447
if (auto bai =
448448
dyn_cast_or_null<BeginApplyInst>(operand->getDefiningInstruction())) {
449449
for (auto *inst : bai->getTokenResult()->getUsers()) {
450-
assert(isa<EndApplyInst>(inst) || isa<AbortApplyInst>(inst));
450+
assert(isa<EndApplyInst>(inst) || isa<AbortApplyInst>(inst) ||
451+
isa<EndBorrowInst>(inst));
451452
visit(inst, ScopeRequiringFinalInit::Coroutine);
452453
}
453454
return true;

lib/SILOptimizer/Utils/CanonicalizeBorrowScope.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,7 @@ SILValue CanonicalizeBorrowScope::getCanonicalBorrowedDef(SILValue def) {
140140

141141
case BorrowedValueKind::LoadBorrow:
142142
case BorrowedValueKind::Phi:
143+
case BorrowedValueKind::BeginApplyToken:
143144
break;
144145
}
145146
}
@@ -170,6 +171,7 @@ bool CanonicalizeBorrowScope::computeBorrowLiveness() {
170171
// can handle persistentCopies.
171172
return false;
172173
case BorrowedValueKind::BeginBorrow:
174+
case BorrowedValueKind::BeginApplyToken:
173175
break;
174176
}
175177
// Note that there is no need to look through any reborrows. The reborrowed

lib/SILOptimizer/Utils/LoopUtils.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -307,7 +307,8 @@ bool swift::canDuplicateLoopInstruction(SILLoop *L, SILInstruction *I) {
307307
if (auto BAI = dyn_cast<BeginApplyInst>(I)) {
308308
for (auto UI : BAI->getTokenResult()->getUses()) {
309309
auto User = UI->getUser();
310-
assert(isa<EndApplyInst>(User) || isa<AbortApplyInst>(User));
310+
assert(isa<EndApplyInst>(User) || isa<AbortApplyInst>(User) ||
311+
isa<EndBorrowInst>(User));
311312
if (!L->contains(User))
312313
return false;
313314
}

lib/SILOptimizer/Utils/SILInliner.cpp

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,10 +41,11 @@ static bool canInlineBeginApply(BeginApplyInst *BA) {
4141
if (isa<EndApplyInst>(user)) {
4242
if (hasEndApply) return false;
4343
hasEndApply = true;
44-
} else {
45-
assert(isa<AbortApplyInst>(user));
44+
} else if (isa<AbortApplyInst>(user)) {
4645
if (hasAbortApply) return false;
4746
hasAbortApply = true;
47+
} else {
48+
assert(isa<EndBorrowInst>(user));
4849
}
4950
}
5051

@@ -95,6 +96,8 @@ class BeginApplySite {
9596
SILBasicBlock *AbortApplyBB = nullptr;
9697
SILBasicBlock *AbortApplyReturnBB = nullptr;
9798

99+
SmallVector<EndBorrowInst *, 2> EndBorrows;
100+
98101
public:
99102
BeginApplySite(BeginApplyInst *BeginApply, SILLocation Loc,
100103
SILBuilder *Builder)
@@ -112,7 +115,8 @@ class BeginApplySite {
112115
SmallVectorImpl<SILInstruction *> &endBorrowInsertPts) {
113116
SmallVector<EndApplyInst *, 1> endApplyInsts;
114117
SmallVector<AbortApplyInst *, 1> abortApplyInsts;
115-
BeginApply->getCoroutineEndPoints(endApplyInsts, abortApplyInsts);
118+
BeginApply->getCoroutineEndPoints(endApplyInsts, abortApplyInsts,
119+
&EndBorrows);
116120
while (!endApplyInsts.empty()) {
117121
auto *endApply = endApplyInsts.pop_back_val();
118122
collectEndApply(endApply);
@@ -257,6 +261,8 @@ class BeginApplySite {
257261
EndApply->eraseFromParent();
258262
if (AbortApply)
259263
AbortApply->eraseFromParent();
264+
for (auto *EndBorrow : EndBorrows)
265+
EndBorrow->eraseFromParent();
260266

261267
assert(!BeginApply->hasUsesOfAnyResult());
262268
}

test/SILOptimizer/inline_begin_apply.sil

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -690,3 +690,118 @@ bb_unwind:
690690
unwind
691691

692692
}
693+
694+
sil [ossa] [transparent] @simplest_in_guaranteed : $@yield_once @convention(thin) () -> @yields @in_guaranteed () {
695+
entry:
696+
%tuple = tuple ()
697+
%addr = alloc_stack $()
698+
store %tuple to [trivial] %addr : $*()
699+
yield %addr : $*(), resume resume_block, unwind unwind_block
700+
701+
resume_block:
702+
dealloc_stack %addr : $*()
703+
%retval = tuple ()
704+
return %retval : $()
705+
706+
unwind_block:
707+
dealloc_stack %addr : $*()
708+
unwind
709+
}
710+
711+
// CHECK-LABEL: sil [ossa] @ends_1 : {{.*}} {
712+
// CHECK: [[VOID:%[^,]+]] = tuple ()
713+
// CHECK: [[ADDR:%[^,]+]] = alloc_stack $()
714+
// CHECK: store [[VOID]] to [trivial] [[ADDR]]
715+
// CHECK: cond_br undef, [[GOOD:bb[0-9]+]], [[BAD:bb[0-9]+]]
716+
// CHECK: [[GOOD]]:
717+
// CHECK: dealloc_stack [[ADDR]]
718+
// CHECK: br [[EXIT:bb[0-9]+]]
719+
// CHECK: [[BAD]]:
720+
// CHECK: cond_br undef, [[NORMAL_BAD:bb[0-9]+]], [[REALLY_BAD:bb[0-9]+]]
721+
// CHECK: [[NORMAL_BAD]]:
722+
// CHECK: dealloc_stack [[ADDR]]
723+
// CHECK: br [[EXIT]]
724+
// CHECK: [[REALLY_BAD]]:
725+
// CHECK: unreachable
726+
// CHECK: [[EXIT]]:
727+
// CHECK-LABEL: } // end sil function 'ends_1'
728+
sil [ossa] @ends_1 : $@convention(thin) () -> () {
729+
entry:
730+
%simplest = function_ref @simplest_in_guaranteed : $@yield_once @convention(thin) () -> @yields @in_guaranteed ()
731+
(%_, %token) = begin_apply %simplest() : $@yield_once @convention(thin) () -> @yields @in_guaranteed ()
732+
cond_br undef, good, bad
733+
734+
good:
735+
end_apply %token as $()
736+
br exit
737+
738+
bad:
739+
cond_br undef, normal_bad, really_bad
740+
741+
normal_bad:
742+
abort_apply %token
743+
br exit
744+
745+
really_bad:
746+
end_borrow %token : $Builtin.SILToken
747+
unreachable
748+
749+
exit:
750+
%retval = tuple ()
751+
return %retval : $()
752+
}
753+
754+
// Test handling for multiple end_borrows of a token.
755+
// CHECK-LABEL: sil [ossa] @ends_2 : {{.*}} {
756+
// CHECK: [[VOID:%[^,]+]] = tuple ()
757+
// CHECK: [[ADDR:%[^,]+]] = alloc_stack $()
758+
// CHECK: store [[VOID]] to [trivial] [[ADDR]] : $*()
759+
// CHECK: cond_br undef, [[GOOD:bb[0-9]+]], [[BAD:bb[0-9]+]]
760+
// CHECK: [[GOOD]]:
761+
// CHECK: dealloc_stack [[ADDR]]
762+
// CHECK: br [[EXIT:bb[0-9]+]]
763+
// CHECK: [[BAD]]:
764+
// CHECK: cond_br undef, [[NORMAL_BAD:bb[0-9]+]], [[REALLY_BAD:bb[0-9]+]]
765+
// CHECK: [[NORMAL_BAD]]:
766+
// CHECK: dealloc_stack [[ADDR]]
767+
// CHECK: br [[EXIT]]
768+
// CHECK: [[REALLY_BAD]]:
769+
// CHECK: cond_br undef, [[BAD_NEWS_BUNNIES:bb[0-9]+]], [[BAD_NEWS_BEARS:bb[0-9]+]]
770+
// CHECK: [[BAD_NEWS_BUNNIES]]:
771+
// CHECK: unreachable
772+
// CHECK: [[BAD_NEWS_BEARS]]:
773+
// CHECK: unreachable
774+
// CHECK: [[EXIT]]:
775+
// CHECK-LABEL: } // end sil function 'ends_2'
776+
sil [ossa] @ends_2 : $@convention(thin) () -> () {
777+
entry:
778+
%simplest = function_ref @simplest_in_guaranteed : $@yield_once @convention(thin) () -> @yields @in_guaranteed ()
779+
(%_, %token) = begin_apply %simplest() : $@yield_once @convention(thin) () -> @yields @in_guaranteed ()
780+
cond_br undef, good, bad
781+
782+
good:
783+
end_apply %token as $()
784+
br exit
785+
786+
bad:
787+
cond_br undef, normal_bad, really_bad
788+
789+
normal_bad:
790+
abort_apply %token
791+
br exit
792+
793+
really_bad:
794+
cond_br undef, bad_news_bunnies, bad_news_bears
795+
796+
bad_news_bunnies:
797+
end_borrow %token : $Builtin.SILToken
798+
unreachable
799+
800+
bad_news_bears:
801+
end_borrow %token : $Builtin.SILToken
802+
unreachable
803+
804+
exit:
805+
%retval = tuple ()
806+
return %retval : $()
807+
}

test/SILOptimizer/ossa_lifetime_completion.sil

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -513,3 +513,27 @@ entry(%instance : @owned $C):
513513
store %instance to [init] %addr : $*C
514514
unreachable
515515
}
516+
517+
// CHECK-LABEL: sil [ossa] @begin_apply : {{.*}} {
518+
// CHECK: ({{%[^,]+}}, [[TOKEN:%[^,]+]]) = begin_apply undef()
519+
// CHECK: cond_br undef, [[LEFT:bb[0-9]+]], [[RIGHT:bb[0-9]+]]
520+
// CHECK: [[LEFT]]:
521+
// CHECK: end_apply [[TOKEN]]
522+
// CHECK: [[RIGHT]]:
523+
// CHECK: end_borrow [[TOKEN]]
524+
// CHECK: unreachable
525+
// CHECK-LABEL: } // end sil function 'begin_apply'
526+
sil [ossa] @begin_apply : $@convention(thin) () -> () {
527+
entry:
528+
(%_, %token) = begin_apply undef() : $@yield_once @convention(thin) () -> (@yields @in_guaranteed ())
529+
specify_test "ossa-lifetime-completion %token"
530+
cond_br undef, left, right
531+
532+
left:
533+
end_apply %token as $()
534+
%retval = tuple ()
535+
return %retval : $()
536+
537+
right:
538+
unreachable
539+
}

0 commit comments

Comments
 (0)