Merge pull request swiftlang#36211 from eeckstein/memory-lifetime

SIL: support store_borrow and partial_apply in memory lifetime verification

Author: eeckstein

docs/SIL.rst

@@ -3547,6 +3547,27 @@ by an `end_borrow`_ instruction. All `load_borrow`_ instructions must be
 paired with exactly one `end_borrow`_ instruction along any path through the
 program. Until `end_borrow`_, it is illegal to invalidate or store to ``%0``.
 
+store_borrow
+````````````
+
+::
+
+  sil-instruction ::= 'store_borrow' sil-value 'to' sil-operand
+
+  store_borrow %0 to %1 : $*T
+  // $T must be a loadable type
+  // %1 must be an alloc_stack $T
+
+Stores the value ``%0`` to a stack location ``%1``, which must be an
+``alloc_stack $T``.
+The stored value is alive until the ``dealloc_stack`` or until another
+``store_borrow`` overwrites the value. During the its lifetime, the stored
+value must not be modified or destroyed.
+The source value ``%0`` is borrowed (i.e. not copied) and it's borrow scope
+must outlive the lifetime of the stored value.
+
+Note: This is the current implementation and the design is not final.
+
 begin_borrow
 ````````````
 

include/swift/SIL/ApplySite.h

@@ -348,6 +348,37 @@ class ApplySite {
     return getSubstCalleeConv().getSILArgumentConvention(calleeArgIdx);
   }
 
