Merge pull request #71025 from jckarter/borrowing-switch-1

SILGen: Always match noncopyable values by borrowing.

Author: jckarter

include/swift/Basic/Features.def

@@ -273,6 +273,9 @@ EXPERIMENTAL_FEATURE(TransferringArgsAndResults, true)
 // Enable `@preconcurrency` attribute on protocol conformances.
 EXPERIMENTAL_FEATURE(PreconcurrencyConformances, false)
 
+// Allow for `switch` of noncopyable values to be borrowing or consuming.
+EXPERIMENTAL_FEATURE(BorrowingSwitch, true)
+
 #undef EXPERIMENTAL_FEATURE_EXCLUDED_FROM_MODULE_INTERFACE
 #undef EXPERIMENTAL_FEATURE
 #undef UPCOMING_FEATURE

include/swift/SIL/SILInstruction.h

@@ -6736,8 +6736,14 @@ class InjectEnumAddrInst
   }
 };
 
-/// Invalidate an enum value and take ownership of its payload data
-/// without moving it in memory.
+/// Project an enum's payload data without checking the case of the enum or
+/// moving it in memory.
+///
+/// For some classes of enum, this is a destructive operation that invalidates
+/// the enum, particularly in cases where the layout algorithm can potentially
+/// use the common spare bits out of the payloads of a multi-payload enum
+/// to store the tag without allocating additional space. The `isDestructive`
+/// static method returns true for enums where this is potentially the case.
 class UncheckedTakeEnumDataAddrInst
   : public UnaryInstructionBase<SILInstructionKind::UncheckedTakeEnumDataAddrInst,
                                 SingleValueInstruction>
@@ -6755,6 +6761,15 @@ class UncheckedTakeEnumDataAddrInst
   }
 
 public:
+  // Returns true if the projection operation is possibly destructive for
+  // instances of the given enum declaration.
+  static bool isDestructive(EnumDecl *forEnum, SILModule &M);
+
+  // Returns true if this projection operation is possibly destructive.
+  bool isDestructive() const {
+    return isDestructive(Element->getParentEnum(), getModule());
+  }
+
   EnumElementDecl *getElement() const { return Element; }
 
   unsigned getCaseIndex() {
@@ -7025,6 +7040,7 @@ class TupleExtractInst
                    ValueOwnershipKind forwardingOwnershipKind)
       : UnaryInstructionBase(DebugLoc, Operand, ResultTy,
                              forwardingOwnershipKind) {
+    assert(Operand->getType().castTo<TupleType>());
     sharedUInt32().TupleExtractInst.fieldNo = FieldNo;
   }
 

lib/AST/ASTPrinter.cpp

@@ -3938,6 +3938,8 @@ static bool usesFeaturePreconcurrencyConformances(Decl *decl) {
   return false;
 }
 
+static bool usesFeatureBorrowingSwitch(Decl *decl) { return false; }
+
 /// Suppress the printing of a particular feature.
 static void suppressingFeature(PrintOptions &options, Feature feature,
                                llvm::function_ref<void()> action) {

lib/SIL/IR/SILInstruction.cpp

@@ -1070,6 +1070,10 @@ MemoryBehavior SILInstruction::getMemoryBehavior() const {
     }
     llvm_unreachable("Covered switch isn't covered?!");
   }
+  
+  // TODO: An UncheckedTakeEnumDataAddr instruction has no memory behavior if
+  // it is nondestructive. Setting this currently causes LICM to miscompile
+  // because access paths do not account for enum projections.
 
   switch (getKind()) {
 #define FULL_INST(CLASS, TEXTUALNAME, PARENT, MEMBEHAVIOR, RELEASINGBEHAVIOR)  \
@@ -1878,6 +1882,35 @@ DestroyValueInst::getNonescapingClosureAllocation() const {
   }
 }
 
