[pred-deadalloc-elim] Fix memory safety issue and handle promoting paths where there is dynamically no value by inserting compensating destroys.

This commit is fixing two things:

1. In certain cases, we are seeing cases where either SILGen or the optimizer
are eliminating destroy_addr along paths where we know that an enum is
dynamically trivial. This can not be expressed in OSSA, so I added code to
pred-deadalloc-elim so that I check if any of our available values after we
finish promoting away an allocation now need to have their consuming use set
completed.

2. That led me to discover that in certain cases load [take] that we were
promoting were available values of other load [take]. This means that we have a
memory safety issue if we promote one load before the other. Consider the
following SIL:

```
%mem = alloc_stack
store %arg to [init] %mem
%0 = load [take] %mem
store %0 to [init] %mem
%1 = load [take] %mem
destroy_value %1
dealloc_stack %mem
```

In this case, if we eliminate %0 before we eliminate %1, we will have a stale
pointer to %0.

I also took this as an opportunity to turn off predictable mem access opt on SIL
that was deserialized canonicalized and non-OSSA SIL. We evidently need to still
do this for pred mem opts for perf reasons (not sure why). But I am pretty sure
this isn't needed and allows me to avoid some nasty code.

Author: gottesmm

lib/SILOptimizer/Mandatory/PredictableMemOpt.cpp

@@ -13,18 +13,20 @@
 #define DEBUG_TYPE "predictable-memopt"
 
 #include "PMOMemoryUseCollector.h"
+#include "swift/Basic/BlotMapVector.h"
 #include "swift/Basic/BlotSetVector.h"
 #include "swift/Basic/FrozenMultiMap.h"
 #include "swift/Basic/STLExtras.h"
+#include "swift/SIL/BasicBlockBits.h"
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/LinearLifetimeChecker.h"
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILBuilder.h"
-#include "swift/SIL/BasicBlockBits.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/CFGOptUtils.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
+#include "swift/SILOptimizer/Utils/OwnershipOptUtils.h"
 #include "swift/SILOptimizer/Utils/SILSSAUpdater.h"
 #include "swift/SILOptimizer/Utils/ValueLifetime.h"
 #include "llvm/ADT/SmallBitVector.h"
