[AddressLowering] Shorten stack lifetimes.

During storage allocation, create all alloc_stacks at the front of the
function entry block, except those for opened existential types.  Do not
create any dealloc_stacks, even for opened existential types.

Then, after function rewriting, position the alloc_stacks according to
uses, except those for opened existential types.  This means allocating
in the block that's the least common ancestor of all uses. At this
point, create dealloc_stacks on the dominance frontier of the
alloc_stacks, even for allocations for opened existential types.

The behavior for opened existential types diverges from all others in
order to maintain SIL validity (as much as possible) while the pass
runs.  It would be invalid to create the alloc_stacks for opened
existential types at the front of the entry block because the type which
is opened is not yet available, but that type's opening must dominate
the corresponding alloc_stack.

Author: nate-chandler

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -1002,16 +1002,28 @@ namespace {
 /// function in postorder. If the definition is an argument of this function,
 /// simply replace the function argument with an address representing the
 /// caller's storage.
-///
-/// TODO: shrink lifetimes by inserting alloc_stack at the dominance LCA and
-/// finding the lifetime boundary with a simple backward walk from uses.
 class OpaqueStorageAllocation {
   AddressLoweringState &pass;
+  /// The alloc_stacks that have been created which eventually need to have
+  /// corresponding dealloc_stacks created.
+  ///
+  /// Supports erasure because created alloc_stacks may be erased when block
+  /// arguments are coalesced.
+  SmallBlotSetVector<AllocStackInst *, 16> allocs;
+  /// The alloc_stacks that have been created which eventually need to be
+  /// positioned appropriately.
+  ///
+  /// The subset of allocs which aren't for opened existentials.
+  InstructionSet allocsToReposition;
 
 public:
-  explicit OpaqueStorageAllocation(AddressLoweringState &pass) : pass(pass) {}
+  explicit OpaqueStorageAllocation(AddressLoweringState &pass)
+      : pass(pass), allocsToReposition(pass.function) {}
 
   void allocateOpaqueStorage();
+  /// Position alloc_stacks according to uses and create dealloc_stacks that
+  /// jointly postdominate.
+  void finalizeOpaqueStorage();
 
 protected:
   void allocateValue(SILValue value);
@@ -1037,6 +1049,8 @@ class OpaqueStorageAllocation {
 
   AllocStackInst *createStackAllocation(SILValue value);
 
+  SILBasicBlock *getLeastCommonAncestorOfUses(SILValue value);
+
   void createStackAllocationStorage(SILValue value) {
     pass.valueStorageMap.getStorage(value).storageAddress =
         createStackAllocation(value);
@@ -1226,9 +1240,20 @@ void OpaqueStorageAllocation::removeAllocation(SILValue value) {
   for (Operand *use : uses) {
     pass.deleter.forceDelete(cast<DeallocStackInst>(use->getUser()));
   }
+  allocs.erase(allocInst);
   pass.deleter.forceDelete(allocInst);
 }
 
+SILBasicBlock *
+OpaqueStorageAllocation::getLeastCommonAncestorOfUses(SILValue value) {
+  SILBasicBlock *lca = nullptr;
+  for (auto *use : value->getUses()) {
+    auto *block = use->getParentBlock();
+    lca = lca ? pass.domInfo->findNearestCommonDominator(lca, block) : block;
+  }
+  return lca;
+}
+
 // Create alloc_stack that dominates an owned value \p value. Create
 // jointly-postdominating dealloc_stack instructions.  Nesting will be fixed
 // later.
@@ -1250,8 +1275,8 @@ AllocStackInst *OpaqueStorageAllocation::createStackAllocation(SILValue value) {
 
   // For opened existential types, allocate stack space at the type
   // definition. Allocating as early as possible provides more opportunity for
-  // creating use projections into value. But allocation must be no earlier then
-  // the latest type definition.
+  // creating use projections into value. But allocation must be no earlier
+  // then the latest type definition.
   SILInstruction *latestOpeningInst = nullptr;
   allocTy.getASTType().visit([&](CanType type) {
     auto archetype = dyn_cast<ArchetypeType>(type);
@@ -1274,31 +1299,44 @@ AllocStackInst *OpaqueStorageAllocation::createStackAllocation(SILValue value) {
       latestOpeningInst = openingInst;
     }
   });
-  auto allocPt = latestOpeningInst ? std::next(latestOpeningInst->getIterator())
-                                   : pass.function->begin()->begin();
+
+  auto allocPt = latestOpeningInst
+                     ? latestOpeningInst->getNextInstruction()->getIterator()
+                     : pass.function->getEntryBlock()->front().getIterator();
   auto allocBuilder = pass.getBuilder(allocPt);
   AllocStackInst *alloc = allocBuilder.createAllocStack(pass.genLoc(), allocTy);
+  allocs.insert(alloc);
+  if (latestOpeningInst == nullptr) {
+    allocsToReposition.insert(alloc);
+  }
+  return alloc;
+}
+
+void OpaqueStorageAllocation::finalizeOpaqueStorage() {
+  SmallVector<SILBasicBlock *, 4> boundary;
+  for (auto maybeAlloc : allocs) {
+    // An allocation may be erased when coalescing block arguments.
+    if (!maybeAlloc.hasValue())
+      continue;
+
+    auto *alloc = maybeAlloc.value();
+
+    if (allocsToReposition.contains(alloc)) {
+      auto allocPt = &*getLeastCommonAncestorOfUses(alloc)->begin();
+      alloc->moveBefore(allocPt);
+    }
 
-  auto dealloc = [&](SILBasicBlock::iterator insertPt) {
-    auto deallocBuilder = pass.getBuilder(insertPt);
-    deallocBuilder.createDeallocStack(pass.genLoc(), alloc);
-  };
-  if (latestOpeningInst) {
     // Deallocate at the predecessors of dominance frontier blocks that are
     // dominated by the alloc to ensure that allocation encloses not only the
     // uses of the current value, but also of any values reusing this storage as
     // a use projection.
-    SmallVector<SILBasicBlock *, 4> boundary;
     computeDominatedBoundaryBlocks(alloc->getParent(), pass.domInfo, boundary);
     for (SILBasicBlock *deallocBlock : boundary) {
-      dealloc(deallocBlock->getTerminator()->getIterator());
-    }
-  } else {
-    for (SILInstruction *deallocPoint : pass.exitingInsts) {
-      dealloc(deallocPoint->getIterator());
+      auto deallocBuilder = pass.getBuilder(deallocBlock->back().getIterator());
+      deallocBuilder.createDeallocStack(pass.genLoc(), alloc);
     }
+    boundary.clear();
   }
-  return alloc;
 }
 
 //===----------------------------------------------------------------------===//
@@ -3850,6 +3888,8 @@ void AddressLowering::runOnFunction(SILFunction *function) {
   // forward order, setting 'storageAddress' for each projection as it goes.
   rewriteFunction(pass);
 
+  allocator.finalizeOpaqueStorage();
+
   deleteRewrittenInstructions(pass);
 
   StackNesting::fixNesting(function);

test/SILOptimizer/opaque_values_Onone.swift

@@ -15,9 +15,9 @@ enum Maybe1<T : Equatable> {
   case yep(T)
   // CHECK-LABEL: sil hidden @maybe1_compare {{.*}} {
   // CHECK:         [[LHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe1
-  // CHECK:         [[RHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe1
   // CHECK:         switch_enum_addr [[LHS_ADDR]] : $*Maybe1<T>, case #Maybe1.yep!enumelt: [[L_YEP:bb[0-9]+]], case #Maybe1.nope!enumelt: {{bb[0-9]+}}
   // CHECK:       [[L_YEP]]:
+  // CHECK:         [[RHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe1
   // CHECK:         unchecked_take_enum_data_addr [[LHS_ADDR]] : $*Maybe1<T>, #Maybe1.yep!enumelt
   // CHECK:         switch_enum_addr [[RHS_ADDR]] : $*Maybe1<T>, case #Maybe1.yep!enumelt: [[L_AND_R_YEP:bb[0-9]+]], default {{bb[0-9]+}}
   // CHECK:       [[L_AND_R_YEP]]:
@@ -42,9 +42,9 @@ enum Maybe2<T : Equatable> {
 
   // CHECK-LABEL: sil hidden @maybe2_compare {{.*}} {
   // CHECK:         [[LHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe2<T>
-  // CHECK:         [[RHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe2<T>
   // CHECK:         switch_enum_addr [[LHS_ADDR]] : $*Maybe2<T>, case #Maybe2.yep!enumelt: [[L_YEP:bb[0-9]+]], case #Maybe2.nope!enumelt: {{bb[0-9]+}}
   // CHECK:       [[L_YEP]]:
+  // CHECK:         [[RHS_ADDR:%[^,]+]] = alloc_stack [lexical] $Maybe2<T>
   // CHECK:         unchecked_take_enum_data_addr [[LHS_ADDR]] : $*Maybe2<T>, #Maybe2.yep!enumelt
   // CHECK:         switch_enum_addr [[RHS_ADDR]] : $*Maybe2<T>, case #Maybe2.yep!enumelt: [[R_YEP:bb[0-9]+]], default {{bb[0-9]+}}
   // CHECK:       [[L_AND_R_YEP]]: