Always manage subobject projections with formal-access cleanups

include/swift/Basic/DiverseStack.h

@@ -23,6 +23,7 @@
 #define SWIFT_BASIC_DIVERSESTACK_H
 
 #include "swift/Basic/Malloc.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 #include <cassert>
 #include <cstring>
@@ -293,6 +294,15 @@ template <class T> class DiverseStackImpl : private DiverseStackBase {
     return stable_iterator(End - it.Ptr);
   } 
 
+  T &findAndAdvance(stable_iterator &i) {
+    auto unstable_i = find(i);
+    assert(unstable_i != end());
+    T &value = *unstable_i;
+    ++unstable_i;
+    i = stabilize(unstable_i);
+    return value;
+  }
+
   class const_iterator {
     const char *Ptr;
     friend class DiverseStackImpl;
@@ -377,6 +387,22 @@ template <class T> class DiverseStackImpl : private DiverseStackBase {
   }
 };
 
+/// A helper class for copying value off a DiverseStack.
+template <class T>
+class DiverseValueBuffer {
+  llvm::SmallVector<char, sizeof(T) + 10 * sizeof(void*)> data;
+
+public:
+  DiverseValueBuffer(const T &value) {
+    size_t size = value.allocated_size();
+    data.reserve(size);
+    data.set_size(size);
+    memcpy(data.data(), reinterpret_cast<const void *>(&value), size);
+  }
+
+  T &getCopy() { return *reinterpret_cast<T *>(data.data()); }
+};
+
 } // end namespace swift
 
 /// Allow stable_iterators to be put in things like TinyPtrVectors.

lib/SILGen/Cleanup.cpp

@@ -53,24 +53,6 @@ static bool hasAnyActiveCleanups(DiverseStackImpl<Cleanup>::iterator begin,
   return false;
 }
 
-namespace {
-  /// A CleanupBuffer is a location to which to temporarily copy a
-  /// cleanup.
-  class CleanupBuffer {
-    SmallVector<char, sizeof(Cleanup) + 10 * sizeof(void *)> data;
-
-  public:
-    CleanupBuffer(const Cleanup &cleanup) {
-      size_t size = cleanup.allocated_size();
-      data.reserve(size);
-      data.set_size(size);
-      memcpy(data.data(), reinterpret_cast<const void *>(&cleanup), size);
-    }
-
-    Cleanup &getCopy() { return *reinterpret_cast<Cleanup *>(data.data()); }
-  };
-} // end anonymous namespace
-
 void CleanupManager::popTopDeadCleanups() {
   auto end = (innermostScope ? innermostScope->depth : stack.stable_end());
   assert(end.isValid());
@@ -83,6 +65,8 @@ void CleanupManager::popTopDeadCleanups() {
   }
 }
 
