Merge pull request #73159 from nate-chandler/cherrypick/release/6.0/rdar126715654

6.0: [NoncopyablePartialConsumption] Promote to feature.

Author: nate-chandler

include/swift/Basic/Features.def

@@ -175,6 +175,7 @@ LANGUAGE_FEATURE(ExpressionMacroDefaultArguments, 422, "Expression macro as call
 LANGUAGE_FEATURE(BuiltinStoreRaw, 0, "Builtin.storeRaw")
 LANGUAGE_FEATURE(BuiltinCreateTask, 0, "Builtin.createTask and Builtin.createDiscardingTask")
 SUPPRESSIBLE_LANGUAGE_FEATURE(AssociatedTypeImplements, 0, "@_implements on associated types")
+LANGUAGE_FEATURE(MoveOnlyPartialConsumption, 429, "Partial consumption of noncopyable values")
 
 // Swift 6
 UPCOMING_FEATURE(ConciseMagicFile, 274, 6)
@@ -190,7 +191,6 @@ UPCOMING_FEATURE(InferSendableFromCaptures, 418, 6)
 UPCOMING_FEATURE(ImplicitOpenExistentials, 352, 6)
 UPCOMING_FEATURE(RegionBasedIsolation, 414, 6)
 UPCOMING_FEATURE(DynamicActorIsolation, 423, 6)
-UPCOMING_FEATURE(MoveOnlyPartialConsumption, 429, 6)
 
 // Swift 7
 UPCOMING_FEATURE(ExistentialAny, 335, 7)

include/swift/SIL/FieldSensitivePrunedLiveness.h

@@ -239,8 +239,11 @@ struct SubElementOffset {
 /// s a struct => count(s) := sum(s.fields, { f in count(type(f)) })
 ///                             + s.hasDeinit
 /// e an enum  => count(e) := sum(e.elements, { elt in count(type(elt)) })
-///                             + e.hasDeinit
 ///                             + 1 // discriminator
+///                             + e.hasDeinit
+///
+/// The deinit bit is at the end to make drop_deinit produce a value whose
+/// leaves are contiguous.
 struct TypeSubElementCount {
   unsigned number;
 

lib/SIL/Utils/FieldSensitivePrunedLiveness.cpp

@@ -375,13 +375,50 @@ SubElementOffset::computeForValue(SILValue projectionDerivedFromRoot,
 //                        MARK: TypeTreeLeafTypeRange
 //===----------------------------------------------------------------------===//
 
+/// Whether \p targetInst is dominated by one of the provided switch_enum_addr's
+/// destination blocks whose corresponding enum element has no associated values
+/// which need to be destroyed (i.e. either it has no associated values or they
+/// are trivial).
+static bool isDominatedByPayloadlessSwitchEnumAddrDests(
+    SILInstruction *targetInst, ArrayRef<SwitchEnumAddrInst *> seais,
+    DominanceInfo *domTree) {
+  if (seais.empty())
+    return false;
+  auto *target = targetInst->getParent();
+  for (auto *seai : seais) {
+    if (!domTree->dominates(seai, targetInst)) {
+      continue;
+    }
+    auto size = seai->getNumCases();
+    auto ty = seai->getOperand()->getType();
+    for (unsigned index = 0; index < size; ++index) {
+      auto pair = seai->getCase(index);
+      auto *eltDecl = pair.first;
+      if (eltDecl->hasAssociatedValues()) {
+        auto eltTy = ty.getEnumElementType(eltDecl, seai->getFunction());
+        if (!eltTy.isTrivial(*seai->getFunction())) {
+          continue;
+        }
+      }
+      auto *block = pair.second;
+      if (domTree->dominates(block, target))
+        return true;
+    }
+  }
+  return false;
+}
+
 void TypeTreeLeafTypeRange::constructFilteredProjections(
     SILValue value, SILInstruction *insertPt, SmallBitVector &filterBitVector,
     DominanceInfo *domTree,
     llvm::function_ref<bool(SILValue, TypeTreeLeafTypeRange, NeedsDestroy_t)>
         callback) {
   auto *fn = insertPt->getFunction();
   SILType type = value->getType();
+  auto loc =
+      (insertPt->getLoc().getKind() != SILLocation::ArtificialUnreachableKind)
+          ? insertPt->getLoc()
+          : RegularLocation::getAutoGeneratedLocation();
 
   PRUNED_LIVENESS_LOG(llvm::dbgs() << "ConstructFilteredProjection. Bv: "
                           << filterBitVector << '\n');
@@ -409,8 +446,7 @@ void TypeTreeLeafTypeRange::constructFilteredProjections(
         continue;
       }
 
-      auto newValue =
-          builder.createStructElementAddr(insertPt->getLoc(), value, varDecl);
+      auto newValue = builder.createStructElementAddr(loc, value, varDecl);
       callback(newValue, TypeTreeLeafTypeRange(start, next), NeedsDestroy);
       start = next;
     }
@@ -429,60 +465,72 @@ void TypeTreeLeafTypeRange::constructFilteredProjections(
       unsigned next;
     };
     SmallVector<ElementRecord, 2> projectedElements;
-    unsigned start = startEltOffset;
+    unsigned runningStart = startEltOffset;
     for (auto *eltDecl : enumDecl->getAllElements()) {
       if (!eltDecl->hasAssociatedValues())
         continue;
 
-      auto nextType = type.getEnumElementType(eltDecl, fn);
-      unsigned next = start + TypeSubElementCount(nextType, fn);
-      if (noneSet(filterBitVector, start, next)) {
-        start = next;
+      auto eltTy = type.getEnumElementType(eltDecl, fn);
+      unsigned next = runningStart + TypeSubElementCount(eltTy, fn);
+      if (noneSet(filterBitVector, runningStart, next)) {
+        runningStart = next;
         continue;
       }
 
-      projectedElements.push_back({eltDecl, start, next});
-      start = next;
+      projectedElements.push_back({eltDecl, runningStart, next});
+      runningStart = next;
     }
+    assert((runningStart + 1 + (type.isValueTypeWithDeinit() ? 1 : 0)) ==
+           endEltOffset);
 
-    // Add a bit for the discriminator.
-    unsigned next = start + 1;
-
-    if (!allSet(filterBitVector, start, next)) {
+    if (!allSet(filterBitVector, startEltOffset, endEltOffset)) {
+      TinyPtrVector<SwitchEnumAddrInst *> seais;
       for (auto record : projectedElements) {
         // Find a preexisting unchecked_take_enum_data_addr that dominates
         // insertPt.
         bool foundProjection = false;
-        for (auto *user : value->getUsers()) {
-          auto *utedai = dyn_cast<UncheckedTakeEnumDataAddrInst>(user);
-          if (!utedai) {
-            continue;
-          }
-          if (utedai->getElement() == record.element) {
-            continue;
+        StackList<SILValue> worklist(value->getFunction());
+        worklist.push_back(value);
+        while (!worklist.empty()) {
+          auto v = worklist.pop_back_val();
+          for (auto *user : v->getUsers()) {
+            if (auto *ddi = dyn_cast<DropDeinitInst>(user)) {
+              worklist.push_back(ddi);
+              continue;
+            }
+            if (auto *seai = dyn_cast<SwitchEnumAddrInst>(user)) {
+              seais.push_back(seai);
+            }
+            auto *utedai = dyn_cast<UncheckedTakeEnumDataAddrInst>(user);
+            if (!utedai) {
+              continue;
+            }
+            if (utedai->getElement() != record.element) {
+              continue;
+            }
+            if (!domTree->dominates(utedai, insertPt)) {
+              continue;
+            }
+
+            callback(utedai, TypeTreeLeafTypeRange(record.start, record.next),
+                     NeedsDestroy);
+            foundProjection = true;
           }
-          if (!domTree->dominates(utedai, insertPt)) {
-            continue;
-          }
-
-          callback(utedai, TypeTreeLeafTypeRange(record.start, record.next),
-                   DoesNotNeedDestroy);
-          foundProjection = true;
         }
+        (void)foundProjection;
         assert(foundProjection ||
-               llvm::count_if(enumDecl->getAllElements(), [](auto *elt) {
-                 return elt->hasAssociatedValues();
-               }) == 1);
+               llvm::count_if(
+                   enumDecl->getAllElements(),
+                   [](auto *elt) { return elt->hasAssociatedValues(); }) == 1 ||
+               isDominatedByPayloadlessSwitchEnumAddrDests(insertPt, seais,
+                                                           domTree));
       }
       return;
     }
 
     // Then just pass back our enum base value as the pointer.
-    callback(value, TypeTreeLeafTypeRange(start, next), NeedsDestroy);
-
-    // Then set start to next and assert we covered the entire end elt offset.
-    start = next;
-    assert(start == endEltOffset);
+    callback(value, TypeTreeLeafTypeRange(startEltOffset, endEltOffset),
+             NeedsDestroy);
     return;
   }
 
@@ -497,8 +545,7 @@ void TypeTreeLeafTypeRange::constructFilteredProjections(
         continue;
       }
 
-      auto newValue =
-          builder.createTupleElementAddr(insertPt->getLoc(), value, index);
+      auto newValue = builder.createTupleElementAddr(loc, value, index);
       callback(newValue, TypeTreeLeafTypeRange(start, next), NeedsDestroy);
       start = next;
     }
@@ -526,9 +573,12 @@ void TypeTreeLeafTypeRange::get(
 
   // An `inject_enum_addr` only initializes the enum tag.
   if (auto inject = dyn_cast<InjectEnumAddrInst>(op->getUser())) {
-    auto upperBound = *startEltOffset + TypeSubElementCount(projectedValue);
-    // TODO: account for deinit component if enum has deinit.
-    assert(!projectedValue->getType().isValueTypeWithDeinit());
+    // Subtract the deinit bit, if any: the discriminator bit is before it:
+    //
+    // [ case1 bits ..., case2 bits, ..., discriminator bit, deinit bit ]
+    auto deinitBits = projectedValue->getType().isValueTypeWithDeinit() ? 1 : 0;
+    auto upperBound =
+        *startEltOffset + TypeSubElementCount(projectedValue) - deinitBits;
     ranges.push_back({upperBound - 1, upperBound});
     return;
   }
@@ -551,10 +601,11 @@ void TypeTreeLeafTypeRange::get(
       }
       numAtoms += elementAtoms;
     }
-    // TODO: account for deinit component if enum has deinit.
-    assert(!projectedValue->getType().isValueTypeWithDeinit());
+    // The discriminator bit is consumed.
     ranges.push_back(
         {*startEltOffset + numAtoms, *startEltOffset + numAtoms + 1});
+    // The deinit bit is _not_ consumed.  A drop_deinit is required to
+    // consumingly switch an enum with a deinit.
     return;
   }
 

lib/SILGen/SILGenDestructor.cpp

@@ -264,14 +264,11 @@ 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);
-        }
+  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);
       }
     }
   }