+bool
+UncheckedTakeEnumDataAddrInst::isDestructive(EnumDecl *forEnum, SILModule &M) {
+  // We only potentially use spare bit optimization when an enum is always
+  // loadable.
+  auto sig = forEnum->getGenericSignature().getCanonicalSignature();
+  if (SILType::isAddressOnly(forEnum->getDeclaredInterfaceType()->getReducedType(sig),
+                             M.Types, sig,
+                             TypeExpansionContext::minimal())) {
+    return false;
+  }
+  
+  // We only overlap spare bits with valid payload values when an enum has
+  // multiple payloads.
+  bool sawPayloadCase = false;
+  for (auto element : forEnum->getAllElements()) {
+    if (element->hasAssociatedValues()) {
+      if (sawPayloadCase) {
+        // TODO: If the associated value's type is always visibly empty then it
+        // would get laid out like a no-payload case.
+        return true;
+      } else {
+        sawPayloadCase = true;
+      }
+    }
+  }
+  
+  return false;
+}
+
 #ifndef NDEBUG
 
 //---

lib/SIL/Verifier/SILVerifier.cpp

@@ -689,9 +689,7 @@ struct ImmutableAddressUseVerifier {
         }
         break;
       case SILInstructionKind::UncheckedTakeEnumDataAddrInst: {
-        auto type =
-            cast<UncheckedTakeEnumDataAddrInst>(inst)->getOperand()->getType();
-        if (type.getOptionalObjectType()) {
+        if (!cast<UncheckedTakeEnumDataAddrInst>(inst)->isDestructive()) {
           for (auto result : inst->getResults()) {
             llvm::copy(result->getUses(), std::back_inserter(worklist));
           }

lib/SILGen/ManagedValue.h

@@ -469,9 +469,6 @@ class ConsumableManagedValue {
   /*implicit*/ ConsumableManagedValue(ManagedValue value,
                                       CastConsumptionKind finalConsumption)
       : Value(value), FinalConsumption(finalConsumption) {
-    assert((value.getType().isObject() ||
-            finalConsumption != CastConsumptionKind::BorrowAlways) &&
-           "Can not borrow always a value");
     assert((value.getType().isAddress() ||
             finalConsumption != CastConsumptionKind::CopyOnSuccess) &&
            "Can not copy on success a value.");

lib/SILGen/SILGenPattern.cpp

@@ -2058,15 +2058,21 @@ void PatternMatchEmission::emitEnumElementDispatch(
   // After this point we now that we must have an address only type.
   assert(src.getType().isAddressOnly(SGF.F) &&
          "Should have an address only type here");
+  assert(!UncheckedTakeEnumDataAddrInst::isDestructive(src.getType().getEnumOrBoundGenericEnum(),
+                                                       SGF.getModule()) &&
+         "address only enum projection should never be destructive");
 
   CanType sourceType = rows[0].Pattern->getType()->getCanonicalType();
 
   // Collect the cases and specialized rows.
   CaseBlocks blocks{SGF, rows, sourceType, SGF.B.getInsertionBB()};
 
-  // We lack a SIL instruction to nondestructively project data from an
+  // We (used to) lack a SIL instruction to nondestructively project data from an
   // address-only enum, so we can only do so in place if we're allowed to take
   // the source always. Copy the source if we can't.
+  //
+  // TODO: This should no longer be necessary now that we guarantee that
+  // potentially address-only enums never use spare bit optimization.
   switch (src.getFinalConsumption()) {
   case CastConsumptionKind::TakeAlways:
   case CastConsumptionKind::CopyOnSuccess:
@@ -2086,8 +2092,6 @@ void PatternMatchEmission::emitEnumElementDispatch(
   }
 
   // Emit the switch_enum_addr instruction.
-  //
-  // NOTE: switch_enum_addr does not actually consume the underlying value.
   SGF.B.createSwitchEnumAddr(loc, src.getValue(), blocks.getDefaultBlock(),
                              blocks.getCaseBlocks(), blocks.getCounts(),
                              defaultCaseCount);
@@ -2171,22 +2175,21 @@ void PatternMatchEmission::emitEnumElementDispatch(
       ManagedValue eltValue;
       // We can only project destructively from an address-only enum, so
       // copy the value if we can't consume it.
-      // TODO: Should have a more efficient way to copy payload
-      // nondestructively from an enum.
+      // TODO: Copying should be avoidable now that we guarantee that address-
+      // only enums never use spare bit optimization.
       switch (eltConsumption) {
       case CastConsumptionKind::TakeAlways: {
         auto finalValue = src.getFinalManagedValue();
         eltValue = SGF.B.createUncheckedTakeEnumDataAddr(loc, finalValue,
                                                          eltDecl, eltTy);
         break;
       }
-      case CastConsumptionKind::BorrowAlways:
-        // If we reach this point, we know that we have a loadable
-        // element type from an enum with mixed address
-        // only/loadable cases. Since we had an address only type,
-        // we assume that we will not have BorrowAlways since
-        // address only types do not support BorrowAlways.
-        llvm_unreachable("not allowed");
+      case CastConsumptionKind::BorrowAlways: {
+        eltValue = ManagedValue::forBorrowedAddressRValue(
+          SGF.B.createUncheckedTakeEnumDataAddr(loc, src.getValue(),
+                                                eltDecl, eltTy));
+        break;
+      }
       case CastConsumptionKind::CopyOnSuccess: {
         auto temp = SGF.emitTemporary(loc, SGF.getTypeLowering(src.getType()));
         SGF.B.createCopyAddr(loc, src.getValue(), temp->getAddress(), IsNotTake,
@@ -2212,12 +2215,22 @@ void PatternMatchEmission::emitEnumElementDispatch(
       if (eltTL->isLoadable()) {
         // If we do not have a loadable value, just use getManagedSubobject
         // Load a loadable data value.
-        if (eltConsumption == CastConsumptionKind::CopyOnSuccess) {
+        switch (eltConsumption) {
+        case CastConsumptionKind::CopyOnSuccess:
           eltValue = SGF.B.createLoadBorrow(loc, eltValue);
           eltConsumption = CastConsumptionKind::BorrowAlways;
-        } else {
-          assert(eltConsumption == CastConsumptionKind::TakeAlways);
+          break;
+
+        case CastConsumptionKind::TakeAlways:
           eltValue = SGF.B.createLoadTake(loc, eltValue);
+          break;
+          
+        case CastConsumptionKind::BorrowAlways:
+          eltValue = SGF.B.createLoadBorrow(loc, eltValue);
+          break;
+          
+        case CastConsumptionKind::TakeOnSuccess:
+          llvm_unreachable("not possible");
         }
         origCMV = {eltValue, eltConsumption};
       } else {
@@ -2847,33 +2860,88 @@ void SILGenFunction::emitSwitchStmt(SwitchStmt *S) {
   ManagedValue subjectMV = emitRValueAsSingleValue(
       S->getSubjectExpr(), SGFContext::AllowGuaranteedPlusZero);
 
+  llvm::Optional<ValueOwnership> noncopyableSwitchOwnership;
+
   // Inline constructor for subject.
   auto subject = ([&]() -> ConsumableManagedValue {
-    // If we have a noImplicitCopy value, ensure plus one and convert
-    // it. Switches always consume move only values.
-    //
-    // NOTE: We purposely do not do this for pure move only types since for them
-    // we emit everything at +0 and then run the BorrowToDestructure transform
-    // upon them. The reason that we do this is that internally to
-    // SILGenPattern, we always attempt to move from +1 -> +0 meaning that even
-    // if we start at +1, we will go back to +0 given enough patterns to go
-    // through. It is simpler to just let SILGenPattern do what it already wants
-    // to do, rather than fight it or try to resusitate the "fake owned borrow"
-    // path that we still use for address only types (and that we want to delete
-    // once we have opaque values).
-    if (subjectMV.getType().isMoveOnly() && subjectMV.getType().isObject()) {
-      if (subjectMV.getType().isMoveOnlyWrapped()) {
-        subjectMV = B.createOwnedMoveOnlyWrapperToCopyableValue(
-            S, subjectMV.ensurePlusOne(*this, S));
-      } else {
-        // If we have a pure move only type and it is owned, borrow it so that
-        // BorrowToDestructure can handle it.
-        if (subjectMV.getOwnershipKind() == OwnershipKind::Owned) {
+    if (subjectMV.getType().isMoveOnly()) {
+      if (getASTContext().LangOpts.hasFeature(Feature::BorrowingSwitch)) {
+        // Determine the overall ownership behavior of the switch, based on the
+        // patterns' ownership behavior.
+        auto ownership = ValueOwnership::Shared;
+        for (auto caseLabel : S->getCases()) {
+          for (auto item : caseLabel->getCaseLabelItems()) {
+            ownership = std::max(ownership, item.getPattern()->getOwnership());
+          }
+        }
+        noncopyableSwitchOwnership = ownership;
+        
+        // Based on the ownership behavior, prepare the subject.
+        // The pattern match itself will always be performed on a borrow, to
+        // ensure that the order of pattern evaluation doesn't prematurely
+        // consume or modify the value until we commit to a match. But if the
+        // match consumes the value, then we need a +1 value to go back to in
+        // order to consume the parts we match to, so we force a +1 value then
+        // borrow that for the pattern match.
+        switch (ownership) {
+        case ValueOwnership::Default:
+          llvm_unreachable("invalid");
+        
+        case ValueOwnership::Shared:
+          if (subjectMV.getType().isAddress() &&
+              subjectMV.getType().isLoadable(F)) {
+            // Load a borrow if the type is loadable.
+            subjectMV = B.createLoadBorrow(S, subjectMV);
+          } else if (!subjectMV.isPlusZero()) {
+            subjectMV = subjectMV.borrow(*this, S);
+          }
+          return {subjectMV, CastConsumptionKind::BorrowAlways};
+          
+        case ValueOwnership::InOut:
+          // TODO: mutating switches
+          llvm_unreachable("not implemented");
+          
+        case ValueOwnership::Owned:
+          // Make sure we own the subject value.
+          subjectMV = subjectMV.ensurePlusOne(*this, S);
+          if (subjectMV.getType().isAddress() &&
+              subjectMV.getType().isLoadable(F)) {
+            // Move the value into memory if it's loadable.
+            subjectMV = B.createLoadTake(S, subjectMV);
+          }
+          // Perform the pattern match on a borrow of the subject.
           subjectMV = subjectMV.borrow(*this, S);
+          return {subjectMV, CastConsumptionKind::BorrowAlways};
+        }
+        llvm_unreachable("unhandled value ownership");
+      } else {
+        // If we have a noImplicitCopy value, ensure plus one and convert
+        // it. Switches always consume move only values.
+        //
+        // NOTE: We purposely do not do this for pure move only types since for them
+        // we emit everything at +0 and then run the BorrowToDestructure transform
+        // upon them. The reason that we do this is that internally to
+        // SILGenPattern, we always attempt to move from +1 -> +0 meaning that even
+        // if we start at +1, we will go back to +0 given enough patterns to go
+        // through. It is simpler to just let SILGenPattern do what it already wants
+        // to do, rather than fight it or try to resusitate the "fake owned borrow"
+        // path that we still use for address only types (and that we want to delete
+        // once we have opaque values).
+        if (subjectMV.getType().isObject()) {
+          if (subjectMV.getType().isMoveOnlyWrapped()) {
+            subjectMV = B.createOwnedMoveOnlyWrapperToCopyableValue(
+                S, subjectMV.ensurePlusOne(*this, S));
+          } else {
+            // If we have a pure move only type and it is owned, borrow it so that
+            // BorrowToDestructure can handle it.
+            if (subjectMV.getOwnershipKind() == OwnershipKind::Owned) {
+              subjectMV = subjectMV.borrow(*this, S);
+            }
+          }
         }
       }
     }
-
+  
     // If we have a plus one value...
     if (subjectMV.isPlusOne(*this)) {
       // And we have an address that is loadable, perform a load [take].

lib/Sema/MiscDiagnostics.cpp

@@ -4112,7 +4112,8 @@ diagnoseMoveOnlyPatternMatchSubject(ASTContext &C,
   if (auto load = dyn_cast<LoadExpr>(subjectExpr)) {
     subjectExpr = load->getSubExpr()->getSemanticsProvidingExpr();
   }
-  if (isa<DeclRefExpr>(subjectExpr)) {
+  if (!C.LangOpts.hasFeature(Feature::BorrowingSwitch)
+      && isa<DeclRefExpr>(subjectExpr)) {
     C.Diags.diagnose(subjectExpr->getLoc(),
                            diag::move_only_pattern_match_not_consumed)
       .fixItInsert(subjectExpr->getStartLoc(), "consume ");