+using CleanupBuffer = DiverseValueBuffer<Cleanup>;
+
 void CleanupManager::popAndEmitCleanup(CleanupHandle handle,
                                        CleanupLocation loc,
                                        ForUnwind_t forUnwind) {
@@ -107,22 +91,30 @@ void CleanupManager::emitCleanups(CleanupsDepth depth, CleanupLocation loc,
   auto topOfStack = cur;
 #endif
   while (cur != depth) {
+    // Advance the iterator.
+    auto cleanupHandle = cur;
+    auto iter = stack.find(cleanupHandle);
+    Cleanup &stackCleanup = *iter;
+    cur = stack.stabilize(++iter);
+
     // Copy the cleanup off the stack if it needs to be emitted.
     // This is necessary both because we might need to pop the cleanup and
     // because the cleanup might push other cleanups that will invalidate
     // references onto the stack.
-    auto iter = stack.find(cur);
-    Cleanup &stackCleanup = *iter;
     Optional<CleanupBuffer> copiedCleanup;
     if (stackCleanup.isActive() && SGF.B.hasValidInsertionPoint()) {
       copiedCleanup.emplace(stackCleanup);
     }
 
-    // Advance the iterator.
-    cur = stack.stabilize(++iter);
-
     // Pop now if that was requested.
     if (popCleanups) {
+#ifndef NDEBUG
+      // This is an implicit deactivation.
+      if (stackCleanup.isActive()) {
+        SGF.FormalEvalContext.checkCleanupDeactivation(cleanupHandle);
+      }
+#endif
+
       stack.pop();
 
 #ifndef NDEBUG

lib/SILGen/Cleanup.h

@@ -165,6 +165,10 @@ class LLVM_LIBRARY_VISIBILITY CleanupManager {
     return *stack.find(iter);
   }
 
+  Cleanup &findAndAdvance(CleanupsDepth &iter) {
+    return stack.findAndAdvance(iter);
+  }
+
   /// \brief Emit a branch to the given jump destination,
   /// threading out through any cleanups we need to run. This does not pop the
   /// cleanup stack.

lib/SILGen/FormalEvaluation.cpp

@@ -106,25 +106,28 @@ void FormalEvaluationScope::popImpl() {
          "popping formal-evaluation scopes out of order");
   context.innermostScope = previous;
 
-  // Check to see if there is anything going on here.
-
-  using iterator = FormalEvaluationContext::iterator;
-  using stable_iterator = FormalEvaluationContext::stable_iterator;
+  auto endDepth = *savedDepth;
 
-  iterator unwrappedSavedDepth = context.find(savedDepth.getValue());
-  iterator iter = context.begin();
-  if (iter == unwrappedSavedDepth)
+  // Check to see if there is anything going on here.
+  if (endDepth == context.stable_begin())
     return;
 
+#ifndef NDEBUG
+  // Verify that all the accesses are valid.
+  for (auto i = context.begin(), e = context.find(endDepth); i != e; ++i) {
+    i->verify(SGF);
+  }
+#endif
+
   // Save our start point to make sure that we are not adding any new cleanups
   // to the front of the stack.
-  stable_iterator originalBegin = context.stable_begin();
-  (void)originalBegin;
+  auto originalBegin = context.stable_begin();
 
   // Then working down the stack until we visit unwrappedSavedDepth...
-  for (; iter != unwrappedSavedDepth; ++iter) {
-    // Grab the next evaluation...
-    FormalAccess &access = *iter;
+  auto i = originalBegin;
+  do {
+    // Grab the next evaluation.
+    FormalAccess &access = context.findAndAdvance(i);
 
     // If this access was already finished, continue. This can happen if an
     // owned formal access was forwarded.
@@ -149,10 +152,9 @@ void FormalEvaluationScope::popImpl() {
     // code. We do a simple N^2 comparison here to detect this because it is
     // extremely unlikely more than a few writebacks are active at once.
     if (access.getKind() == FormalAccess::Exclusive) {
-      iterator j = iter;
-      ++j;
-
-      for (; j != unwrappedSavedDepth; ++j) {
+      // Note that we already advanced 'iter' above, so we can just start
+      // iterating from there.  Also, this doesn't invalidate the iterators.
+      for (auto j = context.find(i), je = context.find(endDepth); j != je; ++j){
         FormalAccess &other = *j;
         if (other.getKind() != FormalAccess::Exclusive)
           continue;
@@ -167,33 +169,25 @@ void FormalEvaluationScope::popImpl() {
     //
     // This evaluates arbitrary code, so it's best to be paranoid
     // about iterators on the context.
-    access.finish(SGF);
-  }
+    DiverseValueBuffer<FormalAccess> copiedAccess(access);
+    copiedAccess.getCopy().finish(SGF);
+
+  } while (i != endDepth);
 
   // Then check that we did not add any additional cleanups to the beginning of
   // the stack...
   assert(originalBegin == context.stable_begin() &&
-         "more formal eval cleanups placed onto context during formal eval scope pop?!");
+         "pushed more formal evaluations while popping formal evaluations?!");
 
   // And then pop off all stack elements until we reach the savedDepth.
-  context.pop(savedDepth.getValue());
+  context.pop(endDepth);
 }
 
 void FormalEvaluationScope::verify() const {
-  // Check to see if there is anything going on here.
+  // Walk up the stack to the saved depth.
   auto &context = SGF.FormalEvalContext;
-  using iterator = FormalEvaluationContext::iterator;
-
-  iterator unwrappedSavedDepth = context.find(savedDepth.getValue());
-  iterator iter = context.begin();
-  if (iter == unwrappedSavedDepth)
-    return;
-
-  // Then working down the stack until we visit unwrappedSavedDepth...
-  for (; iter != unwrappedSavedDepth; ++iter) {
-    // Grab the next evaluation verify that we can successfully access this
-    // formal access.
-    (*iter).verify(SGF);
+  for (auto i = context.begin(), e = context.find(*savedDepth); i != e; ++i) {
+    i->verify(SGF);
   }
 }
 

lib/SILGen/FormalEvaluation.h

@@ -139,6 +139,10 @@ class FormalEvaluationContext {
   stable_iterator stable_begin() { return stabilize(begin()); }
   iterator find(stable_iterator iter) { return stack.find(iter); }
 
+  FormalAccess &findAndAdvance(stable_iterator &stable) {
+    return stack.findAndAdvance(stable);
+  }
+
   template <class U, class... ArgTypes> void push(ArgTypes &&... args) {
     stack.push<U>(std::forward<ArgTypes>(args)...);
   }
@@ -236,12 +240,7 @@ class FormalEvaluationScope {
 #ifndef NDEBUG
 inline void
 FormalEvaluationContext::checkCleanupDeactivation(CleanupHandle handle) {
-  // Start at the innermost scope depth.  Note that we pop scopes off the
-  // stack before we start emitting their cleanups.
-  if (!innermostScope) return;
-  assert(!innermostScope->isPopped());
-  for (auto i = find(*innermostScope->savedDepth), e = end(); i != e; ++i) {
-    auto &access = *i;
+  for (auto &access : *this) {
     assert((access.isFinished() || access.getCleanup() != handle) &&
            "popping active formal-evaluation cleanup");
   }

lib/SILGen/ManagedValue.cpp

@@ -85,6 +85,8 @@ ManagedValue ManagedValue::copyUnmanaged(SILGenFunction &SGF, SILLocation loc) {
 /// have cleanups.  It returns a +1 value with one.
 ManagedValue ManagedValue::formalAccessCopyUnmanaged(SILGenFunction &SGF,
                                                      SILLocation loc) {
+  assert(SGF.isInFormalEvaluationScope());
+
   if (getType().isObject()) {
     return SGF.B.createFormalAccessCopyValue(loc, *this);
   }
@@ -141,6 +143,7 @@ ManagedValue ManagedValue::borrow(SILGenFunction &SGF, SILLocation loc) const {
 
 ManagedValue ManagedValue::formalAccessBorrow(SILGenFunction &SGF,
                                               SILLocation loc) const {
+  assert(SGF.isInFormalEvaluationScope());
   assert(getValue() && "cannot borrow an invalid or in-context value");
   if (isLValue())
     return *this;