@@ -1959,7 +1961,17 @@ class AllocOptimize {
   bool promoteLoadCopy(LoadInst *li);
   bool promoteLoadBorrow(LoadBorrowInst *lbi);
   bool promoteCopyAddr(CopyAddrInst *cai);
-  void promoteLoadTake(LoadInst *Inst, MutableArrayRef<AvailableValue> values);
+
+  /// Promote a load take cleaning up everything except for RAUWing the
+  /// instruction with the aggregated result. The routine returns the new
+  /// aggreaged result to the caller and expects the caller to eventually RAUW
+  /// \p inst with the return value. The reason why we do this is to allow for
+  /// the caller to work around invalidation issues by not deleting the load
+  /// [take] until after all load [take] have been cleaned up.
+  ///
+  /// \returns the value that the caller will RAUW with \p inst.
+  SILValue promoteLoadTake(LoadInst *inst,
+                           MutableArrayRef<AvailableValue> values);
   void promoteDestroyAddr(DestroyAddrInst *dai,
                           MutableArrayRef<AvailableValue> values);
   bool canPromoteTake(SILInstruction *i,
@@ -2293,7 +2305,7 @@ void AllocOptimize::promoteDestroyAddr(
   dai->eraseFromParent();
 }
 
-void AllocOptimize::promoteLoadTake(
+SILValue AllocOptimize::promoteLoadTake(
     LoadInst *li, MutableArrayRef<AvailableValue> availableValues) {
   assert(li->getOwnershipQualifier() == LoadOwnershipQualifier::Take &&
          "load [copy], load [trivial], load should be handled by "
@@ -2311,15 +2323,15 @@ void AllocOptimize::promoteLoadTake(
   AvailableValueAggregator agg(li, availableValues, Uses, deadEndBlocks,
                                AvailableValueExpectedOwnership::Take);
   SILValue newVal = agg.aggregateValues(loadTy, address, firstElt);
+  assert(newVal);
 
   ++NumLoadTakePromoted;
 
   LLVM_DEBUG(llvm::dbgs() << "  *** Promoting load_take: " << *li);
   LLVM_DEBUG(llvm::dbgs() << "      To value: " << *newVal);
 
-  // Then perform the RAUW.
-  li->replaceAllUsesWith(newVal);
-  li->eraseFromParent();
+  // Our parent RAUWs with newVal/erases li.
+  return newVal;
 }
 
 namespace {
@@ -2335,6 +2347,33 @@ struct TakePromotionState {
 
   unsigned size() const { return takeInstIndices.size(); }
 
+  void verify() {
+#ifndef NDEBUG
+    for (unsigned i : range(size())) {
+      SILInstruction *inst;
+      MutableArrayRef<AvailableValue> data;
+      std::tie(inst, data) = getData(i);
+      assert(inst);
+      inst->verifyOperandOwnership();
+      assert(!data.empty() && "Value without any available values?!");
+    }
+#endif
+  }
+
+  void verify(unsigned startOffset) {
+#ifndef NDEBUG
+    assert(startOffset < size());
+    for (unsigned i : range(startOffset, size())) {
+      SILInstruction *inst;
+      MutableArrayRef<AvailableValue> data;
+      std::tie(inst, data) = getData(i);
+      assert(inst);
+      inst->verifyOperandOwnership();
+      assert(!data.empty() && "Value without any available values?!");
+    }
+#endif
+  }
+
   void initializeForTakeInst(unsigned takeInstIndex) {
     availableValueStartOffsets.push_back(availableValueList.size());
     takeInstIndices.push_back(takeInstIndex);
@@ -2352,9 +2391,8 @@ struct TakePromotionState {
       count = availableValueList.size() - startOffset;
     }
 
-    MutableArrayRef<AvailableValue> values(&availableValueList[startOffset],
-                                           count);
-    return {takeInsts[takeInstIndex], values};
+    auto values = MutableArrayRef<AvailableValue>(availableValueList);
+    return {takeInsts[takeInstIndex], values.slice(startOffset, count)};
   }
 };
 
@@ -2424,6 +2462,8 @@ static bool isRemovableAutogeneratedAllocation(AllocationInst *TheMemory) {
 }
 
 bool AllocOptimize::tryToRemoveDeadAllocation() {
+  assert(TheMemory->getFunction()->hasOwnership() &&
+         "Can only eliminate dead allocations with ownership enabled");
   assert((isa<AllocBoxInst>(TheMemory) || isa<AllocStackInst>(TheMemory)) &&
          "Unhandled allocation case");
 
@@ -2492,20 +2532,92 @@ bool AllocOptimize::tryToRemoveDeadAllocation() {
   }
 
   // If we reached this point, we can promote all of our destroy_addr and load
-  // take. Since our load [take] may be available values for our destroy_addr,
-  // we promote the destroy_addr first.
+  // take. Before we begin, gather up all found available values before we do
+  // anything so we can fix up lifetimes later if we need to.
+  SmallBlotSetVector<SILValue, 32> valuesNeedingLifetimeCompletion;
+  for (auto pmoMemUse : Uses) {
+    if (pmoMemUse.Inst && pmoMemUse.Kind == PMOUseKind::Initialization) {
+      // Today if we promote, this is always a store, since we would have
+      // blown up the copy_addr otherwise. Given that, always make sure we
+      // clean up the src as appropriate after we optimize.
+      auto *si = cast<StoreInst>(pmoMemUse.Inst);
+      auto src = si->getSrc();
+
+      // Bail if src has any uses that are forwarding unowned uses. This
+      // allows us to know that we never have to deal with forwarding unowned
+      // instructions like br. These are corner cases that complicate the
+      // logic below.
+      for (auto *use : src->getUses()) {
+        if (use->getOperandOwnership() == OperandOwnership::ForwardingUnowned)
+          return false;
+      }
+      valuesNeedingLifetimeCompletion.insert(src);
+    }
+  }
+
+  // Since our load [take] may be available values for our
+  // destroy_addr/load [take], we promote the destroy_addr first and then handle
+  // load [take] with extra rigour later to handle that possibility.
   for (unsigned i : range(destroyAddrState.size())) {
     SILInstruction *dai;
     MutableArrayRef<AvailableValue> values;
     std::tie(dai, values) = destroyAddrState.getData(i);
     promoteDestroyAddr(cast<DestroyAddrInst>(dai), values);
     // We do not need to unset releases, since we are going to exit here.
   }
+
+  llvm::SmallMapVector<LoadInst *, SILValue, 32> loadsToDelete;
   for (unsigned i : range(loadTakeState.size())) {
     SILInstruction *li;
     MutableArrayRef<AvailableValue> values;
     std::tie(li, values) = loadTakeState.getData(i);
-    promoteLoadTake(cast<LoadInst>(li), values);
+
+    for (unsigned i : indices(values)) {
+      auto v = values[i].Value;
+      auto *li = dyn_cast<LoadInst>(v);
+      if (!li)
+        continue;
+
+      auto iter = loadsToDelete.find(li);
+      if (iter == loadsToDelete.end())
+        continue;
+
+      SILValue newValue = iter->second;
+      assert(newValue && "We should neer store a nil SILValue into this map");
+      values[i].Value = newValue;
+    }
+
+    auto *liCast = cast<LoadInst>(li);
+    SILValue result = promoteLoadTake(liCast, values);
+    assert(result);
+
+    // We need to erase liCast here before we erase it since a load [take] that
+    // we are promoting could be an available value for another load
+    // [take]. Consider the following SIL:
+    //
+    // %mem = alloc_stack
+    // store %arg to [init] %mem
+    // %0 = load [take] %mem
+    // store %0 to [init] %mem
+    // %1 = load [take] %mem
+    // destroy_value %1
+    // dealloc_stack %mem
+    //
+    // In such a case, we are going to delete %0 here, but %0 is an available
+    // value for %1, so we will
+    auto insertIter = loadsToDelete.insert({liCast, result});
+    valuesNeedingLifetimeCompletion.erase(liCast);
+    (void)insertIter;
+    assert(insertIter.second && "loadTakeState doesn't have unique loads?!");
+  }
+
+  // Now that we have promoted all of our load [take], perform the actual
+  // RAUW/removal.
+  for (auto p : loadsToDelete) {
+    LoadInst *li = p.first;
+    SILValue newValue = p.second;
+    li->replaceAllUsesWith(newValue);
+    li->eraseFromParent();
   }
 
   LLVM_DEBUG(llvm::dbgs() << "*** Removing autogenerated non-trivial alloc: "
@@ -2516,6 +2628,144 @@ bool AllocOptimize::tryToRemoveDeadAllocation() {
   // caller remove the allocation itself to avoid iterator invalidation.
   eraseUsesOfInstruction(TheMemory);
 
+  // Now look at all of our available values and complete any of their
+  // post-dominating consuming use sets. This can happen if we have an enum that
+  // is known dynamically none along a path. This is dynamically correct, but
+  // can not be represented in OSSA so we insert these destroys along said path.
+  SmallVector<SILBasicBlock *, 32> consumingUseBlocks;
+  while (!valuesNeedingLifetimeCompletion.empty()) {
+    auto optV = valuesNeedingLifetimeCompletion.pop_back_val();
+    if (!optV)
+      continue;
+    SILValue v = *optV;
+    if (v.getOwnershipKind() != OwnershipKind::Owned)
+      continue;
+
+    // First see if our value doesn't have any uses. In such a case, just
+    // insert a destroy_value at the next instruction and return.
+    if (v->use_empty()) {
+      auto *next = v->getNextInstruction();
+      auto loc = RegularLocation::getAutoGeneratedLocation();
+      SILBuilderWithScope localBuilder(next);
+      localBuilder.createDestroyValue(loc, v);
+      continue;
+    }
+
+    // Otherwise, we first see if we have any consuming uses at all. If we do,
+    // then we know that any such consuming uses since we have an owned value
+    // /must/ be strongly control equivalent to our value and unreachable from
+    // each other, so we can just use findJointPostDominatingSet to complete
+    // the set.
+    consumingUseBlocks.clear();
+    for (auto *use : v->getConsumingUses())
+      consumingUseBlocks.push_back(use->getParentBlock());
+
+    if (!consumingUseBlocks.empty()) {
+      findJointPostDominatingSet(
+          v->getParentBlock(), consumingUseBlocks, [](SILBasicBlock *) {},
+          [&](SILBasicBlock *result) {
+            auto loc = RegularLocation::getAutoGeneratedLocation();
+            SILBuilderWithScope builder(result);
+            builder.createDestroyValue(loc, v);
+          });
+      continue;
+    }
+
+    // If we do not have at least one consuming use, we need to do something
+    // different. This situation can occur given a non-trivial enum typed
+    // stack allocation that:
+    //
+    // 1. Had a destroy_addr eliminated along a path where we dynamically know
+    //    that the stack allocation is storing a trivial case.
+    //
+    // 2. Had some other paths where due to dead end blocks, no destroy_addr
+    //    is needed.
+    //
+    // To fix this, we just treat all uses as consuming blocks and insert
+    // destroys using the joint post dominance set computer and insert
+    // destroys at the end of all input blocks in the post dom set and at the
+    // beginning of any leaking blocks.
+    {
+      // TODO: Can we just pass this in to findJointPostDominatingSet instead
+      // of recomputing it there? Maybe an overload that lets us do this?
+      BasicBlockSet foundUseBlocks(v->getFunction());
+      for (auto *use : v->getUses()) {
+        auto *block = use->getParentBlock();
+        if (!foundUseBlocks.insert(block))
+          continue;
+        consumingUseBlocks.push_back(block);
+      }
+    }
+    findJointPostDominatingSet(
+        v->getParentBlock(), consumingUseBlocks,
+        [&](SILBasicBlock *foundInputBlock) {
+          // This is a block that is reachable from another use. We are not
+          // interested in these.
+        },
+        [&](SILBasicBlock *leakingBlock) {
+          auto loc = RegularLocation::getAutoGeneratedLocation();
+          SILBuilderWithScope builder(leakingBlock);
+          builder.createDestroyValue(loc, v);
+        },
+        [&](SILBasicBlock *inputBlockInPostDomSet) {
+          auto *termInst = inputBlockInPostDomSet->getTerminator();
+          switch (termInst->getTermKind()) {
+          case TermKind::UnreachableInst:
+            // We do not care about input blocks that end in unreachables. We
+            // are going to leak down them so do not insert a destroy_value
+            // there.
+            return;
+
+          // NOTE: Given that our input value is owned, our branch can only
+          // accept the use as a non-consuming use if the branch is forwarding
+          // unowned ownership. Luckily for use, we checked early if we had
+          // any such uses and bailed, so we know the branch can not use our
+          // value. This is just avoiding a corner case that we don't need to
+          // handle.
+          case TermKind::BranchInst:
+            LLVM_FALLTHROUGH;
+          // NOTE: We put cond_br here since in OSSA, cond_br can never have
+          // a non-trivial value operand, meaning we can insert before.
+          case TermKind::CondBranchInst:
+            LLVM_FALLTHROUGH;
+          case TermKind::ReturnInst:
+          case TermKind::ThrowInst:
+          case TermKind::UnwindInst:
+          case TermKind::YieldInst: {
+            // These terminators can never be non-consuming uses of an owned
+            // value since we would be leaking the owned value no matter what
+            // we do. Given that, we can assume that what ever the
+            // non-consuming use actually was, must be before this
+            // instruction. So insert the destroy_value at the end of the
+            // block, before the terminator.
+            auto loc = RegularLocation::getAutoGeneratedLocation();
+            SILBuilderWithScope localBuilder(termInst);
+            localBuilder.createDestroyValue(loc, v);
+            return;
+          }
+          case TermKind::TryApplyInst:
+          case TermKind::SwitchValueInst:
+          case TermKind::SwitchEnumInst:
+          case TermKind::SwitchEnumAddrInst:
+          case TermKind::DynamicMethodBranchInst:
+          case TermKind::AwaitAsyncContinuationInst:
+          case TermKind::CheckedCastBranchInst:
+          case TermKind::CheckedCastAddrBranchInst:
+          case TermKind::CheckedCastValueBranchInst: {
+            // Otherwise, we insert the destroy_addr /after/ the
+            // terminator. All of these are guaranteed to have each successor
+            // to have the block as its only predecessor block.
+            SILBuilderWithScope::insertAfter(termInst, [&](auto &b) {
+              auto loc = RegularLocation::getAutoGeneratedLocation();
+              b.createDestroyValue(loc, v);
+            });
+            return;
+          }
+          }
+          llvm_unreachable("Case that did not return in its body?!");
+        });
+  }
+
   return true;
 }
 
@@ -2687,6 +2937,10 @@ class PredictableMemoryAccessOptimizations : public SILFunctionTransform {
 
 class PredictableDeadAllocationElimination : public SILFunctionTransform {
   void run() override {
+    // If we are already canonical or do not have ownership, just bail.
+    if (getFunction()->wasDeserializedCanonical() ||
+        !getFunction()->hasOwnership())
+      return;
     if (eliminateDeadAllocations(*getFunction()))
       invalidateAnalysis(SILAnalysis::InvalidationKind::FunctionBody);
   }