Merge pull request #40793 from nate-chandler/silgen/emit-borrow-scopes-for-alloc_boxes

[SIL] Added borrow scopes to alloc_boxes.

Author: nate-chandler

lib/SIL/IR/SILBuilder.cpp

@@ -285,6 +285,10 @@ SILBuilder::emitDestroyAddr(SILLocation Loc, SILValue Operand) {
     if (isa<DeallocStackInst>(Inst))
       continue;
 
+    // end_borrow insts also don't affect take-ability
+    if (isa<EndBorrowInst>(Inst))
+      continue;
+
     // An end_access of the same address may be able to be rewritten as a
     // [deinit] access.
     if (auto endAccess = dyn_cast<EndAccessInst>(Inst)) {

lib/SILGen/ResultPlan.cpp

@@ -66,7 +66,13 @@ class IndirectOpenedSelfCleanup final : public Cleanup {
   void emit(SILGenFunction &SGF, CleanupLocation loc, ForUnwind_t forUnwind)
   override {
     assert(box && "buffer never emitted before activating cleanup?!");
-    SGF.B.createDeallocBox(loc, box);
+    auto theBox = box;
+    if (SGF.getASTContext().SILOpts.supportsLexicalLifetimes(SGF.getModule())) {
+      auto *bbi = cast<BeginBorrowInst>(theBox);
+      SGF.B.createEndBorrow(loc, bbi);
+      theBox = bbi->getOperand();
+    }
+    SGF.B.createDeallocBox(loc, theBox);
   }
 
   void dump(SILGenFunction &SGF) const override {
@@ -160,7 +166,10 @@ class IndirectOpenedSelfResultPlan final : public ResultPlan {
       SILBoxType::get(SGF.getASTContext(),
                       boxLayout,
                       layoutSubs));
-    
+    if (SGF.getASTContext().SILOpts.supportsLexicalLifetimes(SGF.getModule())) {
+      resultBox = SGF.B.createBeginBorrow(loc, resultBox, /*isLexical=*/true);
+    }
+
     // Complete the cleanup to deallocate this buffer later, after we're
     // finished with the argument.
     static_cast<IndirectOpenedSelfCleanup&>(SGF.Cleanups.getCleanup(handle))

lib/SILGen/SILGenApply.cpp

@@ -3505,7 +3505,13 @@ class DeallocateUninitializedBox : public Cleanup {
   DeallocateUninitializedBox(SILValue box) : box(box) {}
 
   void emit(SILGenFunction &SGF, CleanupLocation l, ForUnwind_t forUnwind) override {
-    SGF.B.createDeallocBox(l, box);
+    auto theBox = box;
+    if (SGF.getASTContext().SILOpts.supportsLexicalLifetimes(SGF.getModule())) {
+      auto *bbi = cast<BeginBorrowInst>(theBox);
+      SGF.B.createEndBorrow(l, bbi);
+      theBox = bbi->getOperand();
+    }
+    SGF.B.createDeallocBox(l, theBox);
   }
 
   void dump(SILGenFunction &SGF) const override {

lib/SILGen/SILGenDecl.cpp

@@ -297,7 +297,13 @@ class DeallocateUninitializedLocalVariable : public Cleanup {
 
   void emit(SILGenFunction &SGF, CleanupLocation l,
             ForUnwind_t forUnwind) override {
-    SGF.B.createDeallocBox(l, Box);
+    auto box = Box;
+    if (SGF.getASTContext().SILOpts.supportsLexicalLifetimes(SGF.getModule())) {
+      auto *bbi = cast<BeginBorrowInst>(box);
+      SGF.B.createEndBorrow(l, bbi);
+      box = bbi->getOperand();
+    }
+    SGF.B.createDeallocBox(l, box);
   }
 
   void dump(SILGenFunction &) const override {
@@ -360,6 +366,10 @@ class LocalVariableInitialization : public SingleBufferInitialization {
     if (kind)
       Box = SGF.B.createMarkUninitialized(decl, Box, kind.getValue());
 
+    if (SGF.getASTContext().SILOpts.supportsLexicalLifetimes(SGF.getModule())) {
+      Box = SGF.B.createBeginBorrow(decl, Box, /*isLexical=*/true);
+    }
+
     Addr = SGF.B.createProjectBox(decl, Box, 0);
 
     // Push a cleanup to destroy the local variable.  This has to be
@@ -1751,7 +1761,15 @@ void SILGenFunction::destroyLocalVariable(SILLocation silLoc, VarDecl *vd) {
   // For a heap variable, the box is responsible for the value. We just need
   // to give up our retain count on it.
   if (loc.box) {
-    B.emitDestroyValueOperation(silLoc, loc.box);
+    if (!getASTContext().SILOpts.supportsLexicalLifetimes(getModule())) {
+      B.emitDestroyValueOperation(silLoc, loc.box);
+      return;
+    }
+
+    auto *bbi = cast<BeginBorrowInst>(loc.box);
+    B.createEndBorrow(silLoc, bbi);
+    B.emitDestroyValueOperation(silLoc, bbi->getOperand());
+
     return;
   }
 

lib/SILGen/Scope.cpp

@@ -57,6 +57,9 @@ static void lifetimeExtendAddressOnlyRValueSubValues(
            "addresses must be address only.");
     auto boxTy = SILBoxType::get(v->getType().getASTType());
     SILValue box = SGF.B.createAllocBox(loc, boxTy);
+    if (SGF.getASTContext().SILOpts.supportsLexicalLifetimes(SGF.getModule())) {
+      box = SGF.B.createBeginBorrow(loc, box, /*isLexical=*/true);
+    }
     SILValue addr = SGF.B.createProjectBox(loc, box, 0);
     SGF.B.createCopyAddr(loc, v, addr, IsTake, IsInitialization);
 

lib/SILOptimizer/Mandatory/AccessEnforcementSelection.cpp

@@ -227,6 +227,11 @@ void SelectEnforcement::analyzeUsesOfBox(SingleValueInstruction *source) {
   for (auto use : source->getUses()) {
     auto user = use->getUser();
 
+    if (auto BBI = dyn_cast<BeginBorrowInst>(user)) {
+      analyzeUsesOfBox(BBI);
+      continue;
+    }
+
     if (auto MUI = dyn_cast<MarkUninitializedInst>(user)) {
       analyzeUsesOfBox(MUI);
       continue;
@@ -238,10 +243,9 @@ void SelectEnforcement::analyzeUsesOfBox(SingleValueInstruction *source) {
     }
 
     // Ignore certain other uses that do not capture the value.
-    if (isa<StrongRetainInst>(user) ||
-        isa<StrongReleaseInst>(user) ||
-        isa<DestroyValueInst>(user) ||
-        isa<DeallocBoxInst>(user))
+    if (isa<StrongRetainInst>(user) || isa<StrongReleaseInst>(user) ||
+        isa<DestroyValueInst>(user) || isa<DeallocBoxInst>(user) ||
+        isa<EndBorrowInst>(user))
       continue;
 
     // Treat everything else as an escape.
@@ -667,6 +671,8 @@ void AccessEnforcementSelection::processFunction(SILFunction *F) {
 SourceAccess
 AccessEnforcementSelection::getAccessKindForBox(ProjectBoxInst *projection) {
   SILValue source = projection->getOperand();
+  if (auto *BBI = dyn_cast<BeginBorrowInst>(source))
+    source = BBI->getOperand();
   if (auto *MUI = dyn_cast<MarkUninitializedInst>(source))
     source = MUI->getOperand();
 

lib/SILOptimizer/Mandatory/CapturePromotion.cpp

@@ -1040,6 +1040,7 @@ class NonEscapingUserVisitor
   ALWAYS_NON_ESCAPING_INST(Load)
   ALWAYS_NON_ESCAPING_INST(StrongRelease)
   ALWAYS_NON_ESCAPING_INST(DestroyValue)
+  ALWAYS_NON_ESCAPING_INST(EndBorrow)
 #undef ALWAYS_NON_ESCAPING_INST
 
   bool visitDeallocBoxInst(DeallocBoxInst *dbi) {
@@ -1213,7 +1214,8 @@ static bool findEscapeOrMutationUses(Operand *op,
   // we want to be more conservative around non-top level copies (i.e. a copy
   // derived from a projection like instruction). In fact such a thing may not
   // even make any sense!
-  if (isa<CopyValueInst>(user) || isa<MarkUninitializedInst>(user)) {
+  if (isa<CopyValueInst>(user) || isa<MarkUninitializedInst>(user) ||
+      isa<BeginBorrowInst>(user)) {
     bool foundSomeMutations = false;
     for (auto *use : cast<SingleValueInstruction>(user)->getUses()) {
       foundSomeMutations |= findEscapeOrMutationUses(use, state);

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp

@@ -51,7 +51,11 @@ static void gatherDestroysOfContainer(const DIMemoryObjectInfo &memoryInfo,
   // TODO: This should really be tracked separately from other destroys so that
   // we distinguish the lifetime of the container from the value itself.
   assert(isa<ProjectBoxInst>(uninitMemory));
-  auto *mui = cast<MarkUninitializedInst>(uninitMemory->getOperand(0));
+  auto value = uninitMemory->getOperand(0);
+  if (auto *bbi = dyn_cast<BeginBorrowInst>(value)) {
+    value = bbi->getOperand();
+  }
+  auto *mui = cast<MarkUninitializedInst>(value);
   for (auto *user : mui->getUsersOfType<DestroyValueInst>()) {
     useInfo.trackDestroy(user);
   }
@@ -114,6 +118,12 @@ DIMemoryObjectInfo::DIMemoryObjectInfo(MarkUninitializedInst *MI)
   auto &Module = MI->getModule();
 
   SILValue Address = MemoryInst;
+  if (auto BBI = MemoryInst->getSingleUserOfType<BeginBorrowInst>()) {
+    if (auto PBI = BBI->getSingleUserOfType<ProjectBoxInst>()) {
+      IsBox = true;
+      Address = PBI;
+    }
+  }
   if (auto PBI = MemoryInst->getSingleUserOfType<ProjectBoxInst>()) {
     IsBox = true;
     Address = PBI;

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.h

@@ -102,10 +102,14 @@ class DIMemoryObjectInfo {
   /// with the memory info.
   SingleValueInstruction *getUninitializedValue() const {
     if (IsBox) {
+      SILValue inst = MemoryInst;
+      if (auto *bbi = MemoryInst->getSingleUserOfType<BeginBorrowInst>()) {
+        inst = bbi;
+      }
       // TODO: consider just storing the ProjectBoxInst in this case.
-      auto *pbi = MemoryInst->getSingleUserOfType<ProjectBoxInst>();
-      assert(pbi);
-      return pbi;
+      SingleValueInstruction *svi = inst->getSingleUserOfType<ProjectBoxInst>();
+      assert(svi);
+      return svi;
     }
     return MemoryInst;
   }

lib/SILOptimizer/Mandatory/DiagnoseInvalidEscapingCaptures.cpp

@@ -157,7 +157,11 @@ const ParamDecl *getParamDeclFromOperand(SILValue value) {
 
 bool isUseOfSelfInInitializer(Operand *oper) {
   if (auto *PBI = dyn_cast<ProjectBoxInst>(oper->get())) {
-    if (auto *MUI = dyn_cast<MarkUninitializedInst>(PBI->getOperand())) {
+    SILValue value = PBI->getOperand();
+    if (auto *bbi = dyn_cast<BeginBorrowInst>(value)) {
+      value = bbi->getOperand();
+    }
+    if (auto *MUI = dyn_cast<MarkUninitializedInst>(value)) {
       switch (MUI->getMarkUninitializedKind()) {
       case MarkUninitializedInst::Kind::Var:
         return false;