Merge pull request #83009 from jckarter/case-binding-addressable-scopes

SILGen: Properly set up addressability scopes for case pattern bindings.

Author: jckarter

lib/SILGen/SILGenDecl.cpp

@@ -36,6 +36,7 @@
 #include "swift/SIL/SILType.h"
 #include "swift/SIL/TypeLowering.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/Support/ErrorHandling.h"
 #include <iterator>
 
 using namespace swift;
@@ -699,20 +700,23 @@ class DeallocateLocalVariableAddressableBuffer : public Cleanup {
 
   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;
     }
-    auto &loc = found->second;
 
-    if (auto &state = loc.addressableBuffer.state) {
+    if (auto *state = addressableBuffer->getState()) {
       // The addressable buffer was forced, so clean it up now.
       deallocateAddressable(SGF, l, *state);
     } else {
       // Remember this insert location in case we need to force the addressable
       // buffer later.
       SILInstruction *marker = SGF.B.createTuple(l, {});
-      loc.addressableBuffer.cleanupPoints.emplace_back(marker);
+      addressableBuffer->cleanupPoints.emplace_back(marker);
     }
   }
 
@@ -2254,7 +2258,7 @@ SILGenFunction::getLocalVariableAddressableBuffer(VarDecl *decl,
 
   auto value = foundVarLoc->second.value;
   auto access = foundVarLoc->second.access;
-  auto *state = foundVarLoc->second.addressableBuffer.state.get();
+  auto *state = getAddressableBufferInfo(decl)->getState();
 
   SILType fullyAbstractedTy = getLoweredType(AbstractionPattern::getOpaque(),
                                      decl->getTypeInContext()->getRValueType());
@@ -2292,9 +2296,26 @@ SILGenFunction::getLocalVariableAddressableBuffer(VarDecl *decl,
   SILValue reabstraction, allocStack, storeBorrow;
   {
     SavedInsertionPointRAII save(B);
-    ASSERT(AddressableBuffers.find(decl) != AddressableBuffers.end()
-           && "local variable did not have an addressability scope set");
-    auto insertPoint = AddressableBuffers[decl].insertPoint;
+    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;
+      }
+
+      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);
     auto allocStackTy = fullyAbstractedTy;
     if (value->getType().isMoveOnlyWrapped()) {
@@ -2313,8 +2334,12 @@ SILGenFunction::getLocalVariableAddressableBuffer(VarDecl *decl,
     SavedInsertionPointRAII save(B);
     if (isa<ParamDecl>(decl)) {
       B.setInsertionPoint(allocStack->getNextInstruction());
+    } else if (auto inst = value->getDefiningInstruction()) {
+      B.setInsertionPoint(inst->getParent(), std::next(inst->getIterator()));
+    } else if (auto arg = dyn_cast<SILArgument>(value)) {
+      B.setInsertionPoint(arg->getParent()->begin());
     } else {
-      B.setInsertionPoint(value->getNextInstruction());
+      llvm_unreachable("unexpected value source!");
     }
     auto declarationLoc = value->getDefiningInsertionPoint()->getLoc();
 
@@ -2334,17 +2359,15 @@ SILGenFunction::getLocalVariableAddressableBuffer(VarDecl *decl,
   }
 
   // Record the addressable representation.
-  auto &addressableBuffer = VarLocs[decl].addressableBuffer;
-  addressableBuffer.state
-    = std::make_unique<VarLoc::AddressableBuffer::State>(reabstraction,
-                                                         allocStack,
-                                                         storeBorrow);
-  auto *newState = addressableBuffer.state.get();
+  auto *addressableBuffer = getAddressableBufferInfo(decl);
+  auto *newState
+    = new VarLoc::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);
+  decltype(addressableBuffer->cleanupPoints) cleanupPoints;
+  cleanupPoints.swap(addressableBuffer->cleanupPoints);
 
   for (SILInstruction *cleanupPoint : cleanupPoints) {
     SavedInsertionPointRAII insertCleanup(B, cleanupPoint);
@@ -2391,4 +2414,7 @@ SILGenFunction::VarLoc::AddressableBuffer::~AddressableBuffer() {
   for (auto cleanupPoint : cleanupPoints) {
     cleanupPoint->eraseFromParent();
   }
+  if (auto state = stateOrAlias.dyn_cast<State*>()) {
+    delete state;
+  }
 }

lib/SILGen/SILGenFunction.cpp

@@ -2004,3 +2004,20 @@ void SILGenFunction::emitAssignOrInit(SILLocation loc, ManagedValue selfValue,
                        newValue.forward(*this), initFRef, setterFRef,
                        AssignOrInitInst::Unknown);
 }
+
+SILGenFunction::VarLoc::AddressableBuffer *
+SILGenFunction::getAddressableBufferInfo(ValueDecl *vd) {
+  do {
+    auto found = VarLocs.find(vd);
+    if (found == VarLocs.end()) {
+      return nullptr;
+    }
+
+    if (auto orig = found->second.addressableBuffer.stateOrAlias
+                      .dyn_cast<VarDecl*>()) {
+      vd = orig;
+      continue;
+    }
+    return &found->second.addressableBuffer;
+  } while (true);
+}

lib/SILGen/SILGenFunction.h

@@ -27,6 +27,7 @@
 #include "swift/Basic/ProfileCounter.h"
 #include "swift/Basic/Statistic.h"
 #include "swift/SIL/SILBuilder.h"
+#include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILType.h"
 #include "llvm/ADT/PointerIntPair.h"
 
@@ -497,27 +498,42 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
         {}
       };
 
-      std::unique_ptr<State> state = nullptr;
-      
+      llvm::PointerUnion<State *, VarDecl*> stateOrAlias = (State*)nullptr;
+
       // If the variable cleanup is triggered before the addressable
       // representation is demanded, but the addressable representation
       // gets demanded later, we save the insertion points where the
       // representation would be cleaned up so we can backfill them.
       llvm::SmallVector<SILInstruction*, 1> cleanupPoints;
 
       AddressableBuffer() = default;
+
+      AddressableBuffer(VarDecl *original)
+        : stateOrAlias(original)
+      {
+      }
 
       AddressableBuffer(AddressableBuffer &&other)
-        : state(std::move(other.state))
+        : stateOrAlias(other.stateOrAlias)
       {
+        other.stateOrAlias = (State*)nullptr;
         cleanupPoints.swap(other.cleanupPoints);
       }
 
       AddressableBuffer &operator=(AddressableBuffer &&other) {
-        state = std::move(other.state);
+        if (auto state = stateOrAlias.dyn_cast<State*>()) {
+          delete state;
+        }
+        stateOrAlias = other.stateOrAlias;
         cleanupPoints.swap(other.cleanupPoints);
         return *this;
       }
+
+      State *getState() {
+        ASSERT(!stateOrAlias.is<VarDecl*>()
+               && "must get state from original AddressableBuffer");
+        return stateOrAlias.dyn_cast<State*>();
+      }
 
       ~AddressableBuffer();
     };
