[silgenpattern] When emitting address only enum element dispatch, do not leak the switch's operand along the default path.

For the following discussion, let OP be the switch's operand. This is
implemented by:

* Modeling switch_enum_addr as not forwarding OP and instead delegate the
  forwarding of OP to be done by each enum case. This matches what SILGen is
  actually doing since the actual taking of the address in SIL is done in each
  enum case block by an unchecked_take_enum_data_addr.

* In each enum case, I treat OP as being forwarded into an irrefutable
  sub-tree. I follow the pattern of other places this is done by creating a
  CleanupStateRestorationScope and using forwardIntoIrrefutableSubTree. This
  ensures that if I forward OP in the enum case, it just becomes dormant instead
  of being thrown away.

* Inside each case, there is a bunch of code that does some final preparations
  to src before dispatching to the inner dispatch. This code was written using
  old low-level SILValue APIs where ownership is done by hand. I replaced all of
  that by hand ownership with higher level Managed Value APIs that automatically
  handle ownership for the user. This simplified the implementation and ensured
  correctness via SILGenBuilder API invariants.

The end result of all of these together is that:

1. The cleanup on OP is still live when we emit the default case later than the
   cleanups. This eliminates the leak that I am fixing.

2. We have greater correctness since the SILGenBuilder APIs automatically handle
   ownership for us.

