SILGen: Always match noncopyable values by borrowing.

Even if the final pattern ends up consuming the value, the match itself
must be nondestructive, because any match condition could fail and cause
us to have to go back to the original aggregate. For copyable values,
we can always copy our way out of consuming operations, but we don't
have that luxury for noncopyable types, so the entire match operation
has to be done as a borrow.

For address-only enums, this requires codifying part of our tag layout
algorithm in SIL, namely that an address-only enum will never use
spare bits or other overlapping storage for the enum tag. This allows
us to assume that `unchecked_take_enum_data_addr` is safely non-side-
effecting and match an address-only noncopyable enum as a borrow.
I put TODOs to remove defensive copies from various parts of our
copyable enum codegen, as well as to have the instruction report
its memory behavior as `None` when the projection is nondestructive,
but this disturbs SILGen for existing code in ways SIL passes aren't
yet ready for, so I'll leave those as is for now.

This patch is enough to get simple examples of noncopyable enum switches
to SILGen correctly. Additional work is necessary to stage in the binding
step of the pattern match; for a consuming switch, we'll need to end
the borrow(s) and then reproject the matched components so we can
consume them moving them into the owned bindings. The move-only checker
also needs to be updated because it currently always tries to convert
a switch into a consuming operation.

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

@@ -3940,6 +3940,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 ");