AllocBoxToStack: Improve alloc_stack/dealloc_stack scoping

include/swift/SILOptimizer/Utils/StackNesting.h

@@ -0,0 +1,156 @@
+//===--- StackNesting.h - Utility for stack nesting -------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SILOPTIMIZER_UTILS_STACKNESTING_H
+#define SWIFT_SILOPTIMIZER_UTILS_STACKNESTING_H
+
+#include "swift/SIL/SILInstruction.h"
+#include "llvm/ADT/SmallBitVector.h"
+
+#include <vector>
+
+namespace swift {
+
+/// A utility to correct the nesting of stack allocating/deallocating
+/// instructions.
+///
+/// This utility is useful for optimizations which create stack allocations
+/// without caring about correct nesting with existing allocations. This may
+/// result in code like
+/// \code
+///   %1 = alloc_stack
+///     ...
+///   %2 = alloc_stack
+///     ...
+///   dealloc_stack %1
+///     ...
+///   dealloc_stack %2
+/// \endcode
+///
+/// The StackNesting utility is able to correct the code:
+/// \code
+///   %1 = alloc_stack
+///     ...
+///   %2 = alloc_stack
+///     ...
+///   dealloc_stack %2
+///   dealloc_stack %1
+/// \endcode
+///
+class StackNesting {
+
+  typedef llvm::SmallBitVector BitVector;
+
+  /// Data stored for each block (actually for each block which is not dead).
+  struct BlockInfo {
+
+    /// Back-link to the block.
+    SILBasicBlock *Block;
+
+    /// The cached list of successors.
+    llvm::SmallVector<BlockInfo *, 8> Successors;
+
+    /// The list of stack allocating/deallocating instructions in the block.
+    llvm::SmallVector<SILInstruction *, 8> StackInsts;
+
+    /// The bit-set of alive stack locations at the block entry.
+    BitVector AliveStackLocsAtEntry;
+
+    /// True if there is a path from this block to a function-exit.
+    ///
+    /// In other words: this block does not end in an unreachable-instruction.
+    /// This flag is only used for verifying that the lifetime of a stack
+    /// location does not end at the end of a block.
+    bool ExitReachable = false;
+
+    BlockInfo(SILBasicBlock *Block) : Block(Block) { }
+  };
+
+  /// Data stored for each stack location (= allocation).
+  ///
+  /// Each stack location is allocated by a single allocation instruction.
+  struct StackLoc {
+    StackLoc(SILInstruction *Alloc) : Alloc(Alloc) { }
+
+    /// Back-link to the allocation instruction.
+    SILInstruction *Alloc;
+
+    /// Bit-set which represents all alive locations at this allocation.
+    /// It obviously includes this location itself. And it includes all "outer"
+    /// locations which surround this location.
+    BitVector AliveLocs;
+  };
+
+  /// Mapping from stack allocations (= locations) to bit numbers.
+  llvm::DenseMap<SILInstruction *, unsigned> StackLoc2BitNumbers;
+
+  /// The list of stack locations. The index into this array is also the bit
+  /// number in the bit-sets.
+  llvm::SmallVector<StackLoc, 8> StackLocs;
+
+  /// Block data for all (non-dead) blocks.
+  std::vector<BlockInfo> BlockInfos;
+
+public:
+
+  /// The possible return values of correctStackNesting().
+  enum class Changes {
+    /// No changes are made.
+    None,
+
+    /// Only instructions were inserted or deleted.
+    Instructions,
+
+    /// Instructions were inserted or deleted and new blocks were inserted.
+    CFG
+  };
+
+  StackNesting() { }
+
+  /// Performs correction of stack nesting by moving stack-deallocation
+  /// instructions down the control flow.
+  ///
+  /// Returns the status of what changes were made.
+  Changes correctStackNesting(SILFunction *F);
+
+  /// For debug dumping.
+  void dump() const;
+
+  static void dumpBits(const BitVector &Bits);
+
+private:
+  /// Initializes the data structures.
+  void setup(SILFunction *F);
+
+  /// Solves the dataflow problem.
+  ///
+  /// Returns true if there is a nesting of locations in any way, which can
+  /// potentially in the wrong order.
+  bool solve();
+
+  /// Insert deallocation instructions for all locations which are alive before
+  /// the InsertionPoint (AliveBefore) but not alive after the InsertionPoint
+  /// (AliveAfter).
+  ///
+  /// Returns true if any deallocations were inserted.
+  bool insertDeallocs(const BitVector &AliveBefore, const BitVector &AliveAfter,
+                      SILInstruction *InsertionPoint);
+
+  /// Modifies the SIL to end up with a correct stack nesting.
+  ///
+  /// Returns the status of what changes were made.
+  Changes adaptDeallocs();
+};
+
+} // end namespace swift
+
+#endif // SWIFT_SILOPTIMIZER_UTILS_STACKNESTING_H

lib/SILOptimizer/Transforms/AllocBoxToStack.cpp

@@ -19,6 +19,7 @@
 #include "swift/SIL/SILCloner.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/Local.h"
