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

In order to accommodate case bodies with multiple case labels, the AST
represents the bindings in each pattern as a distinct declaration from
the matching binding in the case body, and SILGen shares the variable
representation between the two declarations. That means that the two
declarations also need to be able to share an addressable representation.
Add an "alias" state to the addressable buffer data structures so that
we can refer back to the original case label var decl when the case body
var decl is brought into scope, so that accesses through either decl
properly force the addressable representation.

Fixes rdar://154543619.

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;
@@ -698,20 +699,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);
     }
   }
 
@@ -2253,7 +2257,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());
@@ -2291,9 +2295,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()) {
@@ -2312,8 +2333,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();
 
@@ -2333,17 +2358,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);
@@ -2390,4 +2413,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
+    }
+}