+  /// Return the SILArgumentConvention for the given applied argument operand at
+  /// the apply instruction.
+  ///
+  /// For full applies, this is equivalent to `getArgumentConvention`. But for
+  /// a partial_apply, the argument ownership convention at the partial_apply
+  /// instruction itself is different from the argument convention of the callee.
+  /// For details see the partial_apply documentation in SIL.rst.
+  SILArgumentConvention getArgumentOperandConvention(const Operand &oper) const {
+    SILArgumentConvention conv = getArgumentConvention(oper);
+    auto *pai = dyn_cast<PartialApplyInst>(Inst);
+    if (!pai)
+      return conv;
+    switch (conv) {
+    case SILArgumentConvention::Indirect_Inout:
+    case SILArgumentConvention::Indirect_InoutAliasable:
+      return conv;
+    case SILArgumentConvention::Direct_Owned:
+    case SILArgumentConvention::Direct_Unowned:
+    case SILArgumentConvention::Direct_Guaranteed:
+      return pai->isOnStack() ? SILArgumentConvention::Direct_Guaranteed
+                              : SILArgumentConvention::Direct_Owned;
+    case SILArgumentConvention::Indirect_In:
+    case SILArgumentConvention::Indirect_In_Constant:
+    case SILArgumentConvention::Indirect_In_Guaranteed:
+      return pai->isOnStack() ? SILArgumentConvention::Indirect_In_Guaranteed
+                              : SILArgumentConvention::Indirect_In;
+    case SILArgumentConvention::Indirect_Out:
+      llvm_unreachable("partial_apply cannot have an @out operand");
+    }
+  }
+
   /// Return true if 'self' is an applied argument.
   bool hasSelfArgument() const {
     switch (ApplySiteKind(Inst->getKind())) {

lib/SIL/Verifier/MemoryLifetime.cpp

@@ -347,9 +347,11 @@ bool MemoryLocations::analyzeLocationUsesRecursively(SILValue V, unsigned locIdx
         break;
       case SILInstructionKind::LoadInst:
       case SILInstructionKind::StoreInst:
+      case SILInstructionKind::StoreBorrowInst:
       case SILInstructionKind::EndAccessInst:
       case SILInstructionKind::LoadBorrowInst:
       case SILInstructionKind::DestroyAddrInst:
+      case SILInstructionKind::PartialApplyInst:
       case SILInstructionKind::ApplyInst:
       case SILInstructionKind::TryApplyInst:
       case SILInstructionKind::BeginApplyInst:
@@ -607,6 +609,9 @@ class MemoryLifetimeVerifier {
   SILFunction *function;
   MemoryLocations locations;
 
+  /// alloc_stack memory locations which are used for store_borrow.
+  Bits storeBorrowLocations;
+
   /// Returns true if the enum location \p locIdx can be proven to hold a
   /// hold a trivial value (e non-payload case) at \p atInst.
   bool isEnumTrivialAt(int locIdx, SILInstruction *atInst);
@@ -639,6 +644,18 @@ class MemoryLifetimeVerifier {
   /// \p addr, are set in \p bits.
   void requireBitsSet(const Bits &bits, SILValue addr, SILInstruction *where);
 
+  bool isStoreBorrowLocation(SILValue addr) {
+    auto *loc = locations.getLocation(addr);
+    return loc && storeBorrowLocations.anyCommon(loc->subLocations);
+  }
+
+  /// Require that the location of addr is not an alloc_stack used for a
+  /// store_borrow.
+  void requireNoStoreBorrowLocation(SILValue addr, SILInstruction *where);
+
+  /// Register the destination address of a store_borrow as borrowed location.
+  void registerStoreBorrowLocation(SILValue addr);
+
   /// Handles locations of the predecessor's terminator, which are only valid
   /// in \p block.
   /// Example: @out results of try_apply. They are only valid in the
@@ -799,6 +816,21 @@ void MemoryLifetimeVerifier::requireBitsSet(const Bits &bits, SILValue addr,
   }
 }
 
+void MemoryLifetimeVerifier::requireNoStoreBorrowLocation(SILValue addr,
+                                                  SILInstruction *where) {
+  if (isStoreBorrowLocation(addr)) {
+    reportError("store-borrow location cannot be written",
+                locations.getLocation(addr)->selfAndParents.find_first(), where);
+  }
+}
+
+void MemoryLifetimeVerifier::registerStoreBorrowLocation(SILValue addr) {
+  if (auto *loc = locations.getLocation(addr)) {
+    storeBorrowLocations.resize(locations.getNumLocations());
+    storeBorrowLocations |= loc->subLocations;
+  }
+}
+
 void MemoryLifetimeVerifier::initDataflow(MemoryDataflow &dataFlow) {
   // Initialize the entry and exit sets to all-bits-set. Except for the function
   // entry.
@@ -848,6 +880,12 @@ void MemoryLifetimeVerifier::initDataflowInBlock(SILBasicBlock *block,
       case SILInstructionKind::StoreInst:
         state.genBits(cast<StoreInst>(&I)->getDest(), locations);
         break;
+      case SILInstructionKind::StoreBorrowInst: {
+        SILValue destAddr = cast<StoreBorrowInst>(&I)->getDest();
+        state.genBits(destAddr, locations);
+        registerStoreBorrowLocation(destAddr);
+        break;
+      }
       case SILInstructionKind::CopyAddrInst: {
         auto *CAI = cast<CopyAddrInst>(&I);
         if (CAI->isTakeOfSrc())
@@ -870,12 +908,13 @@ void MemoryLifetimeVerifier::initDataflowInBlock(SILBasicBlock *block,
       case SILInstructionKind::DeallocStackInst:
         state.killBits(I.getOperand(0), locations);
         break;
+      case SILInstructionKind::PartialApplyInst:
       case SILInstructionKind::ApplyInst:
       case SILInstructionKind::TryApplyInst: {
-        FullApplySite FAS(&I);
+        ApplySite AS(&I);
         for (Operand &op : I.getAllOperands()) {
-          if (FAS.isArgumentOperand(op)) {
-            setFuncOperandBits(state, op, FAS.getArgumentConvention(op),
+          if (AS.isArgumentOperand(op)) {
+            setFuncOperandBits(state, op, AS.getArgumentOperandConvention(op),
                               isa<TryApplyInst>(&I));
           }
         }
@@ -1017,6 +1056,7 @@ void MemoryLifetimeVerifier::checkBlock(SILBasicBlock *block, Bits &bits) {
         switch (LI->getOwnershipQualifier()) {
           case LoadOwnershipQualifier::Take:
             locations.clearBits(bits, LI->getOperand());
+            requireNoStoreBorrowLocation(LI->getOperand(), &I);
             break;
           case LoadOwnershipQualifier::Copy:
           case LoadOwnershipQualifier::Trivial:
@@ -1042,19 +1082,29 @@ void MemoryLifetimeVerifier::checkBlock(SILBasicBlock *block, Bits &bits) {
           case StoreOwnershipQualifier::Unqualified:
             llvm_unreachable("unqualified store shouldn't be in ownership SIL");
         }
+        requireNoStoreBorrowLocation(SI->getDest(), &I);
+        break;
+      }
+      case SILInstructionKind::StoreBorrowInst: {
+        SILValue destAddr = cast<StoreBorrowInst>(&I)->getDest();
+        locations.setBits(bits, destAddr);
+        registerStoreBorrowLocation(destAddr);
         break;
       }
       case SILInstructionKind::CopyAddrInst: {
         auto *CAI = cast<CopyAddrInst>(&I);
         requireBitsSet(bits, CAI->getSrc(), &I);
-        if (CAI->isTakeOfSrc())
+        if (CAI->isTakeOfSrc()) {
           locations.clearBits(bits, CAI->getSrc());
+          requireNoStoreBorrowLocation(CAI->getSrc(), &I);
+        }
         if (CAI->isInitializationOfDest()) {
           requireBitsClear(bits & nonTrivialLocations, CAI->getDest(), &I);
         } else {
           requireBitsSet(bits | ~nonTrivialLocations, CAI->getDest(), &I);
         }
         locations.setBits(bits, CAI->getDest());
+        requireNoStoreBorrowLocation(CAI->getDest(), &I);
         break;
       }
       case SILInstructionKind::InjectEnumAddrInst: {
@@ -1066,38 +1116,45 @@ void MemoryLifetimeVerifier::checkBlock(SILBasicBlock *block, Bits &bits) {
           requireBitsClear(bits & nonTrivialLocations, IEAI->getOperand(), &I);
           locations.setBits(bits, IEAI->getOperand());
         }
+        requireNoStoreBorrowLocation(IEAI->getOperand(), &I);
         break;
       }
-      case SILInstructionKind::InitEnumDataAddrInst:
-        requireBitsClear(bits, cast<InitEnumDataAddrInst>(&I)->getOperand(), &I);
+      case SILInstructionKind::InitEnumDataAddrInst: {
+        SILValue enumAddr = cast<InitEnumDataAddrInst>(&I)->getOperand();
+        requireBitsClear(bits, enumAddr, &I);
+        requireNoStoreBorrowLocation(enumAddr, &I);
         break;
+      }
       case SILInstructionKind::UncheckedTakeEnumDataAddrInst: {
         // Note that despite the name, unchecked_take_enum_data_addr does _not_
         // "take" the payload of the Swift.Optional enum. This is a terrible
         // hack in SIL.
-        auto *UTEDAI = cast<UncheckedTakeEnumDataAddrInst>(&I);
-        int enumIdx = locations.getLocationIdx(UTEDAI->getOperand());
+        SILValue enumAddr = cast<UncheckedTakeEnumDataAddrInst>(&I)->getOperand();
+        int enumIdx = locations.getLocationIdx(enumAddr);
         if (enumIdx >= 0)
-          requireBitsSet(bits, UTEDAI->getOperand(), &I);
+          requireBitsSet(bits, enumAddr, &I);
+        requireNoStoreBorrowLocation(enumAddr, &I);
         break;
       }
       case SILInstructionKind::DestroyAddrInst: {
         SILValue opVal = cast<DestroyAddrInst>(&I)->getOperand();
         requireBitsSet(bits | ~nonTrivialLocations, opVal, &I);
         locations.clearBits(bits, opVal);
+        requireNoStoreBorrowLocation(opVal, &I);
         break;
       }
       case SILInstructionKind::EndBorrowInst: {
         if (SILValue orig = cast<EndBorrowInst>(&I)->getSingleOriginalValue())
           requireBitsSet(bits, orig, &I);
         break;
       }
+      case SILInstructionKind::PartialApplyInst:
       case SILInstructionKind::ApplyInst:
       case SILInstructionKind::TryApplyInst: {
-        FullApplySite FAS(&I);
+        ApplySite AS(&I);
         for (Operand &op : I.getAllOperands()) {
-          if (FAS.isArgumentOperand(op))
-            checkFuncArgument(bits, op, FAS.getArgumentConvention(op), &I);
+          if (AS.isArgumentOperand(op))
+            checkFuncArgument(bits, op, AS.getArgumentOperandConvention(op), &I);
         }
         break;
       }
@@ -1114,7 +1171,11 @@ void MemoryLifetimeVerifier::checkBlock(SILBasicBlock *block, Bits &bits) {
         break;
       case SILInstructionKind::DeallocStackInst: {
         SILValue opVal = cast<DeallocStackInst>(&I)->getOperand();
-        requireBitsClear(bits & nonTrivialLocations, opVal, &I);
+        if (isStoreBorrowLocation(opVal)) {
+          requireBitsSet(bits, opVal, &I);
+        } else {
+          requireBitsClear(bits & nonTrivialLocations, opVal, &I);
+        }
         // Needed to clear any bits of trivial locations (which are not required
         // to be zero).
         locations.clearBits(bits, opVal);
@@ -1129,6 +1190,9 @@ void MemoryLifetimeVerifier::checkBlock(SILBasicBlock *block, Bits &bits) {
 void MemoryLifetimeVerifier::checkFuncArgument(Bits &bits, Operand &argumentOp,
                          SILArgumentConvention argumentConvention,
                          SILInstruction *applyInst) {
+  if (argumentConvention != SILArgumentConvention::Indirect_In_Guaranteed)
+    requireNoStoreBorrowLocation(argumentOp.get(), applyInst);
+  
   switch (argumentConvention) {
     case SILArgumentConvention::Indirect_In:
     case SILArgumentConvention::Indirect_In_Constant:
@@ -1166,6 +1230,7 @@ void MemoryLifetimeVerifier::verify() {
   }
   // Second step: handle single-block locations.
   locations.handleSingleBlockLocations([this](SILBasicBlock *block) {
+    storeBorrowLocations.clear();
     Bits bits(locations.getNumLocations());
     checkBlock(block, bits);
   });

lib/SIL/Verifier/SILVerifier.cpp

@@ -2193,6 +2193,23 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     }
   }
 
+  void checkStoreBorrowInst(StoreBorrowInst *SI) {
+    require(SI->getSrc()->getType().isObject(),
+            "Can't store from an address source");
+    require(!fnConv.useLoweredAddresses()
+                || SI->getSrc()->getType().isLoadable(*SI->getFunction()),
+            "Can't store a non loadable type");
+    require(SI->getDest()->getType().isAddress(),
+            "Must store to an address dest");
+    requireSameType(SI->getDest()->getType().getObjectType(),
+                    SI->getSrc()->getType(),
+                    "Store operand type and dest type mismatch");
+
+    // Note: This is the current implementation and the design is not final.
+    require(isa<AllocStackInst>(SI->getDest()),
+            "store_borrow destination can only be an alloc_stack");
+  }
+
   void checkAssignInst(AssignInst *AI) {
     SILValue Src = AI->getSrc(), Dest = AI->getDest();
     require(AI->getModule().getStage() == SILStage::Raw,

lib/SILOptimizer/Mandatory/RawSILInstLowering.cpp

@@ -16,6 +16,7 @@
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
+#include "swift/SILOptimizer/Utils/InstOptUtils.h"
 #include "llvm/ADT/Statistic.h"
 
 STATISTIC(numAssignRewritten, "Number of assigns rewritten");
@@ -177,6 +178,7 @@ static void lowerAssignByWrapperInstruction(SILBuilderWithScope &b,
   SILValue dest = inst->getDest();
   SILLocation loc = inst->getLoc();
   SILBuilderWithScope forCleanup(std::next(inst->getIterator()));
+  SingleValueInstruction *closureToDelete = nullptr;
 
   switch (inst->getAssignDestination()) {
     case AssignByWrapperInst::Destination::BackingWrapper: {
@@ -199,8 +201,7 @@ static void lowerAssignByWrapperInstruction(SILBuilderWithScope &b,
           b.createStore(loc, wrappedSrc, dest, StoreOwnershipQualifier::Assign);
         }
       }
-      // TODO: remove the unused setter function, which usually is a dead
-      // partial_apply.
+      closureToDelete = dyn_cast<SingleValueInstruction>(inst->getSetter());
       break;
     }
     case AssignByWrapperInst::Destination::WrappedValue: {
@@ -218,12 +219,13 @@ static void lowerAssignByWrapperInstruction(SILBuilderWithScope &b,
       // nested access violation.
       if (auto *BA = dyn_cast<BeginAccessInst>(dest))
         accessMarkers.push_back(BA);
-      // TODO: remove the unused init function, which usually is a dead
-      // partial_apply.
+      closureToDelete = dyn_cast<SingleValueInstruction>(inst->getInitializer());
       break;
     }
   }
   inst->eraseFromParent();
+  if (closureToDelete)
+    tryDeleteDeadClosure(closureToDelete);
 }
 
 static void deleteDeadAccessMarker(BeginAccessInst *BA) {

lib/SILOptimizer/Transforms/DestroyHoisting.cpp

@@ -350,8 +350,10 @@ void DestroyHoisting::getUsedLocationsOfInst(Bits &bits, SILInstruction *I) {
       break;
     case SILInstructionKind::LoadInst:
     case SILInstructionKind::StoreInst:
+    case SILInstructionKind::StoreBorrowInst:
     case SILInstructionKind::CopyAddrInst:
     case SILInstructionKind::InjectEnumAddrInst:
+    case SILInstructionKind::PartialApplyInst:
     case SILInstructionKind::ApplyInst:
     case SILInstructionKind::TryApplyInst:
     case SILInstructionKind::YieldInst:

test/Reflection/capture_descriptors.sil

@@ -111,10 +111,11 @@ bb0(%t : $*T):
   return %12 : $()
 }
 
-sil [ossa] @generic_caller3 : $@convention(thin) <A, B, C> () -> @owned @callee_guaranteed () -> () {
-bb0:
+sil [ossa] @generic_caller3 : $@convention(thin) <A, B, C> (@owned Optional<@callee_guaranteed @substituted <X, Y> (@in_guaranteed X) -> @out Y for <A, B>>) -> @owned @callee_guaranteed () -> () {
+bb0(%0 : @owned $Optional<@callee_guaranteed @substituted <X, Y> (@in_guaranteed X) -> @out Y for <A, B>>):
   %f = function_ref @generic_callee3 : $@convention(thin) <T, U> (@in_guaranteed T) -> ()
   %t = alloc_stack $Optional<@callee_guaranteed @substituted <X, Y> (@in_guaranteed X) -> @out Y for <A, B>>
+  store %0 to [init] %t : $*Optional<@callee_guaranteed @substituted <X, Y> (@in_guaranteed X) -> @out Y for <A, B>>
   %c = partial_apply [callee_guaranteed] %f<Optional<(A) -> B>, (B, (C) -> Int)>(%t) : $@convention(thin) <T, U> (@in_guaranteed T) -> ()
   dealloc_stack %t : $*Optional<@callee_guaranteed @substituted <X, Y> (@in_guaranteed X) -> @out Y for <A, B>>
   return %c : $@callee_guaranteed () -> ()