@@ -535,33 +551,50 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   /// emitted. The map is queried to produce the lvalue for a DeclRefExpr to
   /// a local variable.
   llvm::DenseMap<ValueDecl*, VarLoc> VarLocs;
+
+  VarLoc::AddressableBuffer *getAddressableBufferInfo(ValueDecl *vd);
 
   // Represents an addressable buffer that has been allocated but not yet used.
   struct PreparedAddressableBuffer {
-    SILInstruction *insertPoint = nullptr;
+    llvm::PointerUnion<SILInstruction *, VarDecl *> insertPointOrAlias
+      = (SILInstruction*)nullptr;
 
     PreparedAddressableBuffer() = default;
 
     PreparedAddressableBuffer(SILInstruction *insertPoint)
-      : insertPoint(insertPoint)
+      : insertPointOrAlias(insertPoint)
     {
       ASSERT(insertPoint && "null insertion point provided");
     }
+
+    PreparedAddressableBuffer(VarDecl *alias)
+      : insertPointOrAlias(alias)
+    {
+      ASSERT(alias && "null alias provided");
+    }
 
     PreparedAddressableBuffer(PreparedAddressableBuffer &&other)
-      : insertPoint(other.insertPoint)
+      : insertPointOrAlias(other.insertPointOrAlias)
     {
-      other.insertPoint = nullptr;
+      other.insertPointOrAlias = (SILInstruction*)nullptr;
     }
 
     PreparedAddressableBuffer &operator=(PreparedAddressableBuffer &&other) {
-      insertPoint = other.insertPoint;
-      other.insertPoint = nullptr;
+      insertPointOrAlias = other.insertPointOrAlias;
+      other.insertPointOrAlias = nullptr;
       return *this;
     }
+
+    SILInstruction *getInsertPoint() const {
+      return insertPointOrAlias.dyn_cast<SILInstruction*>();
+    }
 
+    VarDecl *getOriginalForAlias() const {
+      return insertPointOrAlias.dyn_cast<VarDecl*>();
+    }
+
     ~PreparedAddressableBuffer() {
-      if (insertPoint) {
+      if (auto insertPoint = getInsertPoint()) {
         // Remove the insertion point if it went unused.
         insertPoint->eraseFromParent();
       }

lib/SILGen/SILGenPattern.cpp

@@ -24,7 +24,6 @@
 #include "swift/AST/SILOptions.h"
 #include "swift/AST/SubstitutionMap.h"
 #include "swift/AST/Types.h"
-#include "swift/Basic/Assertions.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/ProfileCounter.h"
 #include "swift/Basic/STLExtras.h"
@@ -241,7 +240,7 @@ static bool isWildcardPattern(const Pattern *p) {
 
 /// Check to see if the given pattern is a specializing pattern,
 /// and return a semantic pattern for it.
-Pattern *getSpecializingPattern(Pattern *p) {
+static Pattern *getSpecializingPattern(Pattern *p) {
   // Empty entries are basically AnyPatterns.
   if (!p) return nullptr;
 
@@ -975,9 +974,8 @@ class SpecializedArgForwarder : private ArgForwarderBase {
     if (IsFinalUse) {
       ArgForwarderBase::forwardIntoIrrefutable(value);
       return value;
-    } else {
-      return ArgForwarderBase::forward(value, loc);
     }
+    return ArgForwarderBase::forward(value, loc);
   }
 };
 
@@ -3175,6 +3173,9 @@ static void switchCaseStmtSuccessCallback(SILGenFunction &SGF,
           expectedLoc = SILGenFunction::VarLoc(vdLoc->second.value,
                                                vdLoc->second.access,
                                                vdLoc->second.box);
+          expectedLoc.addressableBuffer = vd;
+          // Alias the addressable buffer for the two variables.
+          SGF.AddressableBuffers[expected] = vd;
 
           // Emit a debug description for the variable, nested within a scope
           // for the pattern match.

test/SILGen/addressable_params.swift

@@ -46,3 +46,21 @@ struct Foo {
         }
     }
 }
+
+enum TestEnum {
+    case foo(String)
+    case bar(String)
+}
+
+func addressableParam(_: @_addressable String) -> Bool { true }
+
+func testAddressableSwitchBinding(e: TestEnum) -> Bool {
+    return switch e {
+    case .foo(let f) where addressableParam(f):
+        true
+    case .bar(let b):
+        addressableParam(b)
+    default:
+        false
+    }
+}