[IRGen] Dealloc metadata packs on dom frontier.

Deallocate dynamic allocas done for metadata/wtable packs.  These
stackrestore calls are inserted on the dominance frontier and then stack
nesting is fixed up.  That was achieved as follows:

Added a new IRGen pass PackMetadataMarkerInserter; it
- determines if there are any instructions which might allocate on-stack
  pack metadata
- if there aren't, no changes are made
- if there are, alloc_pack_metadata just before instructions that could
  allocate pack metadata on the stack and dealloc_pack_metadata on the
  dominance frontier of those instructions
- fixup stack nesting

During IRGen, the allocations done for metadata/wtable packs are
recorded and IRGenSILFunction associates them with the instruction that
lowered.  It must be the instruction after some alloc_pack_metadata
instruction.  Then, when visiting the dealloc_pack_metadata instructions
corresponding to that alloc_pack_metadata, deallocate those packs.

Author: nate-chandler

include/swift/IRGen/IRGenSILPasses.h

@@ -19,6 +19,7 @@ namespace irgen {
 /// Create a pass to hoist alloc_stack instructions with non-fixed size.
 SILTransform *createAllocStackHoisting();
 SILTransform *createLoadableByAddress();
+SILTransform *createPackMetadataMarkerInserter();
 
 } // end namespace irgen
 } // end namespace swift

include/swift/SILOptimizer/PassManager/Passes.def

@@ -353,6 +353,8 @@ IRGEN_PASS(LoadableByAddress, "loadable-address",
      "SIL Large Loadable type by-address lowering.")
 PASS(MandatorySILLinker, "mandatory-linker",
      "Deserialize all referenced SIL functions that are shared or transparent")
