@@ -202,9 +202,14 @@ static void promoteDebugValueAddr(DebugValueInst *dvai, SILValue value,
202202
203203// / Returns true if \p I is a load which loads from \p ASI.
204204static bool isLoadFromStack (SILInstruction *i, AllocStackInst *asi) {
205- if (!isa<LoadInst>(i))
205+ if (!isa<LoadInst>(i) && !isa<LoadBorrowInst>(i) )
206206 return false ;
207207
208+ if (auto *lbi = dyn_cast<LoadBorrowInst>(i)) {
209+ if (BorrowedValue (lbi).hasReborrow ())
210+ return false ;
211+ }
212+
208213 // Skip struct and tuple address projections.
209214 ValueBase *op = i->getOperand (0 );
210215 while (op != asi) {
@@ -218,9 +223,9 @@ static bool isLoadFromStack(SILInstruction *i, AllocStackInst *asi) {
218223
219224// / Collects all load instructions which (transitively) use \p I as address.
220225static void collectLoads (SILInstruction *i,
221- SmallVectorImpl<LoadInst *> &foundLoads) {
222- if (auto *load = dyn_cast<LoadInst >(i)) {
223- foundLoads.push_back (load );
226+ SmallVectorImpl<SILInstruction *> &foundLoads) {
227+ if (isa<LoadInst>(i) || isa<LoadBorrowInst >(i)) {
228+ foundLoads.push_back (i );
224229 return ;
225230 }
226231 if (!isa<UncheckedAddrCastInst>(i) && !isa<StructElementAddrInst>(i) &&
@@ -234,22 +239,23 @@ static void collectLoads(SILInstruction *i,
234239}
235240
236241static void
237- replaceLoad (LoadInst *li , SILValue newValue, AllocStackInst *asi,
242+ replaceLoad (SILInstruction *inst , SILValue newValue, AllocStackInst *asi,
238243 SILBuilderContext &ctx, InstructionDeleter &deleter,
239244 SmallVectorImpl<SILInstruction *> &instructionsToDelete) {
245+ assert (isa<LoadInst>(inst) || isa<LoadBorrowInst>(inst));
240246 ProjectionPath projections (newValue->getType ());
241- SILValue op = li ->getOperand ();
242- SILBuilderWithScope builder (li , ctx);
247+ SILValue op = inst ->getOperand (0 );
248+ SILBuilderWithScope builder (inst , ctx);
243249 SILOptScope scope;
244250
245251 while (op != asi) {
246252 assert (isa<UncheckedAddrCastInst>(op) || isa<StructElementAddrInst>(op) ||
247253 isa<TupleElementAddrInst>(op) &&
248254 " found instruction that should have been skipped in "
249255 " isLoadFromStack" );
250- auto *inst = cast<SingleValueInstruction>(op);
251- projections.push_back (Projection (inst ));
252- op = inst ->getOperand (0 );
256+ auto *projInst = cast<SingleValueInstruction>(op);
257+ projections.push_back (Projection (projInst ));
258+ op = projInst ->getOperand (0 );
253259 }
254260
255261 for (const auto &proj : llvm::reverse (projections)) {
@@ -262,33 +268,51 @@ replaceLoad(LoadInst *li, SILValue newValue, AllocStackInst *asi,
262268 // to guaranteed!
263269 if (proj.getKind () == ProjectionKind::Struct ||
264270 proj.getKind () == ProjectionKind::Tuple) {
265- if (auto opVal = scope.borrowValue (li , newValue)) {
271+ if (auto opVal = scope.borrowValue (inst , newValue)) {
266272 assert (*opVal != newValue &&
267273 " Valid value should be different from input value" );
268274 newValue = *opVal;
269275 }
270276 }
271277 newValue =
272- proj.createObjectProjection (builder, li ->getLoc (), newValue).get ();
278+ proj.createObjectProjection (builder, inst ->getLoc (), newValue).get ();
273279 }
274280
275- op = li ->getOperand ();
281+ op = inst ->getOperand (0 );
276282
277- // Replace users of the loaded value with `val`
278- // If we have a load [copy], replace the users with copy_value of `val`
279- if (li->getOwnershipQualifier () == LoadOwnershipQualifier::Copy) {
280- li->replaceAllUsesWith (builder.createCopyValue (li->getLoc (), newValue));
283+ if (auto *lbi = dyn_cast<LoadBorrowInst>(inst)) {
284+ if (shouldAddLexicalLifetime (asi)) {
285+ assert (isa<BeginBorrowInst>(newValue));
286+ SmallVector<SILInstruction *, 4 > endBorrows;
287+ for (auto *ebi : lbi->getUsersOfType <EndBorrowInst>()) {
288+ endBorrows.push_back (ebi);
289+ }
290+ for (auto *ebi : endBorrows) {
291+ ebi->eraseFromParent ();
292+ }
293+ lbi->replaceAllUsesWith (newValue);
294+ }
295+ else {
296+ auto *borrow = SILBuilderWithScope (lbi, ctx).createBeginBorrow (
297+ lbi->getLoc (), newValue, asi->isLexical ());
298+ lbi->replaceAllUsesWith (borrow);
299+ }
281300 } else {
282- assert (!asi->getFunction ()->hasOwnership () ||
283- newValue.getOwnershipKind () != OwnershipKind::Guaranteed);
284- li->replaceAllUsesWith (newValue);
301+ auto *li = cast<LoadInst>(inst);
302+ // Replace users of the loaded value with `newValue`
303+ // If we have a load [copy], replace the users with copy_value of `newValue`
304+ if (li->getOwnershipQualifier () == LoadOwnershipQualifier::Copy) {
305+ li->replaceAllUsesWith (builder.createCopyValue (li->getLoc (), newValue));
306+ } else {
307+ li->replaceAllUsesWith (newValue);
308+ }
285309 }
286310
287311 // Pop the scope so that we emit cleanups.
288312 std::move (scope).popAtEndOfScope (&*builder.getInsertionPoint ());
289313
290314 // Delete the load
291- prepareForDeletion (li , instructionsToDelete);
315+ prepareForDeletion (inst , instructionsToDelete);
292316
293317 while (op != asi && op->use_empty ()) {
294318 assert (isa<UncheckedAddrCastInst>(op) || isa<StructElementAddrInst>(op) ||
@@ -577,6 +601,20 @@ StoreInst *StackAllocationPromoter::promoteAllocationInBlock(
577601
578602 if (isLoadFromStack (inst, asi)) {
579603 assert (!runningVals || runningVals->isStorageValid );
604+ if (auto *lbi = dyn_cast<LoadBorrowInst>(inst)) {
605+ if (runningVals) {
606+ if (shouldAddLexicalLifetime (asi)) {
607+ replaceLoad (lbi, runningVals->value .borrow , asi, ctx,
608+ deleter, instructionsToDelete);
609+ }
610+ else {
611+ replaceLoad (lbi, runningVals->value .replacement (asi), asi, ctx,
612+ deleter, instructionsToDelete);
613+ }
614+ ++NumInstRemoved;
615+ }
616+ continue ;
617+ }
580618 auto *li = cast<LoadInst>(inst);
581619 if (li->getOwnershipQualifier () == LoadOwnershipQualifier::Take) {
582620 if (shouldAddLexicalLifetime (asi)) {
@@ -597,12 +635,13 @@ StoreInst *StackAllocationPromoter::promoteAllocationInBlock(
597635 if (lastStoreInst)
598636 lastStoreInst->isStorageValid = false ;
599637 }
638+
600639 if (runningVals) {
601640 // If we are loading from the AllocStackInst and we already know the
602641 // content of the Alloca then use it.
603642 LLVM_DEBUG (llvm::dbgs () << " *** Promoting load: " << *li);
604- replaceLoad (li , runningVals->value .replacement (asi), asi, ctx, deleter ,
605- instructionsToDelete);
643+ replaceLoad (inst , runningVals->value .replacement (asi), asi, ctx,
644+ deleter, instructionsToDelete);
606645 ++NumInstRemoved;
607646 } else if (li->getOperand () == asi &&
608647 li->getOwnershipQualifier () != LoadOwnershipQualifier::Copy) {
@@ -789,10 +828,8 @@ StackAllocationPromoter::getLiveOutValues(BlockSetVector &phiBlocks,
789828 SILValue borrow = SILValue ();
790829 SILValue copy = SILValue ();
791830 if (shouldAddLexicalLifetime (asi)) {
792- auto *bbi = cast<BeginBorrowInst>(&*std::next (si->getIterator ()));
793- borrow = bbi;
794- auto *cvi = cast<CopyValueInst>(bbi->getNextInstruction ());
795- copy = cvi;
831+ borrow = cast<BeginBorrowInst>(&*std::next (si->getIterator ()));
832+ copy = cast<CopyValueInst>(borrow->getNextInstruction ());
796833 }
797834 LiveValues values = {stored, borrow, copy};
798835 return values;
@@ -953,7 +990,7 @@ void StackAllocationPromoter::propagateLiveness(
953990void StackAllocationPromoter::fixBranchesAndUses (BlockSetVector &phiBlocks,
954991 BlockSetVector &phiBlocksOut) {
955992 // First update uses of the value.
956- SmallVector<LoadInst *, 4 > collectedLoads;
993+ SmallVector<SILInstruction *, 4 > collectedLoads;
957994
958995 for (auto ui = asi->use_begin (), ue = asi->use_end (); ui != ue;) {
959996 auto *user = ui->getUser ();
@@ -1421,6 +1458,13 @@ static bool isAddressForLoad(SILInstruction *load, SILBasicBlock *&singleBlock,
14211458 return true ;
14221459 }
14231460
1461+ if (isa<LoadBorrowInst>(load)) {
1462+ if (involvesUntakableProjection) {
1463+ return false ;
1464+ }
1465+ return true ;
1466+ }
1467+
14241468 if (!isa<UncheckedAddrCastInst>(load) && !isa<StructElementAddrInst>(load) &&
14251469 !isa<TupleElementAddrInst>(load))
14261470 return false ;
@@ -1480,8 +1524,8 @@ static bool isDeadAddrProjection(SILInstruction *inst) {
14801524}
14811525
14821526// / Returns true if this AllocStacks is captured.
1483- // / Sets \p inSingleBlock to true if all uses of \p ASI are in a single block.
1484- static bool isCaptured (AllocStackInst *asi, bool & inSingleBlock) {
1527+ // / Sets \p inSingleBlock to true if all uses of \p asi are in a single block.
1528+ static bool isCaptured (AllocStackInst *asi, bool * inSingleBlock) {
14851529 SILBasicBlock *singleBlock = asi->getParent ();
14861530
14871531 // For all users of the AllocStack instruction.
@@ -1518,7 +1562,7 @@ static bool isCaptured(AllocStackInst *asi, bool &inSingleBlock) {
15181562 }
15191563
15201564 // None of the users capture the AllocStack.
1521- inSingleBlock = (singleBlock != nullptr );
1565+ * inSingleBlock = (singleBlock != nullptr );
15221566 return false ;
15231567}
15241568
@@ -1533,6 +1577,10 @@ bool MemoryToRegisters::isWriteOnlyAllocation(AllocStackInst *asi) {
15331577 if (!isa<AllocStackInst>(si->getSrc ()))
15341578 continue ;
15351579
1580+ if (auto *sbi = dyn_cast<StoreBorrowInst>(user))
1581+ if (!isa<AllocStackInst>(sbi->getSrc ()))
1582+ continue ;
1583+
15361584 // Deallocation is also okay.
15371585 if (isa<DeallocStackInst>(user))
15381586 continue ;
@@ -1557,7 +1605,6 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *asi) {
15571605 LLVM_DEBUG (llvm::dbgs () << " *** Promoting in-block: " << *asi);
15581606
15591607 SILBasicBlock *parentBlock = asi->getParent ();
1560-
15611608 // The default value of the AllocStack is NULL because we don't have
15621609 // uninitialized variables in Swift.
15631610 Optional<StorageStateTracking<LiveValues>> runningVals;
@@ -1580,8 +1627,9 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *asi) {
15801627 /* isStorageValid=*/ true };
15811628 }
15821629 assert (runningVals && runningVals->isStorageValid );
1583- if (cast<LoadInst>(inst)->getOwnershipQualifier () ==
1584- LoadOwnershipQualifier::Take) {
1630+ auto *loadInst = dyn_cast<LoadInst>(inst);
1631+ if (loadInst &&
1632+ loadInst->getOwnershipQualifier () == LoadOwnershipQualifier::Take) {
15851633 if (shouldAddLexicalLifetime (asi)) {
15861634 // End the lexical lifetime at a load [take]. The storage is no
15871635 // longer keeping the value alive.
@@ -1590,8 +1638,8 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *asi) {
15901638 }
15911639 runningVals->isStorageValid = false ;
15921640 }
1593- replaceLoad (cast<LoadInst>( inst) , runningVals->value .replacement (asi),
1594- asi, ctx, deleter, instructionsToDelete);
1641+ replaceLoad (inst, runningVals->value .replacement (asi), asi, ctx, deleter ,
1642+ instructionsToDelete);
15951643 ++NumInstRemoved;
15961644 continue ;
15971645 }
@@ -1656,7 +1704,7 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *asi) {
16561704 // Remove deallocation.
16571705 if (auto *dsi = dyn_cast<DeallocStackInst>(inst)) {
16581706 if (dsi->getOperand () == asi) {
1659- deleter.forceDelete (inst );
1707+ deleter.forceDelete (dsi );
16601708 NumInstRemoved++;
16611709 // No need to continue scanning after deallocation.
16621710 break ;
@@ -1741,7 +1789,7 @@ bool MemoryToRegisters::promoteSingleAllocation(AllocStackInst *alloc) {
17411789
17421790 // Don't handle captured AllocStacks.
17431791 bool inSingleBlock = false ;
1744- if (isCaptured (alloc, inSingleBlock)) {
1792+ if (isCaptured (alloc, & inSingleBlock)) {
17451793 ++NumAllocStackCaptured;
17461794 return false ;
17471795 }
0 commit comments