3. We have eliminated some brittle logic that could in the future introduce
   bugs. Specifically, I noticed that if we were ever given a
   ConsumableManagedValue that was CopyOnSuccess or BorrowAlways but its
   ManagedValue was a +1 value, we would leak. I could not figure out how to
   create a Swift test case that would go down this code path though = (. But
   that being said, it is one new language feature away from being broken. I
   added some asserts to ConsumableManagedValue that ensures this invariant, so
   we are safe. If it is inconvenient, we can also cause ConsumableManagedValue
   to create a new ManagedValue when it detects this condition without the
   cleanup. But lets see how difficult it is to keep this invariant.

In terms of testing, I put in both a SILGen test and also an end<->end
interpreter test to ensure this doesn't break again.

rdar://71992652
SR-13926

Author: gottesmm

lib/SILGen/SILGenPattern.cpp

@@ -1398,6 +1398,25 @@ getManagedSubobject(SILGenFunction &SGF, SILValue value,
   llvm_unreachable("covered switch");
 }
 
+/// Given that we've broken down a source value into this subobject,
+/// and that we were supposed to use the given consumption rules on
+/// it, construct an appropriate managed value.
+static ConsumableManagedValue
+getManagedSubobject(SILGenFunction &SGF, ManagedValue value,
+                    CastConsumptionKind consumption) {
+  switch (consumption) {
+  case CastConsumptionKind::BorrowAlways:
+  case CastConsumptionKind::CopyOnSuccess:
+    return {value.unmanagedBorrow(), consumption};
+  case CastConsumptionKind::TakeAlways:
+  case CastConsumptionKind::TakeOnSuccess: {
+    auto loc = RegularLocation::getAutoGeneratedLocation();
+    return {value.ensurePlusOne(SGF, loc), consumption};
+  }
+  }
+  llvm_unreachable("covered switch");
+}
+
 static ConsumableManagedValue
 emitReabstractedSubobject(SILGenFunction &SGF, SILLocation loc,
                           ConsumableManagedValue value,
@@ -2064,8 +2083,9 @@ void PatternMatchEmission::emitEnumElementDispatch(
   }
 
   // Emit the switch_enum_addr instruction.
-  SILValue srcValue = src.getFinalManagedValue().forward(SGF);
-  SGF.B.createSwitchEnumAddr(loc, srcValue, blocks.getDefaultBlock(),
+  //
+  // NOTE: switch_enum_addr does not actually consume the underlying value.
+  SGF.B.createSwitchEnumAddr(loc, src.getValue(), blocks.getDefaultBlock(),
                              blocks.getCaseBlocks(), blocks.getCounts(),
                              defaultCaseCount);
 
@@ -2077,12 +2097,16 @@ void PatternMatchEmission::emitEnumElementDispatch(
 
     SGF.B.setInsertionPoint(caseBB);
 
+    // We need to make sure our cleanup stays around long enough for us to emit
+    // our destroy, so setup a cleanup state restoration scope for each case.
+    CleanupStateRestorationScope srcScope(SGF.Cleanups);
+    forwardIntoIrrefutableSubtree(SGF, srcScope, src);
+
     // We're in conditionally-executed code; enter a scope.
     Scope scope(SGF.Cleanups, CleanupLocation::get(loc));
-      
+
     // Create a BB argument or 'unchecked_take_enum_data_addr'
     // instruction to receive the enum case data if it has any.
-
     SILType eltTy;
     bool hasElt = false;
     if (eltDecl->hasAssociatedValues()) {
@@ -2108,6 +2132,15 @@ void PatternMatchEmission::emitEnumElementDispatch(
         }
       }
 
+      // Forward src along this path so we don't emit a destroy_addr on our
+      // subject value for this case.
+      //
+      // FIXME: Do we actually want to do this? SILGen tests today assume this
+      // pattern. It might be worth leaving the destroy_addr there to create
+      // additional liveness information. For now though, we maintain the
+      // current behavior.
+      src.getFinalManagedValue().forward(SGF);
+
       SILValue result;
       if (hasNonAny) {
         result = SGF.emitEmptyTuple(loc);
@@ -2132,16 +2165,18 @@ void PatternMatchEmission::emitEnumElementDispatch(
         eltConsumption = CastConsumptionKind::TakeAlways;
       }
 
-      SILValue eltValue;
+      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.
       switch (eltConsumption) {
-      case CastConsumptionKind::TakeAlways:
-        eltValue =
-            SGF.B.createUncheckedTakeEnumDataAddr(loc, srcValue, eltDecl, eltTy);
+      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
@@ -2150,11 +2185,15 @@ void PatternMatchEmission::emitEnumElementDispatch(
         // address only types do not support BorrowAlways.
         llvm_unreachable("not allowed");
       case CastConsumptionKind::CopyOnSuccess: {
-        auto copy = SGF.emitTemporaryAllocation(loc, srcValue->getType());
-        SGF.B.createCopyAddr(loc, srcValue, copy, IsNotTake, IsInitialization);
+        auto temp = SGF.emitTemporary(loc, SGF.getTypeLowering(src.getType()));
+        SGF.B.createCopyAddr(loc, src.getValue(), temp->getAddress(), IsNotTake,
+                             IsInitialization);
+        temp->finishInitialization(SGF);
+
         // We can always take from the copy.
         eltConsumption = CastConsumptionKind::TakeAlways;
-        eltValue = SGF.B.createUncheckedTakeEnumDataAddr(loc, copy, eltDecl, eltTy);
+        eltValue = SGF.B.createUncheckedTakeEnumDataAddr(
+            loc, temp->getManagedAddress(), eltDecl, eltTy);
         break;
       }
 
@@ -2170,17 +2209,16 @@ void PatternMatchEmission::emitEnumElementDispatch(
       if (eltTL->isLoadable()) {
         // If we do not have a loadable value, just use getManagedSubObject
         // Load a loadable data value.
-        auto managedEltValue = ManagedValue::forUnmanaged(eltValue);
         if (eltConsumption == CastConsumptionKind::CopyOnSuccess) {
-          managedEltValue = SGF.B.createLoadBorrow(loc, managedEltValue);
+          eltValue = SGF.B.createLoadBorrow(loc, eltValue);
           eltConsumption = CastConsumptionKind::BorrowAlways;
         } else {
           assert(eltConsumption == CastConsumptionKind::TakeAlways);
-          managedEltValue = SGF.B.createLoadTake(loc, managedEltValue);
+          eltValue = SGF.B.createLoadTake(loc, eltValue);
         }
-        origCMV = {managedEltValue, eltConsumption};
+        origCMV = {eltValue, eltConsumption};
       } else {
-        origCMV = getManagedSubobject(SGF, eltValue, *eltTL, eltConsumption);
+        origCMV = getManagedSubobject(SGF, eltValue, eltConsumption);
       }
 
       eltCMV = origCMV;

test/Interpreter/switch_default.swift

@@ -0,0 +1,28 @@
+// RUN: %target-run-simple-swift
+
+// REQUIRES: executable_test
+
+import StdlibUnittest
+
+var SwitchDefaultTestSuite = TestSuite("Switch.Default")
+defer { runAllTests() }
+
+class Klass {}
+protocol Protocol {}
+
+enum Enum {
+  case value1(LifetimeTracked)
+  case value2(Protocol)
+}
+
+SwitchDefaultTestSuite.test("do not leak default case payload") {
+  // We are passing in Enum.value1 so we go down the default and leak our
+  // lifetime tracked.
+  func f(_ e: Enum?) {
+    switch (e) {
+    case .value2: return
+    default: return
+    }
+  }
+  f(.value1(LifetimeTracked(0)))
+}

test/SILGen/switch_default.swift

@@ -0,0 +1,54 @@
+// RUN: %target-swift-emit-silgen -module-name switch_default %s | %FileCheck %s
+
+class Klass {}
+protocol Protocol {}
+
+enum Enum {
+  case value1(Klass)
+  case value2(Protocol)
+}
+
+// CHECK-LABEL: sil hidden [ossa] @$s14switch_default33testAddressOnlySubjectDefaultCaseyyAA4EnumOSgF : $@convention(thin) (@in_guaranteed Optional<Enum>) -> () {
+// CHECK: bb0([[ARG:%.*]] :
+// CHECK:   [[STACK:%.*]] = alloc_stack $Optional<Enum>
+// CHECK:   copy_addr [[ARG]] to [initialization] [[STACK]]
+// CHECK:   switch_enum_addr [[STACK]] : $*Optional<Enum>, case #Optional.some!enumelt: [[SOME_BB:bb[0-9]*]], default [[DEFAULT_BB:bb[0-9]*]]
+//
+// CHECK: [[SOME_BB]]:
+// CHECK:    [[STACK_1:%.*]] = alloc_stack $Optional<Enum>
+// CHECK:    copy_addr [[STACK]] to [initialization] [[STACK_1]]
+// CHECK:    [[TAKEN_ADDR:%.*]] = unchecked_take_enum_data_addr [[STACK_1]]
+// CHECK:    switch_enum_addr [[TAKEN_ADDR]] : $*Enum, case #Enum.value2!enumelt: [[VALUE2_BB:bb[0-9]*]], default [[DEFAULT_BB_2:bb[0-9]*]]
+//
+// CHECK: [[VALUE2_BB]]:
+// CHECK:    [[TAKEN_TAKEN_ADDR:%.*]] = unchecked_take_enum_data_addr [[TAKEN_ADDR]]
+// CHECK:    destroy_addr [[TAKEN_TAKEN_ADDR]]
+// CHECK:    dealloc_stack [[STACK_1]]
+// CHECK:    destroy_addr [[STACK]]
+// CHECK:    dealloc_stack [[STACK]]
+// CHECK:    br [[EXIT_BB:bb[0-9]+]]
+//
+// We used to leak here!
+// CHECK: [[DEFAULT_BB_2]]:
+// CHECK:    destroy_addr [[TAKEN_ADDR]]
+// CHECK:    dealloc_stack [[STACK_1]]
+// CHECK:    br [[CONT_BB:bb[0-9]*]]
+//
+// CHECK: [[DEFAULT_BB]]:
+// CHECK:    br [[CONT_BB]]
+//
+// CHECK: [[CONT_BB]]:
+// CHECK:    destroy_addr [[STACK]]
+// CHECK:    dealloc_stack [[STACK]]
+// CHECK:    br [[EXIT_BB]]
+//
+// CHECK: [[EXIT_BB]]:
+// CHECK-NEXT: tuple
+// CHECK-NEXT: return
+// } // end sil function '$s14switch_default33testAddressOnlySubjectDefaultCaseyyAA4EnumOSgF'
+func testAddressOnlySubjectDefaultCase(_ e: Enum?) {
+  switch (e) {
+  case .value2: return
+  default: return
+  }
+}