+IRGEN_PASS(PackMetadataMarkerInserter, "pack-metadata-marker-inserter",
+     "Insert markers where pack metadata might be de/allocated.")
 PASS(PerformanceSILLinker, "performance-linker",
      "Deserialize all referenced SIL functions")
 PASS(RawSILInstLowering, "raw-sil-inst-lowering",

lib/IRGen/CMakeLists.txt

@@ -59,6 +59,7 @@ add_swift_host_library(swiftIRGen STATIC
   MetadataLayout.cpp
   MetadataRequest.cpp
   Outlining.cpp
+  PackMetadataMarkerInserter.cpp
   StructLayout.cpp
   SwiftTargetInfo.cpp
   TBDGen.cpp

lib/IRGen/GenPack.cpp

@@ -35,6 +35,11 @@
 using namespace swift;
 using namespace irgen;
 
+static void cleanupTypeMetadataPackImpl(IRGenFunction &IGF, StackAddress pack,
+                                        llvm::Value *shape);
+static void cleanupWitnessTablePackImpl(IRGenFunction &IGF, StackAddress pack,
+                                        llvm::Value *shape);
+
 static CanPackArchetypeType
 getForwardedPackArchetypeType(CanPackType packType) {
   if (auto expansion = packType.unwrapSingletonPackExpansion())
@@ -454,6 +459,7 @@ irgen::emitTypeMetadataPack(IRGenFunction &IGF, CanPackType packType,
     assert(packType->getNumElements() == constantInt->getValue());
     auto pack =
         StackAddress(emitFixedSizeMetadataPackRef(IGF, packType, request));
+    IGF.recordStackPackMetadataAlloc(pack, constantInt);
     return {pack, constantInt};
   }
 
@@ -488,6 +494,7 @@ irgen::emitTypeMetadataPack(IRGenFunction &IGF, CanPackType packType,
     };
 
   visitPackExplosion(IGF, packType, visitFn);
+  IGF.recordStackPackMetadataAlloc(pack, shape);
 
   return {pack, shape};
 }
@@ -599,6 +606,7 @@ irgen::emitWitnessTablePack(IRGenFunction &IGF, CanPackType packType,
     assert(packType->getNumElements() == constantInt->getValue());
     auto pack = StackAddress(
         emitFixedSizeWitnessTablePack(IGF, packType, packConformance));
+    IGF.recordStackPackWitnessTableAlloc(pack, constantInt);
     return {pack, constantInt};
   }
 
@@ -635,12 +643,14 @@ irgen::emitWitnessTablePack(IRGenFunction &IGF, CanPackType packType,
   };
 
   visitPackExplosion(IGF, packType, visitFn);
+  IGF.recordStackPackWitnessTableAlloc(pack, shape);
 
   return {pack, shape};
 }
 
-void irgen::cleanupWitnessTablePack(IRGenFunction &IGF, StackAddress pack,
-                                    llvm::Value *shape) {
+static void cleanupWitnessTablePackImpl(IRGenFunction &IGF, StackAddress pack,
+                                        llvm::Value *shape) {
+
   if (pack.getExtraInfo()) {
     IGF.emitDeallocateDynamicAlloca(pack);
   } else if (auto count = countForShape(shape)) {
@@ -649,6 +659,61 @@ void irgen::cleanupWitnessTablePack(IRGenFunction &IGF, StackAddress pack,
   }
 }
 
+void irgen::cleanupWitnessTablePack(IRGenFunction &IGF, StackAddress pack,
+                                    llvm::Value *shape) {
+  cleanupWitnessTablePackImpl(IGF, pack, shape);
+  IGF.eraseStackPackWitnessTableAlloc(pack, shape);
+}
+
+void irgen::cleanupStackAllocPacks(IRGenFunction &IGF,
+                                   ArrayRef<StackPackAlloc> allocs) {
+  for (auto alloc : llvm::reverse(allocs)) {
+    StackAddress addr;
+    uint8_t kind;
+    llvm::Value *shape;
+    std::tie(addr, shape, kind) = alloc;
+
+    switch ((GenericRequirement::Kind)kind) {
+    case GenericRequirement::Kind::MetadataPack:
+      cleanupTypeMetadataPackImpl(IGF, addr, shape);
+      break;
+    case GenericRequirement::Kind::WitnessTablePack:
+      cleanupWitnessTablePackImpl(IGF, addr, shape);
+      break;
+    default:
+      llvm_unreachable("bad requirement in stack pack alloc");
+    }
+  }
+}
+
+void IRGenFunction::recordStackPackMetadataAlloc(StackAddress addr,
+                                                 llvm::Value *shape) {
+  OutstandingStackPackAllocs.insert(
+      {addr, shape, (uint8_t)GenericRequirement::Kind::MetadataPack});
+}
+
+void IRGenFunction::eraseStackPackMetadataAlloc(StackAddress addr,
+                                                llvm::Value *shape) {
+  auto removed = OutstandingStackPackAllocs.remove(
+      {addr, shape, (uint8_t)GenericRequirement::Kind::MetadataPack});
+  assert(removed && "erased stack pack metadata addr that wasn't recorded!?");
+  (void)removed;
+}
+
+void IRGenFunction::recordStackPackWitnessTableAlloc(StackAddress addr,
+                                                     llvm::Value *shape) {
+  OutstandingStackPackAllocs.insert(
+      {addr, shape, (uint8_t)GenericRequirement::Kind::WitnessTablePack});
+}
+
+void IRGenFunction::eraseStackPackWitnessTableAlloc(StackAddress addr,
+                                                    llvm::Value *shape) {
+  auto removed = OutstandingStackPackAllocs.remove(
+      {addr, shape, (uint8_t)GenericRequirement::Kind::WitnessTablePack});
+  assert(removed && "erased stack pack metadata addr that wasn't recorded!?");
+  (void)removed;
+}
+
 llvm::Value *irgen::emitWitnessTablePackRef(IRGenFunction &IGF,
                                             CanPackType packType,
                                             PackConformance *conformance) {
@@ -1045,8 +1110,8 @@ void irgen::bindOpenedElementArchetypesAtIndex(IRGenFunction &IGF,
   }
 }
 
-void irgen::cleanupTypeMetadataPack(IRGenFunction &IGF, StackAddress pack,
-                                    llvm::Value *shape) {
+static void cleanupTypeMetadataPackImpl(IRGenFunction &IGF, StackAddress pack,
+                                        llvm::Value *shape) {
   if (pack.getExtraInfo()) {
     IGF.emitDeallocateDynamicAlloca(pack);
   } else if (auto count = countForShape(shape)) {
@@ -1055,6 +1120,12 @@ void irgen::cleanupTypeMetadataPack(IRGenFunction &IGF, StackAddress pack,
   }
 }
 
+void irgen::cleanupTypeMetadataPack(IRGenFunction &IGF, StackAddress pack,
+                                    llvm::Value *shape) {
+  cleanupTypeMetadataPackImpl(IGF, pack, shape);
+  IGF.eraseStackPackMetadataAlloc(pack, shape);
+}
+
 Address irgen::emitStorageAddressOfPackElement(IRGenFunction &IGF, Address pack,
                                                llvm::Value *index,
                                                SILType elementType,

lib/IRGen/GenPack.h

@@ -80,6 +80,18 @@ llvm::Value *emitWitnessTablePackRef(IRGenFunction &IGF, CanPackType packType,
 void cleanupWitnessTablePack(IRGenFunction &IGF, StackAddress pack,
                              llvm::Value *shape);
 
+/// An on-stack pack metadata/wtable allocation.
+///
+/// Includes the stack address, the element count, and the kind of requirement
+/// (a GenericRequirement::Kind represented as a raw uint8_t).
+using StackPackAlloc =
+    std::tuple<StackAddress, /*shape*/ llvm::Value *, /*kind*/ uint8_t>;
+
+/// Emits cleanups for an array of on-stack pack metadata/wtable allocations in
+/// reverse order.
+void cleanupStackAllocPacks(IRGenFunction &IGF,
+                            ArrayRef<StackPackAlloc> allocs);
+
 /// Emit the dynamic index of a particular structural component
 /// of the given pack type.  If the component is a pack expansion, this
 /// is the index of the first element of the pack (or where it would be

lib/IRGen/IRGenFunction.cpp

@@ -67,6 +67,8 @@ IRGenFunction::~IRGenFunction() {
 
   // Tear down any side-table data structures.
   if (LocalTypeData) destroyLocalTypeData();
+
+  // All dynamically allocated metadata should have been cleaned up.
 }
 
 OptimizationMode IRGenFunction::getEffectiveOptimizationMode() const {

lib/IRGen/IRGenFunction.h

@@ -18,16 +18,17 @@
 #ifndef SWIFT_IRGEN_IRGENFUNCTION_H
 #define SWIFT_IRGEN_IRGENFUNCTION_H
 
-#include "swift/Basic/LLVM.h"
-#include "swift/AST/Type.h"
+#include "DominancePoint.h"
+#include "GenPack.h"
+#include "IRBuilder.h"
+#include "LocalTypeDataKind.h"
 #include "swift/AST/ReferenceCounting.h"
+#include "swift/AST/Type.h"
+#include "swift/Basic/LLVM.h"
 #include "swift/SIL/SILLocation.h"
 #include "swift/SIL/SILType.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/IR/CallingConv.h"
-#include "IRBuilder.h"
-#include "LocalTypeDataKind.h"
-#include "DominancePoint.h"
 
 namespace llvm {
   class AllocaInst;
@@ -206,6 +207,9 @@ class IRGenFunction {
   // particular.
   bool packMetadataStackPromotionDisabled = false;
 
+  /// The on-stack pack metadata allocations emitted so far awaiting cleanup.
+  llvm::SmallSetVector<StackPackAlloc, 2> OutstandingStackPackAllocs;
+
 private:
   Address asyncContextLocation;
 
@@ -381,6 +385,12 @@ class IRGenFunction {
 
   llvm::Value *emitPackShapeExpression(CanType type);
 
+  void recordStackPackMetadataAlloc(StackAddress addr, llvm::Value *shape);
+  void eraseStackPackMetadataAlloc(StackAddress addr, llvm::Value *shape);
+
+  void recordStackPackWitnessTableAlloc(StackAddress addr, llvm::Value *shape);
+  void eraseStackPackWitnessTableAlloc(StackAddress addr, llvm::Value *shape);
+
   /// Emit a load of a reference to the given Objective-C selector.
   llvm::Value *emitObjCSelectorRefLoad(StringRef selector);
 

lib/IRGen/IRGenSIL.cpp

@@ -441,7 +441,15 @@ class IRGenSILFunction :
 
   // A cached dominance analysis.
   std::unique_ptr<DominanceInfo> Dominance;
-  
+
+  /// For each instruction which did allocate pack metadata on-stack, the stack
+  /// locations at which they were allocated.
+  ///
+  /// Used to emit cleanups for those allocations in
+  /// emitDeallocateDynamicPackMetadataAllocas.
+  llvm::DenseMap<SILInstruction *, llvm::SmallVector<StackPackAlloc, 2>>
+      StackPackAllocs;
+
   IRGenSILFunction(IRGenModule &IGM, SILFunction *f);
   ~IRGenSILFunction();
 
@@ -1479,7 +1487,6 @@ class IRGenSILFunction :
   LOADABLE_REF_STORAGE_HELPER(Name)
 #include "swift/AST/ReferenceStorage.def"
 #undef LOADABLE_REF_STORAGE_HELPER
-  
 };
 
 } // end anonymous namespace
@@ -2549,8 +2556,18 @@ void IRGenSILFunction::visitSILBasicBlock(SILBasicBlock *BB) {
     IGM.emittedRuntimeFuncs.clear();
 #endif
 
+    assert(OutstandingStackPackAllocs.empty());
+
     visit(&I);
 
+    // Record the on-stack pack allocations emitted on behalf of this SIL
+    // instruction.  They will be cleaned up when visiting the corresponding
+    // cleanup markers.
+    for (auto pair : OutstandingStackPackAllocs) {
+      StackPackAllocs[&I].push_back(pair);
+    }
+    OutstandingStackPackAllocs.clear();
+
 #ifdef CHECK_RUNTIME_EFFECT_ANALYSIS
     if (!isa<DebugValueInst>(&I)) {
       SILType impactType;
@@ -5698,6 +5715,10 @@ void IRGenSILFunction::visitDeallocPackInst(swift::DeallocPackInst *i) {
 
 void IRGenSILFunction::visitDeallocPackMetadataInst(
     DeallocPackMetadataInst *i) {
+  auto iter = StackPackAllocs.find(i->getIntroducer());
+  if (iter == StackPackAllocs.end())
+    return;
+  cleanupStackAllocPacks(*this, iter->getSecond());
 }
 
 void IRGenSILFunction::visitDeallocRefInst(swift::DeallocRefInst *i) {