Allow partial consumption of self in deinit.

Basic support for partial consumption in deinit without full support for
all future uses of drop_deinit.

Author: nate-chandler

include/swift/SIL/AddressWalker.h

@@ -30,6 +30,12 @@
 
 namespace swift {
 
+enum class TransitiveAddressWalkerTransitiveUseVisitation : uint8_t {
+  OnlyUses,
+  OnlyUser,
+  BothUserAndUses,
+};
+
 /// A state structure for findTransitiveUsesForAddress. Intended to be only used
 /// a single time. Please always use a new one for every call to
 /// findTransitiveUsesForAddress.
@@ -61,11 +67,16 @@ class TransitiveAddressWalker {
   /// understand. These cause us to return AddressUseKind::Unknown.
   void onError(Operand *use) {}
 
+  using TransitiveUseVisitation =
+      TransitiveAddressWalkerTransitiveUseVisitation;
+
   /// Customization point that causes the walker to treat a specific transitive
   /// use as an end point use.
   ///
   /// Example: Visiting a mutable or immutable open_existential_addr.
-  bool visitTransitiveUseAsEndPointUse(Operand *use) { return false; }
+  TransitiveUseVisitation visitTransitiveUseAsEndPointUse(Operand *use) {
+    return TransitiveUseVisitation::OnlyUses;
+  }
 
   void meet(AddressUseKind other) {
     assert(!didInvalidate);
@@ -112,8 +123,12 @@ TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) && {
       return callVisitUse(use);
     }
 
-    if (asImpl().visitTransitiveUseAsEndPointUse(use))
-      return callVisitUse(use);
+    auto visitation = asImpl().visitTransitiveUseAsEndPointUse(use);
+    if (visitation != TransitiveUseVisitation::OnlyUses)
+      callVisitUse(use);
+
+    if (visitation == TransitiveUseVisitation::OnlyUser)
+      return;
 
     for (auto *use : svi->getUses())
       addToWorklist(use);

include/swift/SIL/FieldSensitivePrunedLiveness.h

@@ -24,6 +24,7 @@
 #include "swift/Basic/Debug.h"
 #include "swift/Basic/FrozenMultiMap.h"
 #include "swift/Basic/STLExtras.h"
+#include "swift/SIL/ApplySite.h"
 #include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
@@ -254,9 +255,7 @@ struct TypeSubElementCount {
   TypeSubElementCount(SILType type, SILFunction *fn)
       : TypeSubElementCount(type, fn->getModule(), TypeExpansionContext(*fn)) {}
 
-  TypeSubElementCount(SILValue value)
-      : TypeSubElementCount(value->getType(), *value->getModule(),
-                            TypeExpansionContext(*value->getFunction())) {}
+  TypeSubElementCount(SILValue value);
 
   operator unsigned() const { return number; }
 
@@ -304,6 +303,39 @@ struct TypeTreeLeafTypeRange {
              *startEltOffset + TypeSubElementCount(projectedValue)}};
   }
 
+  /// Which bits of \p rootValue are involved in \p op.
+  ///
+  /// This is a subset of (usually equal to) the bits of op->getType() in \p
+  /// rootValue.
+  static std::optional<TypeTreeLeafTypeRange> get(Operand *op,
+                                                  SILValue rootValue) {
+    auto projectedValue = op->get();
+    auto startEltOffset = SubElementOffset::compute(projectedValue, rootValue);
+    if (!startEltOffset)
+      return std::nullopt;
+
+    // A drop_deinit only consumes the deinit bit of its operand.
+    auto *ddi = dyn_cast<DropDeinitInst>(op->getUser());
+    if (ddi) {
+      auto upperBound = *startEltOffset + TypeSubElementCount(projectedValue);
+      return {{upperBound - 1, upperBound}};
+    }
+
+    // Uses that borrow a value do not involve the deinit bit.
+    //
+    // FIXME: This shouldn't be limited to applies.
+    unsigned deinitBitOffset = 0;
+    if (op->get()->getType().isValueTypeWithDeinit() &&
+        op->getOperandOwnership() == OperandOwnership::Borrow &&
+        ApplySite::isa(op->getUser())) {
+      deinitBitOffset = 1;
+    }
+
+    return {{*startEltOffset, *startEltOffset +
+                                  TypeSubElementCount(projectedValue) -
+                                  deinitBitOffset}};
+  }
+
   /// Given a type \p rootType and a set of needed elements specified by the bit
   /// vector \p neededElements, place into \p foundContiguousTypeRanges a set of
   /// TypeTreeLeafTypeRanges that are associated with the bit vectors

lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp

@@ -104,6 +104,18 @@ TypeSubElementCount::TypeSubElementCount(SILType type, SILModule &mod,
   // our default value, so we can just return.
 }
 
+TypeSubElementCount::TypeSubElementCount(SILValue value) : number(1) {
+  auto whole = TypeSubElementCount(value->getType(), *value->getModule(),
+                                   TypeExpansionContext(*value->getFunction()));
+  // The value produced by a drop_deinit has one fewer subelement than that of
+  // its type--the deinit bit is not included.
+  if (isa<DropDeinitInst>(value)) {
+    assert(value->getType().isValueTypeWithDeinit());
+    whole = whole - 1;
+  }
+  number = whole;
+}
+
 //===----------------------------------------------------------------------===//
 //                           MARK: SubElementNumber
 //===----------------------------------------------------------------------===//
@@ -201,7 +213,14 @@ SubElementOffset::computeForAddress(SILValue projectionDerivedFromRoot,
       projectionDerivedFromRoot = initData->getOperand();
       continue;
     }
-    
+
+    // A drop_deinit consumes the "self" bit at the end of its type.  The offset
+    // is still to the beginning.
+    if (auto dd = dyn_cast<DropDeinitInst>(projectionDerivedFromRoot)) {
+      projectionDerivedFromRoot = dd->getOperand();
+      continue;
+    }
+
     // Look through wrappers.
     if (auto c2m = dyn_cast<CopyableToMoveOnlyWrapperAddrInst>(projectionDerivedFromRoot)) {
       projectionDerivedFromRoot = c2m->getOperand();

lib/SILGen/SILGenDestructor.cpp

@@ -264,6 +264,18 @@ void SILGenFunction::emitDeallocatingMoveOnlyDestructor(DestructorDecl *dd) {
                                 dd->getDeclContext()->getSelfNominalTypeDecl(),
                                 cleanupLoc);
 
+  if (getASTContext().LangOpts.hasFeature(
+          Feature::MoveOnlyPartialConsumption)) {
+    if (auto *ddi = dyn_cast<DropDeinitInst>(selfValue)) {
+      if (auto *mu =
+              dyn_cast<MarkUnresolvedNonCopyableValueInst>(ddi->getOperand())) {
+        if (auto *asi = dyn_cast<AllocStackInst>(mu->getOperand())) {
+          B.createDeallocStack(loc, asi);
+        }
+      }
+    }
+  }
+
   // Return.
   B.createReturn(loc, emitEmptyTuple(loc));
 }
@@ -506,15 +518,19 @@ void SILGenFunction::emitClassMemberDestruction(ManagedValue selfValue,
 void SILGenFunction::emitMoveOnlyMemberDestruction(SILValue selfValue,
                                                    NominalTypeDecl *nom,
                                                    CleanupLocation cleanupLoc) {
-  // drop_deinit invalidates any user-defined struct/enum deinit
-  // before the individual members are destroyed.
-  selfValue = B.createDropDeinit(cleanupLoc, selfValue);
+  if (!isa<DropDeinitInst>(selfValue)) {
+    // drop_deinit invalidates any user-defined struct/enum deinit
+    // before the individual members are destroyed.
+    selfValue = B.createDropDeinit(cleanupLoc, selfValue);
+  }
   if (selfValue->getType().isObject()) {
     // A destroy value that uses the result of a drop_deinit implicitly performs
     // memberwise destruction.
     B.emitDestroyValueOperation(cleanupLoc, selfValue);
     return;
   }
+  // self has been stored into a temporary
+  assert(!selfValue->getType().isObject());
   if (auto *structDecl = dyn_cast<StructDecl>(nom)) {
     for (VarDecl *vd : nom->getStoredProperties()) {
       const TypeLowering &ti = getTypeLowering(vd->getTypeInContext());

lib/SILGen/SILGenLValue.cpp

@@ -3509,13 +3509,26 @@ RValue SILGenFunction::emitRValueForNonMemberVarDecl(SILLocation loc,
     SILValue accessAddr = UnenforcedFormalAccess::enter(*this, loc, destAddr,
                                                         SILAccessKind::Read);
 
+    auto isEffectivelyMarkUnresolvedInst = [](auto *inst) -> bool {
+      if (!inst)
+        return false;
+      if (isa<MarkUnresolvedNonCopyableValueInst>(inst))
+        return true;
+      auto *ddi = dyn_cast<DropDeinitInst>(inst);
+      if (!ddi)
+        return false;
+      return isa<MarkUnresolvedNonCopyableValueInst>(ddi->getOperand());
+    };
+
     if (accessAddr->getType().isMoveOnly() &&
-        !isa<MarkUnresolvedNonCopyableValueInst>(accessAddr)) {
+        !isEffectivelyMarkUnresolvedInst(
+            accessAddr->getDefiningInstruction())) {
+      auto kind =
+          MarkUnresolvedNonCopyableValueInst::CheckKind::NoConsumeOrAssign;
       // When loading an rvalue, we should never need to modify the place
       // we're loading from.
-      accessAddr = B.createMarkUnresolvedNonCopyableValueInst(
-          loc, accessAddr,
-          MarkUnresolvedNonCopyableValueInst::CheckKind::NoConsumeOrAssign);
+      accessAddr =
+          B.createMarkUnresolvedNonCopyableValueInst(loc, accessAddr, kind);
     }
 
     auto propagateRValuePastAccess = [&](RValue &&rvalue) {

lib/SILGen/SILGenProlog.cpp

@@ -47,28 +47,48 @@ SILValue SILGenFunction::emitSelfDeclForDestructor(VarDecl *selfDecl) {
   selfType = F.mapTypeIntoContext(selfType);
   SILValue selfValue = F.begin()->createFunctionArgument(selfType, selfDecl);
 
+  uint16_t ArgNo = 1; // Hardcoded for destructors.
+  auto dv = SILDebugVariable(selfDecl->isLet(), ArgNo);
+
   // If we have a move only type, then mark it with
   // mark_unresolved_non_copyable_value so we can't escape it.
-  if (selfType.isMoveOnly()) {
-    // For now, we do not handle move only class deinits. This is because we
-    // need to do a bit more refactoring to handle the weird way that it deals
-    // with ownership. But for simple move only deinits (like struct/enum), that
-    // are owned, lets mark them as needing to be no implicit copy checked so
-    // they cannot escape.
-    if (selfValue->getOwnershipKind() == OwnershipKind::Owned) {
-      selfValue = B.createMarkUnresolvedNonCopyableValueInst(
-          selfDecl, selfValue,
+  //
+  // For now, we do not handle move only class deinits. This is because we need
+  // to do a bit more refactoring to handle the weird way that it deals with
+  // ownership. But for simple move only deinits (like struct/enum), that are
+  // owned, lets mark them as needing to be no implicit copy checked so they
+  // cannot escape.
+  if (selfType.isMoveOnly() && !selfType.isAnyClassReferenceType()) {
+    if (getASTContext().LangOpts.hasFeature(
+            Feature::MoveOnlyPartialConsumption)) {
+      SILValue addr = B.createAllocStack(selfDecl, selfValue->getType(), dv);
+      addr = B.createMarkUnresolvedNonCopyableValueInst(
+          selfDecl, addr,
           MarkUnresolvedNonCopyableValueInst::CheckKind::
               ConsumableAndAssignable);
+      if (selfValue->getType().isObject()) {
+        B.createStore(selfDecl, selfValue, addr, StoreOwnershipQualifier::Init);
+      } else {
+        B.createCopyAddr(selfDecl, selfValue, addr, IsTake, IsInitialization);
+      }
+      // drop_deinit invalidates any user-defined struct/enum deinit
+      // before the individual members are destroyed.
+      addr = B.createDropDeinit(selfDecl, addr);
+      selfValue = addr;
+    } else {
+      if (selfValue->getOwnershipKind() == OwnershipKind::Owned) {
+        selfValue = B.createMarkUnresolvedNonCopyableValueInst(
+            selfDecl, selfValue,
+            MarkUnresolvedNonCopyableValueInst::CheckKind::
+                ConsumableAndAssignable);
+      }
     }
   }
 
   VarLocs[selfDecl] = VarLoc::get(selfValue);
   SILLocation PrologueLoc(selfDecl);
   PrologueLoc.markAsPrologue();
-  uint16_t ArgNo = 1; // Hardcoded for destructors.
-  B.createDebugValue(PrologueLoc, selfValue,
-                     SILDebugVariable(selfDecl->isLet(), ArgNo));
+  B.createDebugValue(PrologueLoc, selfValue, dv);
   return selfValue;
 }
 

lib/SILOptimizer/Mandatory/AccessEnforcementSelection.cpp

@@ -711,6 +711,10 @@ SourceAccess AccessEnforcementSelection::getSourceAccess(SILValue address) {
   if (auto *m = dyn_cast<MoveOnlyWrapperToCopyableAddrInst>(address))
     return getSourceAccess(m->getOperand());
 
+  // Recurse through drop_deinit.
+  if (auto *ddi = dyn_cast<DropDeinitInst>(address))
+    return getSourceAccess(ddi->getOperand());
+
   // Recurse through moveonlywrapper_to_copyable_box.
   if (auto *m = dyn_cast<MoveOnlyWrapperToCopyableBoxInst>(address))
     return getAccessKindForBox(m->getOperand());