lib/SILGen/SILGenProlog.cpp

@@ -59,30 +59,19 @@ SILValue SILGenFunction::emitSelfDeclForDestructor(VarDecl *selfDecl) {
   // 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;
+    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 {
-      if (selfValue->getOwnershipKind() == OwnershipKind::Owned) {
-        selfValue = B.createMarkUnresolvedNonCopyableValueInst(
-            selfDecl, selfValue,
-            MarkUnresolvedNonCopyableValueInst::CheckKind::
-                ConsumableAndAssignable);
-      }
+      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;
   }
 
   VarLocs[selfDecl] = VarLoc::get(selfValue);

lib/SILOptimizer/Analysis/SimplifyInstruction.cpp

@@ -172,6 +172,8 @@ SILValue InstSimplifier::visitStructExtractInst(StructExtractInst *sei) {
 
 SILValue
 InstSimplifier::visitUncheckedEnumDataInst(UncheckedEnumDataInst *uedi) {
+  if (uedi->getOperand()->getType().isValueTypeWithDeinit())
+    return SILValue();
   // (unchecked_enum_data (enum payload)) -> payload
   auto opt = lookThroughOwnershipInsts(uedi->getOperand());
   if (auto *ei = dyn_cast<EnumInst>(opt)) {

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerTester.cpp

@@ -109,7 +109,8 @@ class MoveOnlyAddressCheckerTesterPass : public SILFunctionTransform {
       borrowtodestructure::IntervalMapAllocator allocator;
       MoveOnlyAddressChecker checker{getFunction(), diagnosticEmitter,
                                      allocator, domTree, poa};
-      madeChange = checker.check(moveIntroducersToProcess);
+      madeChange = checker.completeLifetimes();
+      madeChange |= checker.check(moveIntroducersToProcess);
     }
 
     // If we did not emit any diagnostics, emit a diagnostic if we missed any

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.cpp

@@ -242,6 +242,7 @@
 #include "swift/SIL/FieldSensitivePrunedLiveness.h"
 #include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/MemAccessUtils.h"
+#include "swift/SIL/OSSALifetimeCompletion.h"
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/PrunedLiveness.h"
 #include "swift/SIL/SILArgument.h"
@@ -1663,10 +1664,20 @@ struct CopiedLoadBorrowEliminationVisitor
       }
 
       case OperandOwnership::ForwardingConsume:
-      case OperandOwnership::DestroyingConsume:
+      case OperandOwnership::DestroyingConsume: {
+        if (auto *dvi = dyn_cast<DestroyValueInst>(nextUse->getUser())) {
+          auto value = dvi->getOperand();
+          auto *pai = dyn_cast_or_null<PartialApplyInst>(
+              value->getDefiningInstruction());
+          if (pai && pai->isOnStack()) {
+            // A destroy_value of an on_stack partial apply isn't actually a
+            // consuming use--it closes a borrow scope.
+            continue;
+          }
+        }
         // We can only hit this if our load_borrow was copied.
         llvm_unreachable("We should never hit this");
-
+      }
       case OperandOwnership::GuaranteedForwarding: {
         SmallVector<SILValue, 8> forwardedValues;
         auto *fn = nextUse->getUser()->getFunction();
@@ -1791,10 +1802,11 @@ shouldEmitPartialMutationError(UseState &useState, PartialMutation::Kind kind,
   }
 
   auto feature = partialMutationFeature(kind);
-  if (!fn->getModule().getASTContext().LangOpts.hasFeature(feature) &&
+  if (feature &&
+      !fn->getModule().getASTContext().LangOpts.hasFeature(*feature) &&
       !isa<DropDeinitInst>(user)) {
     LLVM_DEBUG(llvm::dbgs()
-               << "    " << getFeatureName(feature) << " disabled!\n");
+               << "    " << getFeatureName(*feature) << " disabled!\n");
     // If the types equal, just bail early.
     // Emit the error.
     return {PartialMutationError::featureDisabled(iterType, kind)};
@@ -3975,3 +3987,29 @@ bool MoveOnlyAddressChecker::check(
   }
   return pimpl.changed;
 }
+bool MoveOnlyAddressChecker::completeLifetimes() {
+  // TODO: Delete once OSSALifetimeCompletion is run as part of SILGenCleanup
+  bool changed = false;
+
+  // Lifetimes must be completed inside out (bottom-up in the CFG).
+  PostOrderFunctionInfo *postOrder = poa->get(fn);
+  OSSALifetimeCompletion completion(fn, domTree);
+  for (auto *block : postOrder->getPostOrder()) {
+    for (SILInstruction &inst : reverse(*block)) {
+      for (auto result : inst.getResults()) {
+        if (completion.completeOSSALifetime(result) ==
+            LifetimeCompletion::WasCompleted) {
+          changed = true;
+        }
+      }
+    }
+    for (SILArgument *arg : block->getArguments()) {
+      assert(!arg->isReborrow() && "reborrows not legal at this SIL stage");
+      if (completion.completeOSSALifetime(arg) ==
+          LifetimeCompletion::WasCompleted) {
+        changed = true;
+      }
+    }
+  }
+  return changed;
+}

lib/SILOptimizer/Mandatory/MoveOnlyAddressCheckerUtils.h

@@ -44,6 +44,7 @@ struct MoveOnlyAddressChecker {
   /// \p diagnosticEmitter.getDiagnosticCount().
   bool check(llvm::SmallSetVector<MarkUnresolvedNonCopyableValueInst *, 32>
                  &moveIntroducersToProcess);
+  bool completeLifetimes();
 };
 
 } // namespace siloptimizer