TempRValueElimination: fix a problem with leaking enum values

When optimizing an enum `store` to an `alloc_stack`, require that all uses are in the same block.
Otherwise it could be a `switch_enum` of an optional where the none-case does not have a destroy of the enum value.
After transforming such an `alloc_stack`, the value would leak in the none-case block.

Fixes a OSSA verification error.

Author: eeckstein

include/swift/SIL/SILType.h

@@ -231,6 +231,14 @@ class SILType {
   EnumDecl *getEnumOrBoundGenericEnum() const {
     return getASTType().getEnumOrBoundGenericEnum();
   }
+  
+  /// Returns true if this type is an enum or contains an enum.
+  bool isOrHasEnum() const {
+    return getASTType().findIf([](Type ty) {
+      return ty->getEnumOrBoundGenericEnum() != nullptr;
+    });
+  }
+
   /// Retrieve the NominalTypeDecl for a type that maps to a Swift
   /// nominal or bound generic nominal type.
   NominalTypeDecl *getNominalOrBoundGenericNominal() const {

lib/SIL/Verifier/MemoryLifetimeVerifier.cpp

@@ -185,12 +185,6 @@ static bool injectsNoPayloadCase(InjectEnumAddrInst *IEAI) {
   return elemType.isEmpty(*function);
 }
 
-static bool isOrHasEnum(SILType type) {
-  return type.getASTType().findIf([](Type ty) {
-    return ty->getEnumOrBoundGenericEnum() != nullptr;
-  });
-}
-
 bool MemoryLifetimeVerifier::storesTrivialEnum(int locIdx,
                         SILBasicBlock::reverse_iterator start,
                         SILBasicBlock::reverse_iterator end) {
@@ -204,7 +198,7 @@ bool MemoryLifetimeVerifier::storesTrivialEnum(int locIdx,
     if (auto *SI = dyn_cast<StoreInst>(&inst)) {
       const Location *loc = locations.getLocation(SI->getDest());
       if (loc && loc->isSubLocation(locIdx) &&
-          isOrHasEnum(SI->getSrc()->getType())) {
+          SI->getSrc()->getType().isOrHasEnum()) {
         return SI->getOwnershipQualifier() == StoreOwnershipQualifier::Trivial;
       }
     }

lib/SILOptimizer/Transforms/TempRValueElimination.cpp

@@ -653,6 +653,8 @@ TempRValueOptPass::tryOptimizeStoreIntoTemp(StoreInst *si) {
     return std::next(si->getIterator());
   }
 
+  bool isOrHasEnum = tempObj->getType().isOrHasEnum();
+
   // Scan all uses of the temporary storage (tempObj) to verify they all refer
   // to the value initialized by this copy. It is sufficient to check that the
   // only users that modify memory are the copy_addr [initialization] and
@@ -664,6 +666,16 @@ TempRValueOptPass::tryOptimizeStoreIntoTemp(StoreInst *si) {
     if (user == si)
       continue;
 
+    // For enums we require that all uses are in the same block.
+    // Otherwise it could be a switch_enum of an optional where the none-case
+    // does not have a destroy of the enum value.
+    // After transforming such an alloc_stack the value would leak in the none-
+    // case block.
+    if (isOrHasEnum && user->getParent() != si->getParent() &&
+        !isa<DeallocStackInst>(user)) {
+      return std::next(si->getIterator());
+    }
+
     // Bail if there is any kind of user which is not handled in the code below.
     switch (user->getKind()) {
       case SILInstructionKind::DestroyAddrInst:

test/SILOptimizer/temp_rvalue_opt_ossa.sil

@@ -1293,6 +1293,55 @@ bb0(%0 : @owned $GS<Builtin.NativeObject>):
   return %v : $()
 }
 
+// CHECK-LABEL: sil [ossa] @dont_optimize_store_of_enum
+// CHECK:         alloc_stack
+// CHECK:       } // end sil function 'dont_optimize_store_of_enum'
+sil [ossa] @dont_optimize_store_of_enum : $@convention(method) <Element> (@guaranteed Optional<Klass>) -> () {
+bb0(%0 : @guaranteed $Optional<Klass>):
+  %1 = copy_value %0 : $Optional<Klass>
+  %32 = alloc_stack $Optional<Klass>
+  store %1 to [init] %32 : $*Optional<Klass>
+  switch_enum %0 : $Optional<Klass>, case #Optional.some!enumelt: bb6, case #Optional.none!enumelt: bb5
+
+bb5:
+  dealloc_stack %32 : $*Optional<Klass>
+  br bb7
+
+bb6(%50 : @guaranteed $Klass):
+  %53 = load [take] %32 : $*Optional<Klass>
+  destroy_value %53 : $Optional<Klass>
+  dealloc_stack %32 : $*Optional<Klass>
+  br bb7
+
+bb7:
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil [ossa] @store_of_enum_must_be_in_same_block
+// CHECK-NOT:     alloc_stack
+// CHECK:       } // end sil function 'store_of_enum_must_be_in_same_block'
+sil [ossa] @store_of_enum_must_be_in_same_block : $@convention(method) <Element> (@guaranteed Optional<Klass>) -> () {
+bb0(%0 : @guaranteed $Optional<Klass>):
+  %32 = alloc_stack $Optional<Klass>
+  switch_enum %0 : $Optional<Klass>, case #Optional.some!enumelt: bb6, case #Optional.none!enumelt: bb5
+
+bb5:
+  dealloc_stack %32 : $*Optional<Klass>
+  br bb7
+
+bb6(%50 : @guaranteed $Klass):
+  %1 = copy_value %0 : $Optional<Klass>
+  store %1 to [init] %32 : $*Optional<Klass>
+  %53 = load [take] %32 : $*Optional<Klass>
+  destroy_value %53 : $Optional<Klass>
+  dealloc_stack %32 : $*Optional<Klass>
+  br bb7
+
+bb7:
+  %r = tuple ()
+  return %r : $()
+}
 ////////////////////////////////////////
 // Unchecked Take Enum Data Addr Inst //
 ////////////////////////////////////////