SILMem2Reg: fix a problem with leaking enum values

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -1757,6 +1757,14 @@ bool MemoryToRegisters::promoteSingleAllocation(AllocStackInst *alloc) {
       b.createDeallocStack(next->getLoc(), alloc);
     }
     return true;
+  } else {
+    // For enums we require that all uses are in the same block.
+    // Otherwise there 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 (f.hasOwnership() && alloc->getType().isOrHasEnum())
+      return false;
   }
 
   LLVM_DEBUG(llvm::dbgs() << "*** Need to insert BB arguments for " << *alloc);

test/SILOptimizer/mem2reg_lifetime.sil

@@ -1328,9 +1328,9 @@ entry(%either : @owned $E):
     %addr = alloc_stack [lexical] $E
     store %either to [init] %addr : $*E
     %either2 = load [take] %addr : $*E
+    dealloc_stack %addr : $*E
     switch_enum %either2 : $E, case #E.one!enumelt: one
 one(%instance : @owned $C):
-    dealloc_stack %addr : $*E
     return %instance : $C
 }
 

test/SILOptimizer/mem2reg_lifetime_nontrivial.sil

@@ -1101,7 +1101,8 @@ enum KlassOptional {
 sil @return_optional_or_error : $@convention(thin) () -> (@owned KlassOptional, @error Error)
 
 // CHECK-LABEL: sil [ossa] @test_deadphi4 :
-// CHECK-NOT: alloc_stack
+// mem2reg cannot optimize an enum location which spans over multiple blocks.
+// CHECK:         alloc_stack
 // CHECK-LABEL: } // end sil function 'test_deadphi4'
 sil [ossa] @test_deadphi4 : $@convention(thin) (@owned KlassOptional) -> () {
 bb0(%0 : @owned $KlassOptional):

test/SILOptimizer/mem2reg_ossa.sil

@@ -488,3 +488,53 @@ bb0:
   dealloc_stack %1 : $*((), ())
   return %16 : $((), ())
 }
+
+// CHECK-LABEL: sil [ossa] @dont_optimize_optional_in_multiple_blocks :
+// CHECK:         alloc_stack
+// CHECK:       } // end sil function 'dont_optimize_optional_in_multiple_blocks'
+sil [ossa] @dont_optimize_optional_in_multiple_blocks : $@convention(method) (@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] @optimize_optional_in_single_block :
+// CHECK-NOT:     alloc_stack
+// CHECK:       } // end sil function 'optimize_optional_in_single_block'
+sil [ossa] @optimize_optional_in_single_block : $@convention(method) (@guaranteed Optional<Klass>) -> () {
+bb0(%0 : @guaranteed $Optional<Klass>):
+  switch_enum %0 : $Optional<Klass>, case #Optional.some!enumelt: bb6, case #Optional.none!enumelt: bb5
+
+bb5:
+  br bb7
+
+bb6(%50 : @guaranteed $Klass):
+  %32 = alloc_stack $Optional<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 : $()
+}
+

test/SILOptimizer/mem2reg_ossa_nontrivial.sil

@@ -1021,7 +1021,8 @@ enum KlassOptional {
 sil @return_optional_or_error : $@convention(thin) () -> (@owned KlassOptional, @error Error)
 
 // CHECK-LABEL: sil [ossa] @test_deadphi4 :
-// CHECK-NOT: alloc_stack
+// mem2reg cannot optimize an enum location which spans over multiple blocks.
+// CHECK:         alloc_stack
 // CHECK-LABEL: } // end sil function 'test_deadphi4'
 sil [ossa] @test_deadphi4 : $@convention(thin) (@owned KlassOptional) -> () {
 bb0(%0 : @owned $KlassOptional):