Fix ExistentialSpecializer to correctly cleanup storage.

Handle calling conventions and cleanups in all the places (hopefully).
- when ExistentialSpecializer copies the specialized concrete arg into the
original existential value
- when ExistentialSpecializer generates a think
- when SILCombine substitutes concrete values in place of the opened
existential.

One particularly nasty problem is the existential boxes need to be
destroyed. It is not ok to simply destroy their value. The "leaks"
tool does not catch this problem.

Ownership SIL will make this all much more robust.

Fixes <rdar://problem/50595630> Multiple leaks detected - Swift Perf

Author: atrick

lib/SILOptimizer/FunctionSignatureTransforms/ExistentialSpecializer.cpp

@@ -115,19 +115,6 @@ bool ExistentialSpecializer::findConcreteTypeFromSoleConformingType(
   return true;
 }
 
-/// Check if the argument Arg is used in a destroy_use instruction.
-static void
-findIfCalleeUsesArgInDestroyUse(SILValue Arg,
-                                ExistentialTransformArgumentDescriptor &ETAD) {
-  for (Operand *ArgUse : Arg->getUses()) {
-    auto *ArgUser = ArgUse->getUser();
-    if (isa<DestroyAddrInst>(ArgUser)) {
-      ETAD.DestroyAddrUse = true;
-      break;
-    }
-  }
-}
-
 /// Helper function to ensure that the argument is not InOut or InOut_Aliasable
 static bool isNonInoutIndirectArgument(SILValue Arg,
                                        SILArgumentConvention ArgConvention) {
@@ -193,18 +180,21 @@ bool ExistentialSpecializer::canSpecializeExistentialArgsInFunction(
       continue;
     }
 
-    /// Determine attributes of the existential addr arguments such as
-    /// destroy_use, immutable_access. 
+    /// Determine attributes of the existential addr argument.
     ExistentialTransformArgumentDescriptor ETAD;
     auto paramInfo = origCalleeConv.getParamInfoForSILArg(Idx);
-    ETAD.AccessType = (paramInfo.isIndirectMutating() || paramInfo.isConsumed())
+    // The ExistentialSpecializerCloner copies the incoming generic argument
+    // into an existential. This won't work if the original argument is
+    // mutated. Furthermore, SILCombine would not be able to replace a mutated
+    // existential with a concrete value, so the specialization thunk could not
+    // be optimized away.
+    if (paramInfo.isIndirectMutating())
+      continue;
+
+    ETAD.AccessType = paramInfo.isConsumed()
                           ? OpenedExistentialAccess::Mutable
                           : OpenedExistentialAccess::Immutable;
-    ETAD.DestroyAddrUse = false;
-    if ((CalleeArgs[Idx]->getType().getPreferredExistentialRepresentation(
-            F->getModule()))
-        != ExistentialRepresentation::Class)
-      findIfCalleeUsesArgInDestroyUse(CalleeArg, ETAD);
+    ETAD.isConsumed = paramInfo.isConsumed();
 
     /// Save the attributes
     ExistentialArgDescriptor[Idx] = ETAD;

lib/SILOptimizer/FunctionSignatureTransforms/ExistentialTransform.cpp

@@ -57,9 +57,10 @@ class ExistentialSpecializerCloner
   // Use one OpenedArchetypesTracker while cloning.
   SILOpenedArchetypesTracker OpenedArchetypesTracker;
 
-  // Holds an entry for any indirect argument requiring a new alloc_stack
-  // created during argument cloning.
-  SmallDenseMap<int, AllocStackInst *> ArgToAllocStackMap;
+  // AllocStack instructions introduced in the new prolog that require cleanup.
+  SmallVector<AllocStackInst *, 4> AllocStackInsts;
+  // Temporary values introduced in the new prolog that require cleanup.
+  SmallVector<SILValue, 4> CleanupValues;
 
 protected:
   void postProcess(SILInstruction *Orig, SILInstruction *Cloned) {
@@ -99,36 +100,24 @@ void ExistentialSpecializerCloner::cloneAndPopulateFunction() {
   auto *NewEntryBB = getBuilder().getFunction().getEntryBlock();
   cloneFunctionBody(&Original, NewEntryBB, entryArgs);
 
-  /// If there is an argument with no DestroyUse, insert DeallocStack
-  /// before return Instruction.
-  SmallPtrSet<ReturnInst *, 4> ReturnInsts;
-  /// Find the set of return instructions in a function.
-  for (auto &BB : getBuilder().getFunction()) {
-    TermInst *TI = BB.getTerminator();
-    if (auto *RI = dyn_cast<ReturnInst>(TI)) {
-      ReturnInsts.insert(RI);
-    }
-  }
-
-  /// Walk over destroy_addr instructions and perform fixups.
-  for (auto &ArgDesc : reversed(ArgumentDescList)) {
-    int ArgIndex = ArgDesc.Index;
-    auto iter = ArgToAllocStackMap.find(ArgIndex);
-    if (iter != ArgToAllocStackMap.end()) {
-      // Need to insert DeallocStack before return.
-      for (auto *I : ReturnInsts) {
-        SILBuilderWithScope Builder(I);
-        // A return location can't be used for a non-return instruction.
-        auto loc = RegularLocation::getAutoGeneratedLocation();
-        Builder.createDeallocStack(loc, iter->second);
-      }
-    }
+  // Cleanup allocations created in the new prolog.
+  SmallVector<SILBasicBlock *, 4> exitingBlocks;
+  getBuilder().getFunction().findExitingBlocks(exitingBlocks);
+  for (auto *exitBB : exitingBlocks) {
+    SILBuilderWithScope Builder(exitBB->getTerminator());
+    // A return location can't be used for a non-return instruction.
+    auto loc = RegularLocation::getAutoGeneratedLocation();
+    for (SILValue cleanupVal : CleanupValues)
+      Builder.createDestroyAddr(loc, cleanupVal);
+
+    for (auto *ASI : reversed(AllocStackInsts))
+      Builder.createDeallocStack(loc, ASI);
   }
 }
 
 // Create the entry basic block with the function arguments.
-void ExistentialSpecializerCloner::
-cloneArguments(SmallVectorImpl<SILValue> &entryArgs) {
+void ExistentialSpecializerCloner::cloneArguments(
+    SmallVectorImpl<SILValue> &entryArgs) {
   auto &M = OrigF->getModule();
   auto &Ctx = M.getASTContext();
 
@@ -190,38 +179,42 @@ cloneArguments(SmallVectorImpl<SILValue> &entryArgs) {
       /// bb0(%0 : $*T):
       /// %3 = alloc_stack $P
       /// %4 = init_existential_addr %3 : $*P, $T
-      /// copy_addr %0 to [initialization] %4 : $*T
-      /// destroy_addr %0 : $*T
+      /// copy_addr [take] %0 to [initialization] %4 : $*T
       /// %7 = open_existential_addr immutable_access %3 : $*P to
       /// $*@opened P
       auto *ASI =
           NewFBuilder.createAllocStack(InsertLoc, ArgDesc.Arg->getType());
-      ArgToAllocStackMap.insert(
-          std::pair<int, AllocStackInst *>(ArgDesc.Index, ASI));
+      AllocStackInsts.push_back(ASI);
 
       auto *EAI = NewFBuilder.createInitExistentialAddr(
           InsertLoc, ASI, NewArg->getType().getASTType(), NewArg->getType(),
           Conformances);
-      /// If DestroyAddr is already there, then do not use [take].
-      NewFBuilder.createCopyAddr(
-        InsertLoc, NewArg, EAI,
-        ExistentialArgDescriptor[ArgDesc.Index].AccessType ==
-        OpenedExistentialAccess::Immutable
-        ? IsTake_t::IsNotTake
-        : IsTake_t::IsTake,
-        IsInitialization_t::IsInitialization);
-      if (ExistentialArgDescriptor[ArgDesc.Index].DestroyAddrUse) {
-        NewFBuilder.createDestroyAddr(InsertLoc, NewArg);
-      }
+
+      bool origConsumed = ExistentialArgDescriptor[ArgDesc.Index].isConsumed;
+      // If the existential is not consumed in the function body, then the one
+      // we introduce here needs cleanup.
+      if (!origConsumed)
+        CleanupValues.push_back(ASI);
+
+      NewFBuilder.createCopyAddr(InsertLoc, NewArg, EAI,
+                                 origConsumed ? IsTake_t::IsTake
+                                              : IsTake_t::IsNotTake,
+                                 IsInitialization_t::IsInitialization);
       entryArgs.push_back(ASI);
       break;
     }
     case ExistentialRepresentation::Class: {
+      // FIXME_ownership: init_existential_ref always takes ownership of the
+      // incoming reference. If the argument convention is borrowed
+      // (!isConsumed), then we should create a copy_value here and add this new
+      // existential to the CleanupValues vector.
+
       ///  Simple case: Create an init_existential.
       /// %5 = init_existential_ref %0 : $T : $T, $P
       auto *InitRef = NewFBuilder.createInitExistentialRef(
           InsertLoc, ArgDesc.Arg->getType(), NewArg->getType().getASTType(),
           NewArg, Conformances);
+
       entryArgs.push_back(InitRef);
       break;
     }
@@ -399,12 +392,16 @@ void ExistentialTransform::populateThunkBody() {
   /// Determine arguments to Apply.
   /// Generate opened existentials for generics.
   SmallVector<SILValue, 8> ApplyArgs;
+  // Maintain a list of arg values to be destroyed. These are consumed by the
+  // convention and require a copy.
+  SmallVector<CopyAddrInst *, 8> TempCopyAddrInsts;
   SmallDenseMap<GenericTypeParamType *, Type> GenericToOpenedTypeMap;
   for (auto &ArgDesc : ArgumentDescList) {
     auto iter = ArgToGenericTypeMap.find(ArgDesc.Index);
     auto it = ExistentialArgDescriptor.find(ArgDesc.Index);
     if (iter != ArgToGenericTypeMap.end() &&
         it != ExistentialArgDescriptor.end()) {
+      ExistentialTransformArgumentDescriptor &ETAD = it->second;
       OpenedArchetypeType *Opened;
       auto OrigOperand = ThunkBody->getArgument(ArgDesc.Index);
       auto SwiftType = ArgDesc.Arg->getType().getASTType();
@@ -418,11 +415,26 @@ void ExistentialTransform::populateThunkBody() {
       case ExistentialRepresentation::Opaque: {
         archetypeValue = Builder.createOpenExistentialAddr(
             Loc, OrigOperand, OpenedSILType, it->second.AccessType);
-        ApplyArgs.push_back(archetypeValue);
+        SILValue calleeArg = archetypeValue;
+        if (ETAD.isConsumed) {
+          // open_existential_addr projects a borrowed address into the
+          // existential box. Since the callee consumes the generic value, we
+          // must pass in a copy.
+          auto *ASI =
+            Builder.createAllocStack(Loc, OpenedSILType);
+          auto *CAI =
+              Builder.createCopyAddr(Loc, archetypeValue, ASI, IsNotTake,
+                                     IsInitialization_t::IsInitialization);
+          TempCopyAddrInsts.push_back(CAI);
+          calleeArg = ASI;
+        }
+        ApplyArgs.push_back(calleeArg);
         break;
       }
       case ExistentialRepresentation::Class: {
         /// If the operand is not object type, we would need an explicit load.
+        // OpenExistentialRef forwards ownership, so it does the right thing
+        // regardless of whether the argument is borrowed or consumed.
         assert(OrigOperand->getType().isObject());
         archetypeValue =
             Builder.createOpenExistentialRef(Loc, OrigOperand, OpenedSILType);
@@ -498,7 +510,24 @@ void ExistentialTransform::populateThunkBody() {
     /// Create the Apply with substitutions
     ReturnValue = Builder.createApply(Loc, FRI, SubMap, ApplyArgs);
   }
-
+  auto cleanupLoc = RegularLocation::getAutoGeneratedLocation();
+  for (CopyAddrInst *CAI : reversed(TempCopyAddrInsts)) {
+    // The original argument was copied into a temporary and consumed by the
+    // callee as such:
+    //   bb (%consumedExistential : $*Protocol)
+    //     %valAdr = open_existential_addr %consumedExistential
+    //     %temp = alloc_stack $T
+    //     copy_addr %valAdr to %temp // <== Temp CopyAddr
+    //     apply(%temp)               // <== Temp is consumed by the apply
+    //
+    // Destroy the original arument and deallocation the temporary:
+    //     destroy_addr %consumedExistential : $*Protocol
+    //     dealloc_stack %temp : $*T
+    auto *consumedExistential = cast<SILFunctionArgument>(
+        cast<OpenExistentialAddrInst>(CAI->getSrc())->getOperand());
+    Builder.createDestroyAddr(cleanupLoc, consumedExistential);
+    Builder.createDeallocStack(cleanupLoc, CAI->getDest());
+  }
   /// Set up the return results.
   if (NewF->isNoReturnFunction()) {
     Builder.createUnreachable(Loc);

lib/SILOptimizer/FunctionSignatureTransforms/ExistentialTransform.h

@@ -33,7 +33,7 @@ namespace swift {
 /// to transformation.
 struct ExistentialTransformArgumentDescriptor {
   OpenedExistentialAccess AccessType;
-  bool DestroyAddrUse;
+  bool isConsumed;
 };
 
 /// ExistentialTransform creates a protocol constrained generic and a thunk.