+#include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
@@ -395,8 +396,7 @@ static bool canPromoteAllocBox(AllocBoxInst *ABI,
 
 /// rewriteAllocBoxAsAllocStack - Replace uses of the alloc_box with a
 /// new alloc_stack, but do not delete the alloc_box yet.
-static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI,
-                                   llvm::SmallVectorImpl<TermInst *> &Returns) {
+static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI) {
   DEBUG(llvm::dbgs() << "*** Promoting alloc_box to stack: " << *ABI);
 
   llvm::SmallVector<SILInstruction*, 4> FinalReleases;
@@ -405,8 +405,7 @@ static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI,
 
   // Promote this alloc_box to an alloc_stack. Insert the alloc_stack
   // at the beginning of the function.
-  auto &Entry = ABI->getFunction()->front();
-  SILBuilder BuildAlloc(&Entry, Entry.begin());
+  SILBuilder BuildAlloc(ABI);
   BuildAlloc.setCurrentDebugScope(ABI->getDebugScope());
   assert(ABI->getBoxType()->getLayout()->getFields().size() == 1
          && "rewriting multi-field box not implemented");
@@ -440,25 +439,18 @@ static bool rewriteAllocBoxAsAllocStack(AllocBoxInst *ABI,
     .getTypeLowering(ABI->getBoxType()->getFieldType(ABI->getModule(), 0));
   auto Loc = CleanupLocation::get(ABI->getLoc());
 
-  // For non-trivial types, insert destroys for each final release-like
-  // instruction we found that isn't an explicit dealloc_box.
-  if (!Lowering.isTrivial()) {
-    for (auto LastRelease : FinalReleases) {
-      if (isa<DeallocBoxInst>(LastRelease))
-        continue;
-
-      SILBuilderWithScope BuildDestroy(LastRelease);
-      BuildDestroy.emitDestroyAddrAndFold(Loc, PointerResult);
+  for (auto LastRelease : FinalReleases) {
+    SILBuilderWithScope Builder(LastRelease);
+    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.createDeallocStack(Loc, ASI);
   }
 
-  for (auto Return : Returns) {
-    SILBuilderWithScope BuildDealloc(Return);
-    BuildDealloc.createDeallocStack(Loc, ASI);
-  }
-
-  // Remove any retain and release instructions.  Since all uses of result #1
-  // are gone, this only walks through uses of result #0 (the retain count
+  // Remove any retain and release instructions.  Since all uses of project_box
+  // are gone, this only walks through uses of the box itself (the retain count
   // pointer).
   while (!ABI->use_empty()) {
     auto *User = (*ABI->use_begin())->getUser();
@@ -825,7 +817,6 @@ rewritePartialApplies(llvm::SmallVectorImpl<Operand *> &PromotedOperands,
 static unsigned
 rewritePromotedBoxes(llvm::SmallVectorImpl<AllocBoxInst *> &Promoted,
                      llvm::SmallVectorImpl<Operand *> &PromotedOperands,
-                     llvm::SmallVectorImpl<TermInst *> &Returns,
                      bool &CFGChanged) {
   // First we'll rewrite any partial applies that we can to remove the
   // box container pointer from the operands.
@@ -835,7 +826,7 @@ rewritePromotedBoxes(llvm::SmallVectorImpl<AllocBoxInst *> &Promoted,
   auto rend = Promoted.rend();
   for (auto I = Promoted.rbegin(); I != rend; ++I) {
     auto *ABI = *I;
-    if (rewriteAllocBoxAsAllocStack(ABI, Returns)) {
+    if (rewriteAllocBoxAsAllocStack(ABI)) {
       ++Count;
       ABI->eraseFromParent();
     }
@@ -849,13 +840,8 @@ class AllocBoxToStack : public SILFunctionTransform {
   void run() override {
     llvm::SmallVector<AllocBoxInst *, 8> Promotable;
     llvm::SmallVector<Operand *, 8> PromotedOperands;
-    llvm::SmallVector<TermInst *, 8> Returns;
 
     for (auto &BB : *getFunction()) {
-      auto *Term = BB.getTerminator();
-      if (Term->isFunctionExiting())
-        Returns.push_back(Term);
-
       for (auto &I : BB)
         if (auto *ABI = dyn_cast<AllocBoxInst>(&I))
           if (canPromoteAllocBox(ABI, PromotedOperands))
@@ -864,9 +850,14 @@ class AllocBoxToStack : public SILFunctionTransform {
 
     if (!Promotable.empty()) {
       bool CFGChanged = false;
-      auto Count = rewritePromotedBoxes(Promotable, PromotedOperands, Returns,
+      auto Count = rewritePromotedBoxes(Promotable, PromotedOperands,
                                         CFGChanged);
       NumStackPromoted += Count;
+      if (Count) {
+        StackNesting SN;
+        if (SN.correctStackNesting(getFunction()) == StackNesting::Changes::CFG)
+          CFGChanged = true;
+      }
 
       invalidateAnalysis(CFGChanged ?
                          SILAnalysis::InvalidationKind::FunctionBody :

lib/SILOptimizer/Utils/CMakeLists.txt

@@ -13,5 +13,6 @@ set(UTILS_SOURCES
   Utils/SILInliner.cpp
   Utils/SILSSAUpdater.cpp
   Utils/SpecializationMangler.cpp
+  Utils/StackNesting.cpp
   PARENT_SCOPE)