Fix a memory leak in FSO

Make sure the destructor of the SmallVector in ProjectionTreeNode gets
called when the BumpPtrAllocator is destroy'ed.

Author: trentxintong

include/swift/SIL/Projection.h

@@ -808,7 +808,8 @@ class ProjectionTree {
 
   SILModule &Mod;
 
-  llvm::BumpPtrAllocator &Allocator;
+  /// The allocator we use to allocate ProjectionTreeNodes in the tree.
+  llvm::SpecificBumpPtrAllocator<ProjectionTreeNode> Allocator;
 
   // A common pattern is a 3 field struct.
   llvm::SmallVector<ProjectionTreeNode *, 4> ProjectionTreeNodes;
@@ -818,12 +819,10 @@ class ProjectionTree {
 
 public:
   /// Construct a projection tree from BaseTy.
-  ProjectionTree(SILModule &Mod, llvm::BumpPtrAllocator &Allocator,
-                 SILType BaseTy);
+  ProjectionTree(SILModule &Mod, SILType BaseTy);
   /// Construct an uninitialized projection tree, which can then be
   /// initialized by initializeWithExistingTree.
-  ProjectionTree(SILModule &Mod, llvm::BumpPtrAllocator &Allocator) 
-    : Mod(Mod), Allocator(Allocator) {}
+  ProjectionTree(SILModule &Mod) : Mod(Mod)  {}
   ~ProjectionTree();
   ProjectionTree(const ProjectionTree &) = delete;
   ProjectionTree(ProjectionTree &&) = default;
@@ -936,15 +935,16 @@ class ProjectionTree {
   void createRoot(SILType BaseTy) {
     assert(ProjectionTreeNodes.empty() &&
            "Should only create root when ProjectionTreeNodes is empty");
-    auto *Node = new (Allocator) ProjectionTreeNode(BaseTy);
+    auto *Node = new (Allocator.Allocate()) ProjectionTreeNode(BaseTy);
     ProjectionTreeNodes.push_back(Node);
   }
 
   ProjectionTreeNode *createChild(ProjectionTreeNode *Parent,
                                      SILType BaseTy,
                                      const Projection &P) {
     unsigned Index = ProjectionTreeNodes.size();
-    auto *Node = new (Allocator) ProjectionTreeNode(Parent, Index, BaseTy, P);
+    auto *Node = new (Allocator.Allocate()) ProjectionTreeNode(Parent, Index,
+                                                               BaseTy, P);
     ProjectionTreeNodes.push_back(Node);
     return ProjectionTreeNodes[Index];
   }

include/swift/SILOptimizer/Utils/FunctionSignatureOptUtils.h

@@ -76,13 +76,13 @@ struct ArgumentDescriptor {
   /// to the original argument. The reason why we do this is to make sure we
   /// have access to the original argument's state if we modify the argument
   /// when optimizing.
-  ArgumentDescriptor(llvm::BumpPtrAllocator &BPA, SILArgument *A)
+  ArgumentDescriptor(SILArgument *A)
       : Arg(A), PInfo(Arg->getKnownParameterInfo()), Index(A->getIndex()),
         Decl(A->getDecl()), IsEntirelyDead(false), Explode(false),
         OwnedToGuaranteed(false),
         IsIndirectResult(A->isIndirectResult()),
         CalleeRelease(), CalleeReleaseInThrowBlock(),
-        ProjTree(A->getModule(), BPA, A->getType())  {}
+        ProjTree(A->getModule(), A->getType())  {}
 
   ArgumentDescriptor(const ArgumentDescriptor &) = delete;
   ArgumentDescriptor(ArgumentDescriptor &&) = default;

lib/SIL/Projection.cpp

@@ -1142,8 +1142,7 @@ class NewAggregateBuilderMap {
 //===----------------------------------------------------------------------===//
 
 ProjectionTree::
-ProjectionTree(SILModule &Mod, llvm::BumpPtrAllocator &BPA, SILType BaseTy) 
-  : Mod(Mod), Allocator(BPA) {
+ProjectionTree(SILModule &Mod, SILType BaseTy) : Mod(Mod) {
   DEBUG(llvm::dbgs() << "Constructing Projection Tree For : " << BaseTy);
 
   // Create the root node of the tree with our base type.
@@ -1161,7 +1160,8 @@ void
 ProjectionTree::initializeWithExistingTree(const ProjectionTree &PT) {
   LiveLeafIndices = PT.LiveLeafIndices;
   for (const auto &N : PT.ProjectionTreeNodes) {
-    ProjectionTreeNodes.push_back(new (Allocator) ProjectionTreeNode(*N));
+    ProjectionTreeNodes.push_back(new (Allocator.Allocate())
+                                  ProjectionTreeNode(*N));
   }
 }
 

lib/SILOptimizer/Transforms/FunctionSignatureOpts.cpp

@@ -12,9 +12,9 @@
 ///
 /// \file
 ///
-/// This pass defines function signature related optimizations. Everytime a
-/// function signature optimization is performed, changes are made to the
-/// original function and after all function signature optimizations are
+/// This pass defines function signature related optimizations.
+/// When a function signature optimization is performed, changes are made to
+/// the original function and after all function signature optimizations are
 /// finished, a new function is created and the old function is turned into
 /// a thunk.
 ///
@@ -164,8 +164,9 @@ class FunctionSignatureTransform {
   /// ----------------------------------------------------------///
   /// Find any dead argument opportunities.
   bool DeadArgumentAnalyzeParameters();
-  /// Do the actual dead argument transformations.
-  void DeadArgumentTransformParameters();
+  /// Modify the current function so that later function signature analysis
+  /// are more effective.
+  void DeadArgumentTransformFunction();
   /// Remove the dead argument once the new function is created.
   void DeadArgumentFinalizeOptimizedFunction();
 
@@ -174,8 +175,11 @@ class FunctionSignatureTransform {
   /// ----------------------------------------------------------///
   bool OwnedToGuaranteedAnalyzeResults();
   bool OwnedToGuaranteedAnalyzeParameters();
-  void OwnedToGuaranteedTransformResults();
-  void OwnedToGuaranteedTransformParameters();
+
+  /// Modify the current function so that later function signature analysis
+  /// are more effective.
+  void OwnedToGuaranteedTransformFunctionResults();
+  void OwnedToGuaranteedTransformFunctionParameters();
 
   /// Find any owned to guaranteed opportunities.
   bool OwnedToGuaranteedAnalyze() {
@@ -186,8 +190,8 @@ class FunctionSignatureTransform {
 
   /// Do the actual owned to guaranteeed transformations.
   void OwnedToGuaranteedTransform() {
-    OwnedToGuaranteedTransformResults();
-    OwnedToGuaranteedTransformParameters();
+    OwnedToGuaranteedTransformFunctionResults();
+    OwnedToGuaranteedTransformFunctionParameters();
   }
 
   /// Set up epilogue work for the thunk result based in the given argument.
@@ -274,7 +278,7 @@ class FunctionSignatureTransform {
     // already created a thunk.
     if ((hasCaller || Changed) && DeadArgumentAnalyzeParameters()) {
       Changed = true;
-      DeadArgumentTransformParameters();
+      DeadArgumentTransformFunction();
     }
 
     // Run ArgumentExplosion transformation. We only specialize
@@ -551,7 +555,7 @@ bool FunctionSignatureTransform::DeadArgumentAnalyzeParameters() {
   return SignatureOptimize;
 }
 
-void FunctionSignatureTransform::DeadArgumentTransformParameters() {
+void FunctionSignatureTransform::DeadArgumentTransformFunction() {
   SILBasicBlock *BB = &*F->begin();
   for (const ArgumentDescriptor &AD : ArgumentDescList) {
     if (!AD.IsEntirelyDead)
@@ -643,20 +647,22 @@ bool FunctionSignatureTransform::OwnedToGuaranteedAnalyzeResults() {
   return SignatureOptimize;
 }
 
-void FunctionSignatureTransform::OwnedToGuaranteedTransformParameters() {
+void FunctionSignatureTransform::OwnedToGuaranteedTransformFunctionParameters() {
   // And remove all Callee releases that we found and made redundant via owned
   // to guaranteed conversion.
   for (const ArgumentDescriptor &AD : ArgumentDescList) {
     if (!AD.OwnedToGuaranteed)
       continue;
-    for (auto &X : AD.CalleeRelease) 
+    for (auto &X : AD.CalleeRelease) { 
       X->eraseFromParent();
-    for (auto &X : AD.CalleeReleaseInThrowBlock) 
+    }
+    for (auto &X : AD.CalleeReleaseInThrowBlock) { 
       X->eraseFromParent();
+    }
   }
 }
 
-void FunctionSignatureTransform::OwnedToGuaranteedTransformResults() {
+void FunctionSignatureTransform::OwnedToGuaranteedTransformFunctionResults() {
   // And remove all callee retains that we found and made redundant via owned
   // to unowned conversion.
   for (const ResultDescriptor &RD : ResultDescList) {
@@ -852,7 +858,6 @@ class FunctionSignatureOpts : public SILFunctionTransform {
     if (!canSpecializeFunction(F))
       return;
 
-    llvm::BumpPtrAllocator BPA;
     auto *AA = PM->getAnalysis<AliasAnalysis>();
     auto *RCIA = getAnalysis<RCIdentityAnalysis>();
 
@@ -876,7 +881,7 @@ class FunctionSignatureOpts : public SILFunctionTransform {
     llvm::SmallVector<ResultDescriptor, 4> ResultDescList;
     ArrayRef<SILArgument *> Args = F->begin()->getBBArgs();
     for (unsigned i = 0, e = Args.size(); i != e; ++i) {
-      ArgumentDescList.emplace_back(BPA, Args[i]);
+      ArgumentDescList.emplace_back(Args[i]);
     }
     for (SILResultInfo IR : F->getLoweredFunctionType()->getAllResults()) {
       ResultDescList.emplace_back(IR);