[silgen] Teach CleanupCloner how to handle OwnedValueWritebackCleanups correctly.

Previously we would just forward the cleanup and create a normal "destroy"
cleanup resulting in early deallocation and use after frees along error paths.

As part of this I also had to tweak how DI recognizes self init writebacks to
not use SILLocations. This is approach is more robust (since we aren't relying
on SourceLocs/have verifiers to make sure we don't violate SIL) and also avoids
issues from the write back store not necessarily have the same SILLocation.

<rdar://problem/59830255>

Author: gottesmm

lib/SILGen/Cleanup.cpp

@@ -220,6 +220,21 @@ void CleanupManager::setCleanupState(CleanupsDepth depth, CleanupState state) {
     popTopDeadCleanups();
 }
 
+std::tuple<Cleanup::Flags, Optional<SILValue>>
+CleanupManager::getFlagsAndWritebackBuffer(CleanupHandle depth) {
+  auto iter = stack.find(depth);
+  assert(iter != stack.end() && "can't change end of cleanups stack");
+  assert(iter->getState() != CleanupState::Dead &&
+         "Trying to get writeback buffer of a dead cleanup?!");
+
+  auto resultFlags = iter->getFlags();
+  Optional<SILValue> result;
+  bool foundValue = iter->getWritebackBuffer([&](SILValue v) { result = v; });
+  (void)foundValue;
+  assert(result.hasValue() == foundValue);
+  return std::make_tuple(resultFlags, result);
+}
+
 void CleanupManager::forwardCleanup(CleanupsDepth handle) {
   auto iter = stack.find(handle);
   assert(iter != stack.end() && "can't change end of cleanups stack");
@@ -344,7 +359,16 @@ void CleanupStateRestorationScope::pop() && { popImpl(); }
 //===----------------------------------------------------------------------===//
 
 CleanupCloner::CleanupCloner(SILGenFunction &SGF, const ManagedValue &mv)
-    : SGF(SGF), hasCleanup(mv.hasCleanup()), isLValue(mv.isLValue()) {}
+    : SGF(SGF), hasCleanup(mv.hasCleanup()), isLValue(mv.isLValue()),
+      writebackBuffer(None) {
+  if (hasCleanup) {
+    auto handle = mv.getCleanup();
+    auto state = SGF.Cleanups.getFlagsAndWritebackBuffer(handle);
+    if (SILValue value = std::get<1>(state).getValueOr(SILValue())) {
+      writebackBuffer = value;
+    }
+  }
+}
 
 CleanupCloner::CleanupCloner(SILGenBuilder &builder, const ManagedValue &mv)
     : CleanupCloner(builder.getSILGenFunction(), mv) {}
@@ -372,6 +396,14 @@ ManagedValue CleanupCloner::clone(SILValue value) const {
     return ManagedValue::forUnmanaged(value);
   }
 
+  if (writebackBuffer.hasValue()) {
+    auto loc = RegularLocation::getAutoGeneratedLocation();
+    auto cleanup =
+        SGF.enterOwnedValueWritebackCleanup(loc, *writebackBuffer, value);
+    return ManagedValue::forExclusivelyBorrowedOwnedObjectRValue(value,
+                                                                 cleanup);
+  }
+
   if (value->getType().isAddress()) {
     return SGF.emitManagedBufferWithCleanup(value);
   }

lib/SILGen/Cleanup.h

@@ -19,6 +19,7 @@
 
 #include "swift/Basic/DiverseStack.h"
 #include "swift/SIL/SILLocation.h"
+#include "swift/SIL/SILValue.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace swift {
@@ -75,10 +76,23 @@ llvm::raw_ostream &operator<<(raw_ostream &os, CleanupState state);
 
 class LLVM_LIBRARY_VISIBILITY Cleanup {
   friend class CleanupManager;
+  friend class CleanupCloner;
 
-  unsigned allocatedSize;
+protected:
+  // A set of flags that categorize the type of cleanup such that it can be
+  // recreated via SILGenFunction methods based on the type of argument input.
+  //
+  // Example: Distinguishing in between @owned cleanups with a writeback buffer
+  // (ExclusiveBorrowCleanup) or ones that involve formal access cleanups.
+  enum class Flags : uint8_t {
+    None = 0,
+    ExclusiveBorrowCleanup = 1,
+  };
 
+private:
   CleanupState state;
+  unsigned allocatedSize : 24;
+  Flags flags : 8;
 
 protected:
   Cleanup() {}
@@ -99,6 +113,16 @@ class LLVM_LIBRARY_VISIBILITY Cleanup {
   virtual void emit(SILGenFunction &SGF, CleanupLocation loc,
                     ForUnwind_t forUnwind) = 0;
   virtual void dump(SILGenFunction &SGF) const = 0;
+
+protected:
+  Flags getFlags() const { return flags; }
+
+  /// Call func passing in the SILValue address that this cleanup will write
+  /// back to if supported and any flags associated with the cleanup. Returns
+  /// false otherwise.
+  virtual bool getWritebackBuffer(function_ref<void(SILValue)> func) {
+    return false;
+  }
 };
 
 /// A cleanup depth is generally used to denote the set of cleanups
@@ -117,6 +141,7 @@ typedef DiverseStackImpl<Cleanup>::stable_iterator CleanupHandle;
 
 class LLVM_LIBRARY_VISIBILITY CleanupManager {
   friend class Scope;
+  friend class CleanupCloner;
 
   SILGenFunction &SGF;
 
@@ -229,7 +254,7 @@ class LLVM_LIBRARY_VISIBILITY CleanupManager {
   /// Set the state of the cleanup at the given depth.
   /// The transition must be non-trivial and legal.
   void setCleanupState(CleanupHandle depth, CleanupState state);
-  
+
   /// True if there are any active cleanups in the scope between the two
   /// cleanup handles.
   bool hasAnyActiveCleanups(CleanupsDepth from, CleanupsDepth to);
@@ -246,6 +271,12 @@ class LLVM_LIBRARY_VISIBILITY CleanupManager {
 
   /// Verify that the given cleanup handle is valid.
   void checkIterator(CleanupHandle handle) const;
+
+private:
+  // Look up the flags and optionally the writeback address associated with the
+  // cleanup at \p depth. If
+  std::tuple<Cleanup::Flags, Optional<SILValue>>
+  getFlagsAndWritebackBuffer(CleanupHandle depth);
 };
 
 /// An RAII object that allows the state of a cleanup to be
@@ -274,10 +305,14 @@ class CleanupStateRestorationScope {
   void popImpl();
 };
 
+/// Extract enough information from a managed value to reliably clone its
+/// cleanup (if it has any) on a newly computed type. This includes modeling
+/// writeback buffers.
 class CleanupCloner {
   SILGenFunction &SGF;
   bool hasCleanup;
   bool isLValue;
+  Optional<SILValue> writebackBuffer;
 
 public:
   CleanupCloner(SILGenFunction &SGF, const ManagedValue &mv);

lib/SILGen/SILGenExpr.cpp

@@ -722,6 +722,7 @@ namespace {
 /// cleanup to take ownership of the value and thus prevent it form being
 /// written back.
 struct OwnedValueWritebackCleanup final : Cleanup {
+  using Flags = Cleanup::Flags;
 
   /// We store our own loc so that we can ensure that DI ignores our writeback.
   SILLocation loc;
@@ -733,6 +734,11 @@ struct OwnedValueWritebackCleanup final : Cleanup {
                              SILValue value)
       : loc(loc), lvalueAddress(lvalueAddress), value(value) {}
 
+  bool getWritebackBuffer(function_ref<void(SILValue)> func) override {
+    func(lvalueAddress);
+    return true;
+  }
+
   void emit(SILGenFunction &SGF, CleanupLocation l, ForUnwind_t forUnwind) override {
     SILValue valueToStore = value;
     SILType lvalueObjTy = lvalueAddress->getType().getObjectType();

lib/SILGen/SILGenFunction.h

@@ -1188,6 +1188,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   CleanupHandle enterDeallocateUninitializedArrayCleanup(SILValue array);
   void emitUninitializedArrayDeallocation(SILLocation loc, SILValue array);
 
+  /// Emit a cleanup for an owned value that should be written back at end of
+  /// scope if the value is not forwarded.
   CleanupHandle enterOwnedValueWritebackCleanup(SILLocation loc,
                                                 SILValue address,
                                                 SILValue newValue);

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp

@@ -1639,6 +1639,7 @@ void ClassInitElementUseCollector::collectClassInitSelfUses() {
         }
 
         // A store of a load from the box is ignored.
+        //
         // SILGen emits these if delegation to another initializer was
         // interrupted before the initializer was called.
         SILValue src = SI->getSrc();
@@ -1778,14 +1779,30 @@ void ClassInitElementUseCollector::collectClassInitSelfLoadUses(
       }
     }
 
-    // If this load's value is being stored back into the delegating
-    // mark_uninitialized buffer and it is a self init use, skip the
-    // use. This is to handle situations where due to usage of a metatype to
-    // allocate, we do not actually consume self.
+    // If this load's value is being stored immediately back into the delegating
+    // mark_uninitialized buffer, skip the use.
+    //
+    // This is to handle situations where we do not actually consume self as a
+    // result of situations such as:
+    //
+    // 1. The usage of a metatype to allocate the object.
+    //
+    // 2. If our self init call has a throwing function as an argument that
+    //    actually throws.
     if (auto *SI = dyn_cast<StoreInst>(User)) {
-      if (SI->getDest() == MUI &&
-          (isSelfInitUse(User) || isSuperInitUse(User))) {
-        continue;
+      if (SI->getDest() == MUI) {
+        SILValue src = SI->getSrc();
+
+        // Look through conversions.
+        while (auto *conversion = dyn_cast<ConversionInst>(src)) {
+          src = conversion->getConverted();
+        }
+
+        if (auto *li = dyn_cast<LoadInst>(src)) {
+          if (li->getOperand() == MUI) {
+            continue;
+          }
+        }
       }
     }
 

lib/SILOptimizer/Mandatory/DefiniteInitialization.cpp

@@ -785,6 +785,7 @@ void LifetimeChecker::doIt() {
       break;
     case DIUseKind::LoadForTypeOfSelf:
       handleLoadForTypeOfSelfUse(Use);
+      break;
     }
   }