Merge pull request swiftlang#8825 from gottesmm/move_fixup_past_capture_promotion

[semantic-sil] Update capture promotion for moving mark_uninitialized onto alloc_box instead of project_box and move the fixup pass behind it.

Author: gottesmm

lib/SILOptimizer/IPO/CapturePromotion.cpp

@@ -991,18 +991,7 @@ static bool isProjectBoxNonEscapingUse(ProjectBoxInst *PBI,
                                        EscapeMutationScanningState &State) {
   DEBUG(llvm::dbgs() << "    Found project box: " << *PBI);
 
-  // Check for mutations of the address component.
-  SILValue Addr = PBI;
-
-  // If the AllocBox is used by a mark_uninitialized, scan the MUI for
-  // interesting uses.
-  if (Addr->hasOneUse()) {
-    SILInstruction *SingleAddrUser = Addr->use_begin()->getUser();
-    if (isa<MarkUninitializedInst>(SingleAddrUser))
-      Addr = SILValue(SingleAddrUser);
-  }
-
-  for (Operand *AddrOp : Addr->getUses()) {
+  for (Operand *AddrOp : PBI->getUses()) {
     if (!isNonEscapingUse(AddrOp, State)) {
       DEBUG(llvm::dbgs() << "    FAIL! Has escaping user of addr:"
                          << *AddrOp->getUser());
@@ -1015,11 +1004,13 @@ static bool isProjectBoxNonEscapingUse(ProjectBoxInst *PBI,
 
 static bool scanUsesForEscapesAndMutations(Operand *Op,
                                            EscapeMutationScanningState &State) {
-  if (auto *PAI = dyn_cast<PartialApplyInst>(Op->getUser())) {
+  SILInstruction *User = Op->getUser();
+
+  if (auto *PAI = dyn_cast<PartialApplyInst>(User)) {
     return isPartialApplyNonEscapingUser(Op, PAI, State);
   }
 
-  if (auto *PBI = dyn_cast<ProjectBoxInst>(Op->getUser())) {
+  if (auto *PBI = dyn_cast<ProjectBoxInst>(User)) {
     // It is assumed in later code that we will only have 1 project_box. This
     // can be seen since there is no code for reasoning about multiple
     // boxes. Just put in the restriction so we are consistent.
@@ -1029,20 +1020,16 @@ static bool scanUsesForEscapesAndMutations(Operand *Op,
     return isProjectBoxNonEscapingUse(PBI, State);
   }
 
-  // Given a top level copy value use, check all of its user operands as if
-  // they were apart of the use list of the base operand.
-  //
-  // This is because even though we are copying the box, a copy of the box is
-  // just a retain + bitwise copy of the pointer. This has nothing to do with
-  // whether or not the address escapes in some way.
+  // Given a top level copy value use or mark_uninitialized, check all of its
+  // user operands as if they were apart of the use list of the base operand.
   //
   // This is a separate code path from the non escaping user visitor check since
   // 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 (auto *CVI = dyn_cast<CopyValueInst>(Op->getUser())) {
-    return all_of(CVI->getUses(), [&State](Operand *CopyOp) -> bool {
-      return scanUsesForEscapesAndMutations(CopyOp, State);
+  if (isa<CopyValueInst>(User) || isa<MarkUninitializedInst>(User)) {
+    return all_of(User->getUses(), [&State](Operand *UserOp) -> bool {
+      return scanUsesForEscapesAndMutations(UserOp, State);
     });
   }
 
@@ -1171,36 +1158,18 @@ static SILValue getOrCreateProjectBoxHelper(SILValue PartialOperand) {
     return B.createProjectBox(CVI->getLoc(), CVI, 0);
   }
 
-  // Otherwise, handle the alloc_box case.
-  auto *ABI = cast<AllocBoxInst>(PartialOperand);
-  // Load and copy from the address value, passing the result as an argument
-  // to the new closure.
-  //
-  // *NOTE* This code assumes that we only have one project box user. We
-  // enforce this with the assert below.
-  assert(count_if(ABI->getUses(), [](Operand *Op) -> bool {
-           return isa<ProjectBoxInst>(Op->getUser());
-         }) == 1);
-
-  // If the address is marked uninitialized, load through the mark, so
-  // that DI can reason about it.
-  for (Operand *BoxValueUse : ABI->getUses()) {
-    auto *PBI = dyn_cast<ProjectBoxInst>(BoxValueUse->getUser());
-    if (!PBI)
-      continue;
-
-    auto *OptIter = PBI->getSingleUse();
-    if (!OptIter)
-      continue;
-
-    if (!isa<MarkUninitializedInst>(OptIter->getUser()))
-      continue;
-
-    return OptIter->getUser();
+  // Otherwise, handle the alloc_box case. If we have a mark_uninitialized on
+  // the box, we create the project value through that.
+  SILInstruction *Box = cast<AllocBoxInst>(PartialOperand);
+  if (auto *Op = Box->getSingleUse()) {
+    if (auto *MUI = dyn_cast<MarkUninitializedInst>(Op->getUser())) {
+      Box = MUI;
+    }
   }
 
-  // Otherwise, just return a project_box.
-  return getOrCreateProjectBox(ABI, 0);
+  // Just return a project_box.
+  SILBuilder B(std::next(Box->getIterator()));
+  return B.createProjectBox(Box->getLoc(), Box, 0);
 }
 
 /// \brief Given a partial_apply instruction and a set of promotable indices,

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -74,13 +74,13 @@ static void addOwnershipModelEliminatorPipeline(SILPassPipelinePlan &P) {
 static void addMandatoryOptPipeline(SILPassPipelinePlan &P,
                                     const SILOptions &Options) {
   P.startPipeline("Guaranteed Passes");
-  P.addMarkUninitializedFixup();
   if (Options.EnableMandatorySemanticARCOpts) {
     P.addSemanticARCOpts();
   }
   P.addDiagnoseStaticExclusivity();
   P.addCapturePromotion();
   P.addAllocBoxToStack();
+  P.addMarkUninitializedFixup();
   P.addNoReturnFolding();
   P.addOwnershipModelEliminator();
   P.addDefiniteInitialization();

lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

@@ -109,20 +109,21 @@ static bool addLastRelease(SILValue V, SILBasicBlock *BB,
 // Find the final releases of the alloc_box along any given path.
 // These can include paths from a release back to the alloc_box in a
 // loop.
-static bool getFinalReleases(AllocBoxInst *ABI,
-                             llvm::SmallVectorImpl<SILInstruction*> &Releases) {
+static bool
+getFinalReleases(SILValue Box,
+                 llvm::SmallVectorImpl<SILInstruction *> &Releases) {
   llvm::SmallPtrSet<SILBasicBlock*, 16> LiveIn;
   llvm::SmallPtrSet<SILBasicBlock*, 16> UseBlocks;
 
-  auto *DefBB = ABI->getParent();
+  auto *DefBB = Box->getParentBlock();
 
   auto seenRelease = false;
   SILInstruction *OneRelease = nullptr;
 
   // We'll treat this like a liveness problem where the alloc_box is
   // the def. Each block that has a use of the owning pointer has the
   // value live-in unless it is the block with the alloc_box.
-  llvm::SmallVector<Operand *, 32> Worklist(ABI->use_begin(), ABI->use_end());
+  llvm::SmallVector<Operand *, 32> Worklist(Box->use_begin(), Box->use_end());
   while (!Worklist.empty()) {
     auto *Op = Worklist.pop_back_val();
     auto *User = Op->getUser();
@@ -137,12 +138,10 @@ static bool getFinalReleases(AllocBoxInst *ABI,
     // Also keep track of the blocks with uses.
     UseBlocks.insert(BB);
 
-    // If we have a copy value, add its uses to the work list and continue.
-    //
-    // We are actually performing multiple operations here. We are eliminating
-    // copies of the box and the box itself.
-    if (auto *CVI = dyn_cast<CopyValueInst>(User)) {
-      std::copy(CVI->use_begin(), CVI->use_end(), std::back_inserter(Worklist));
+    // If we have a copy value or a mark_uninitialized, add its uses to the work
+    // list and continue.
+    if (isa<MarkUninitializedInst>(User) || isa<CopyValueInst>(User)) {
+      copy(User->getUses(), std::back_inserter(Worklist));
       continue;
     }
 
@@ -171,7 +170,7 @@ static bool getFinalReleases(AllocBoxInst *ABI,
   // release/dealloc.
   for (auto *BB : UseBlocks)
     if (!successorHasLiveIn(BB, LiveIn))
-      if (!addLastRelease(ABI, BB, Releases))
+      if (!addLastRelease(Box, BB, Releases))
         return false;
 
   return true;
@@ -213,9 +212,8 @@ static bool partialApplyEscapes(SILValue V, bool examineApply) {
     // If we have a copy_value, the copy value does not cause an escape, but its
     // uses might do so... so add the copy_value's uses to the worklist and
     // continue.
-    if (auto *Copy = dyn_cast<CopyValueInst>(User)) {
-      std::copy(Copy->use_begin(), Copy->use_end(),
-                std::back_inserter(Worklist));
+    if (isa<CopyValueInst>(User)) {
+      copy(User->getUses(), std::back_inserter(Worklist));
       continue;
     }
 
@@ -353,11 +351,9 @@ static bool checkPartialApplyBody(Operand *O) {
                                PromotedOperands);
 }
 
-
-/// findUnexpectedBoxUse - Validate that the uses of a pointer to a
-/// box do not eliminate it from consideration for promotion to a
-/// stack element. Optionally examine the body of partial_apply
-/// to see if there is an unexpected use inside.  Return the
+/// Validate that the uses of a pointer to a box do not eliminate it from
+/// consideration for promotion to a stack element. Optionally examine the body
+/// of partial_apply to see if there is an unexpected use inside.  Return the
 /// instruction with the unexpected use if we find one.
 static SILInstruction* findUnexpectedBoxUse(SILValue Box,
                                             bool examinePartialApply,
@@ -386,9 +382,10 @@ static SILInstruction* findUnexpectedBoxUse(SILValue Box,
         (!inAppliedFunction && isa<DeallocBoxInst>(User)))
       continue;
 
-    // If we have a copy_value, visit the users of the copy_value.
-    if (auto *CVI = dyn_cast<CopyValueInst>(User)) {
-      std::copy(CVI->use_begin(), CVI->use_end(), std::back_inserter(Worklist));
+    // If our user instruction is a copy_value or a marked_uninitialized, visit
+    // the users recursively.
+    if (isa<MarkUninitializedInst>(User) || isa<CopyValueInst>(User)) {
+      copy(User->getUses(), std::back_inserter(Worklist));
       continue;
     }
 
@@ -432,19 +429,20 @@ static bool canPromoteAllocBox(AllocBoxInst *ABI,
   return true;
 }
 
-static void replaceProjectBoxUsers(AllocBoxInst *ABI, AllocStackInst *ASI) {
-  llvm::SmallVector<Operand *, 8> Worklist(ABI->use_begin(), ABI->use_end());
+static void replaceProjectBoxUsers(SILValue HeapBox, SILValue StackBox) {
+  llvm::SmallVector<Operand *, 8> Worklist(HeapBox->use_begin(),
+                                           HeapBox->use_end());
   while (!Worklist.empty()) {
     auto *Op = Worklist.pop_back_val();
     if (auto *PBI = dyn_cast<ProjectBoxInst>(Op->getUser())) {
-      PBI->replaceAllUsesWith(ASI);
+      PBI->replaceAllUsesWith(StackBox);
       continue;
     }
 
     auto *CVI = dyn_cast<CopyValueInst>(Op->getUser());
     if (!CVI)
       continue;
-    std::copy(CVI->use_begin(), CVI->use_end(), std::back_inserter(Worklist));
+    copy(CVI->getUses(), std::back_inserter(Worklist));
   }
 }
 
@@ -453,36 +451,40 @@ static void replaceProjectBoxUsers(AllocBoxInst *ABI, AllocStackInst *ASI) {
 static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI) {
   DEBUG(llvm::dbgs() << "*** Promoting alloc_box to stack: " << *ABI);
 
+  SILValue HeapBox = ABI;
+  Optional<MarkUninitializedInst::Kind> Kind;
+  if (HeapBox->hasOneUse()) {
+    auto *User = HeapBox->getSingleUse()->getUser();
+    if (auto *MUI = dyn_cast<MarkUninitializedInst>(User)) {
+      HeapBox = MUI;
+      Kind = MUI->getKind();
+    }
+  }
+
   llvm::SmallVector<SILInstruction *, 4> FinalReleases;
-  if (!getFinalReleases(ABI, FinalReleases))
+  if (!getFinalReleases(HeapBox, FinalReleases))
     return false;
 
   // Promote this alloc_box to an alloc_stack. Insert the alloc_stack
   // at the beginning of the function.
-  SILBuilder BuildAlloc(ABI);
-  BuildAlloc.setCurrentDebugScope(ABI->getDebugScope());
+  SILBuilder Builder(ABI);
+  Builder.setCurrentDebugScope(ABI->getDebugScope());
   assert(ABI->getBoxType()->getLayout()->getFields().size() == 1
          && "rewriting multi-field box not implemented");
-  auto *ASI = BuildAlloc.createAllocStack(ABI->getLoc(),
-                          ABI->getBoxType()->getFieldType(ABI->getModule(), 0),
-                          ABI->getVarInfo());
+  auto *ASI = Builder.createAllocStack(
+      ABI->getLoc(), ABI->getBoxType()->getFieldType(ABI->getModule(), 0),
+      ABI->getVarInfo());
+
+  // Transfer a mark_uninitialized if we have one.
+  SILValue StackBox = ASI;
+  if (Kind) {
+    StackBox =
+        Builder.createMarkUninitialized(ASI->getLoc(), ASI, Kind.getValue());
+  }
 
   // Replace all uses of the address of the box's contained value with
   // the address of the stack location.
-  replaceProjectBoxUsers(ABI, ASI);
-
-  // Check to see if the alloc_box was used by a mark_uninitialized instruction.
-  // If so, any uses of the pointer result need to keep using the MUI, not the
-  // alloc_stack directly.  If we don't do this, DI will miss the uses.
-  SILValue PointerResult = ASI;
-  for (auto UI : ASI->getUses()) {
-    if (auto *MUI = dyn_cast<MarkUninitializedInst>(UI->getUser())) {
-      assert(ASI->hasOneUse() &&
-             "alloc_stack used by mark_uninitialized, but not exclusively!");
-      PointerResult = MUI;
-      break;
-    }
-  }
+  replaceProjectBoxUsers(HeapBox, StackBox);
 
   assert(ABI->getBoxType()->getLayout()->getFields().size() == 1
          && "promoting multi-field box not implemented");
@@ -495,7 +497,7 @@ static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI) {
     if (!isa<DeallocBoxInst>(LastRelease)&& !Lowering.isTrivial()) {
       // For non-trivial types, insert destroys for each final release-like
       // instruction we found that isn't an explicit dealloc_box.
-      Builder.emitDestroyAddrAndFold(Loc, PointerResult);
+      Builder.emitDestroyAddrAndFold(Loc, StackBox);
     }
     Builder.createDeallocStack(Loc, ASI);
   }
@@ -509,10 +511,10 @@ static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI) {
   while (!Worklist.empty()) {
     auto *User = Worklist.pop_back_val();
 
-    if (isa<CopyValueInst>(User)) {
-      std::transform(
-          User->use_begin(), User->use_end(), std::back_inserter(Worklist),
-          [](Operand *Op) -> SILInstruction * { return Op->getUser(); });
+    // Look through any mark_uninitialized, copy_values.
+    if (isa<MarkUninitializedInst>(User) || isa<CopyValueInst>(User)) {
+      transform(User->getUses(), std::back_inserter(Worklist),
+                [](Operand *Op) -> SILInstruction * { return Op->getUser(); });
       User->replaceAllUsesWith(
           SILUndef::get(User->getType(), User->getModule()));
       User->eraseFromParent();