Merge pull request #83712 from jckarter/stmt-condition-addressable

Fixes for `@addressable` references to variables in `if` and `guard` conditions

Author: jckarter

lib/SILGen/Cleanup.h

@@ -74,6 +74,10 @@ enum class CleanupState {
   PersistentlyActive
 };
 
+inline bool isActiveCleanupState(CleanupState state) {
+  return state >= CleanupState::Active;
+}
+
 llvm::raw_ostream &operator<<(raw_ostream &os, CleanupState state);
 
 class LLVM_LIBRARY_VISIBILITY Cleanup {
@@ -109,7 +113,7 @@ class LLVM_LIBRARY_VISIBILITY Cleanup {
   virtual void setState(SILGenFunction &SGF, CleanupState newState) {
     state = newState;
   }
-  bool isActive() const { return state >= CleanupState::Active; }
+  bool isActive() const { return isActiveCleanupState(state); }
   bool isDead() const { return state == CleanupState::Dead; }
 
   virtual void emit(SILGenFunction &SGF, CleanupLocation loc,

lib/SILGen/SILGenDecl.cpp

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "Cleanup.h"
 #include "Initialization.h"
 #include "LValue.h"
 #include "RValue.h"
@@ -682,10 +683,18 @@ namespace {
 
 static void deallocateAddressable(SILGenFunction &SGF,
                 SILLocation l,
-                const SILGenFunction::VarLoc::AddressableBuffer::State &state) {
-  SGF.B.createEndBorrow(l, state.storeBorrow);
+                const SILGenFunction::AddressableBuffer::State &state,
+                CleanupState cleanupState) {
+  // Only delete the value if the variable was active on this path.
+  // The stack slot exists within the variable scope regardless, so always
+  // needs to be cleaned up.
+  bool active = isActiveCleanupState(cleanupState);
+  if (active) {
+    SGF.B.createEndBorrow(l, state.storeBorrow);
+  }
   SGF.B.createDeallocStack(l, state.allocStack);
-  if (state.reabstraction
+  if (active
+      && state.reabstraction
       && !state.reabstraction->getType().isTrivial(SGF.F)) {
     SGF.B.createDestroyValue(l, state.reabstraction);
   }
@@ -695,28 +704,33 @@ static void deallocateAddressable(SILGenFunction &SGF,
 /// binding.
 class DeallocateLocalVariableAddressableBuffer : public Cleanup {
   ValueDecl *vd;
+  CleanupHandle destroyVarHandle;
 public:
-  DeallocateLocalVariableAddressableBuffer(ValueDecl *vd) : vd(vd) {}
+  DeallocateLocalVariableAddressableBuffer(ValueDecl *vd,
+                                           CleanupHandle destroyVarHandle)
+    : vd(vd), destroyVarHandle(destroyVarHandle) {}
 
   void emit(SILGenFunction &SGF, CleanupLocation l,
             ForUnwind_t forUnwind) override {
     auto addressableBuffer = SGF.getAddressableBufferInfo(vd);
     if (!addressableBuffer) {
       return;
     }
-    auto found = SGF.VarLocs.find(vd);
-    if (found == SGF.VarLocs.end()) {
-      return;
-    }
 
+    // Reflect the state of the variable's original cleanup.
+    CleanupState cleanupState = CleanupState::Active;
+    if (destroyVarHandle.isValid()) {
+      cleanupState = SGF.Cleanups.getCleanup(destroyVarHandle).getState();
+    }
+
     if (auto *state = addressableBuffer->getState()) {
-      // The addressable buffer was forced, so clean it up now.
-      deallocateAddressable(SGF, l, *state);
+      // The addressable buffer was already forced, so clean up now.
+      deallocateAddressable(SGF, l, *state, cleanupState);
     } else {
       // Remember this insert location in case we need to force the addressable
       // buffer later.
       SILInstruction *marker = SGF.B.createTuple(l, {});
-      addressableBuffer->cleanupPoints.emplace_back(marker);
+      addressableBuffer->cleanupPoints.push_back({marker, cleanupState});
     }
   }
 
@@ -824,7 +838,7 @@ class LetValueInitialization : public Initialization {
 
     // If the binding has an addressable buffer forced, it should be cleaned
     // up at this scope.
-    SGF.enterLocalVariableAddressableBufferScope(vd);
+    SGF.enterLocalVariableAddressableBufferScope(vd, DestroyCleanup);
   }
 
   ~LetValueInitialization() override {
@@ -2241,10 +2255,11 @@ void SILGenFunction::destroyLocalVariable(SILLocation silLoc, VarDecl *vd) {
 }
 
 void
-SILGenFunction::enterLocalVariableAddressableBufferScope(VarDecl *decl) {
+SILGenFunction::enterLocalVariableAddressableBufferScope(VarDecl *decl,
+                                                         CleanupHandle destroyCleanup) {
   auto marker = B.createTuple(decl, {});
-  AddressableBuffers[decl] = marker;
-  Cleanups.pushCleanup<DeallocateLocalVariableAddressableBuffer>(decl);
+  AddressableBuffers[decl].insertPoint = marker;
+  Cleanups.pushCleanup<DeallocateLocalVariableAddressableBuffer>(decl, destroyCleanup);
 }
 
 SILValue
@@ -2296,27 +2311,12 @@ SILGenFunction::getLocalVariableAddressableBuffer(VarDecl *decl,
   SILValue reabstraction, allocStack, storeBorrow;
   {
     SavedInsertionPointRAII save(B);
-    SILInstruction *insertPoint = nullptr;
-    // Look through bindings that might alias the original addressable buffer
-    // (such as case block variables, which use an alias variable to represent the
-    // incoming value from all of the case label patterns).
-    VarDecl *origDecl = decl;
-    do {
-      auto bufferIter = AddressableBuffers.find(origDecl);
-      ASSERT(bufferIter != AddressableBuffers.end()
-             && "local variable didn't have an addressability scope set");
-
-      insertPoint = bufferIter->second.getInsertPoint();
-      if (insertPoint) {
-        break;
-      }
+    auto *addressableBuffer = getAddressableBufferInfo(decl);
 
-      origDecl = bufferIter->second.getOriginalForAlias();
-      ASSERT(origDecl && "no insert point or alias for addressable declaration!");
-    } while (true);
-      
-    assert(insertPoint && "didn't find an insertion point for the addressable buffer");
-    B.setInsertionPoint(insertPoint);
+    assert(addressableBuffer
+           && addressableBuffer->insertPoint
+           && "didn't find an insertion point for the addressable buffer");
+    B.setInsertionPoint(addressableBuffer->insertPoint);
     auto allocStackTy = fullyAbstractedTy;
     if (value->getType().isMoveOnlyWrapped()) {
       allocStackTy = allocStackTy.addingMoveOnlyWrapper();
@@ -2366,18 +2366,19 @@ SILGenFunction::getLocalVariableAddressableBuffer(VarDecl *decl,
   // Record the addressable representation.
   auto *addressableBuffer = getAddressableBufferInfo(decl);
   auto *newState
-    = new VarLoc::AddressableBuffer::State(reabstraction, allocStack, storeBorrow);
+    = new AddressableBuffer::State(reabstraction, allocStack, storeBorrow);
   addressableBuffer->stateOrAlias = newState;
 
   // Emit cleanups on any paths where we previously would have cleaned up
   // the addressable representation if it had been forced earlier.
   decltype(addressableBuffer->cleanupPoints) cleanupPoints;
   cleanupPoints.swap(addressableBuffer->cleanupPoints);
 
-  for (SILInstruction *cleanupPoint : cleanupPoints) {
-    SavedInsertionPointRAII insertCleanup(B, cleanupPoint);
-    deallocateAddressable(*this, cleanupPoint->getLoc(), *newState);
-    cleanupPoint->eraseFromParent();
+  for (auto &cleanupPoint : cleanupPoints) {
+    SavedInsertionPointRAII insertCleanup(B, cleanupPoint.inst);
+    deallocateAddressable(*this, cleanupPoint.inst->getLoc(), *newState,
+                          cleanupPoint.state);
+    cleanupPoint.inst->eraseFromParent();
   }
 
   return storeBorrow;
@@ -2415,9 +2416,12 @@ void BlackHoleInitialization::copyOrInitValueInto(SILGenFunction &SGF, SILLocati
   SGF.B.createIgnoredUse(loc, value.getValue());
 }
 
-SILGenFunction::VarLoc::AddressableBuffer::~AddressableBuffer() {
-  for (auto cleanupPoint : cleanupPoints) {
-    cleanupPoint->eraseFromParent();
+SILGenFunction::AddressableBuffer::~AddressableBuffer() {
+  if (insertPoint) {
+    insertPoint->eraseFromParent();
+  }
+  for (auto &cleanupPoint : cleanupPoints) {
+    cleanupPoint.inst->eraseFromParent();
   }
   if (auto state = stateOrAlias.dyn_cast<State*>()) {
     delete state;

lib/SILGen/SILGenFunction.cpp

@@ -2005,19 +2005,16 @@ void SILGenFunction::emitAssignOrInit(SILLocation loc, ManagedValue selfValue,
                        AssignOrInitInst::Unknown);
 }
 
-SILGenFunction::VarLoc::AddressableBuffer *
+SILGenFunction::AddressableBuffer *
 SILGenFunction::getAddressableBufferInfo(ValueDecl *vd) {
   do {
-    auto found = VarLocs.find(vd);
-    if (found == VarLocs.end()) {
-      return nullptr;
-    }
+    auto &found = AddressableBuffers[vd];
 
-    if (auto orig = found->second.addressableBuffer.stateOrAlias
+    if (auto orig = found.stateOrAlias
                       .dyn_cast<VarDecl*>()) {
       vd = orig;
       continue;
     }
-    return &found->second.addressableBuffer;
+    return &found;
   } while (true);
 }