@@ -37,6 +37,44 @@ static void diagnose(ASTContext &Context, SourceLoc loc, Diag<T...> diag,
3737 Context.Diags .diagnose (loc, diag, std::forward<U>(args)...);
3838}
3939
40+ namespace {
41+
42+ // / A helper class to update an instruction iterator if
43+ // / removal of instructions would invalidate it.
44+ class DeleteInstructionsHandler : public DeleteNotificationHandler {
45+ SILBasicBlock::iterator &CurrentI;
46+ SILModule &Module;
47+
48+ public:
49+ DeleteInstructionsHandler (SILBasicBlock::iterator &I)
50+ : CurrentI(I), Module(I->getModule ()) {
51+ Module.registerDeleteNotificationHandler (this );
52+ }
53+
54+ ~DeleteInstructionsHandler () {
55+ // Unregister the handler.
56+ Module.removeDeleteNotificationHandler (this );
57+ }
58+
59+ // Handling of instruction removal notifications.
60+ bool needsNotifications () { return true ; }
61+
62+ // Handle notifications about removals of instructions.
63+ void handleDeleteNotification (swift::ValueBase *Value) {
64+ if (auto DeletedI = dyn_cast<SILInstruction>(Value)) {
65+ if (CurrentI == SILBasicBlock::iterator (DeletedI)) {
66+ if (CurrentI != CurrentI->getParent ()->begin ()) {
67+ --CurrentI;
68+ } else {
69+ ++CurrentI;
70+ }
71+ }
72+ }
73+ }
74+ };
75+
76+ } // end of namespace
77+
4078// / \brief Fixup reference counts after inlining a function call (which is a
4179// / no-op unless the function is a thick function). Note that this function
4280// / makes assumptions about the release/retain convention of thick function
@@ -419,8 +457,9 @@ runOnFunctionRecursively(SILFunction *F, FullApplySite AI,
419457 // Reestablish our iterator if it wrapped.
420458 if (I == ApplyBlock->end ())
421459 I = ApplyBlock->begin ();
422- else
423- ++I;
460+
461+ // Update the iterator when instructions are removed.
462+ DeleteInstructionsHandler DeletionHandler (I);
424463
425464 // If the inlined apply was a thick function, then we need to balance the
426465 // reference counts for correctness.
@@ -433,8 +472,9 @@ runOnFunctionRecursively(SILFunction *F, FullApplySite AI,
433472
434473 // Reposition iterators possibly invalidated by mutation.
435474 FI = SILFunction::iterator (ApplyBlock);
436- I = ApplyBlock->begin ();
437475 E = ApplyBlock->end ();
476+ assert (FI == SILFunction::iterator (I->getParent ()) &&
477+ " Mismatch between the instruction and basic block" );
438478 ++NumMandatoryInlines;
439479 }
440480 }
0 commit comments