[Mem2Reg] Don't skip write-only enum locations.

The lifetimes of the values stored into these locations may need to be
completed: if a trivial case of a non-trivial enum is written to the
location, its lifetime must be completed.

rdar://158139991

Author: nate-chandler

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -2214,7 +2214,8 @@ bool MemoryToRegisters::promoteAllocation(AllocStackInst *alloc,
   }
 
   // Remove write-only AllocStacks.
-  if (isWriteOnlyAllocation(alloc) && !lexicalLifetimeEnsured(alloc)) {
+  if (isWriteOnlyAllocation(alloc) && !alloc->getType().isOrHasEnum() &&
+      !lexicalLifetimeEnsured(alloc)) {
     LLVM_DEBUG(llvm::dbgs() << "*** Deleting store-only AllocStack: "<< *alloc);
     deleter.forceDeleteWithUsers(alloc);
     return true;

test/SILOptimizer/mem2reg_ossa_nontrivial.sil

@@ -1591,3 +1591,84 @@ bb3:
   return %r
 }
 
+// CHECK-LABEL: sil [ossa] @end_lifetime_at_trivial_case_dealloc_stack : {{.*}} {
+// CHECK-NOT:     alloc_stack
+// CHECK:         [[IN:%[^,]+]] = load [copy]
+// CHECK:         [[IN_BORROW:%[^,]+]] = begin_borrow [[IN]]
+// CHECK:         switch_enum [[IN_BORROW]]
+// CHECK-SAME:        case #Result.success!enumelt: [[INT:bb[0-9]+]]
+// CHECK:       [[INT]]({{%[^,]+}} : $Int):
+// CHECK:         end_borrow [[IN_BORROW]]
+// CHECK:         end_lifetime [[IN]]
+// CHECK-LABEL: } // end sil function 'end_lifetime_at_trivial_case_dealloc_stack'
+sil [ossa] @end_lifetime_at_trivial_case_dealloc_stack : $@convention(thin) (@in_guaranteed Result<Int, X>) -> () {
+entry(%in_addr : $*Result<Int, X>):
+  %main_addr = alloc_stack $Result<Int, X>
+  %in = load [copy] %in_addr
+  store %in to [init] %main_addr
+  %in_borrow = load_borrow %main_addr
+  switch_enum %in_borrow,
+      case #Result.success!enumelt: int,
+      case #Result.failure!enumelt: full
+
+int(%6 : $Int):
+  end_borrow %in_borrow
+  %temp = alloc_stack $Result<Int, X>
+  %in_copy = load [copy] %main_addr
+  store %in_copy to [init] %temp
+  br int2
+
+int2:
+  dealloc_stack %temp
+  br exit
+
+full(%13 : @guaranteed $X):
+  end_borrow %in_borrow
+  br exit
+
+exit:
+  destroy_addr %main_addr
+  dealloc_stack %main_addr
+  %18 = tuple ()
+  return %18
+}
+
+// CHECK-LABEL: sil [ossa] @end_lifetime_at_trivial_case_dealloc_stack_2 : {{.*}} {
+// CHECK:       bb0([[IN_ADDR:%[^,]+]] :
+// CHECK:         [[IN_COPY:%[^,]+]] = load [copy] [[IN_ADDR]]
+// CHECK:         [[IN_BORROW:%[^,]+]] = begin_borrow [[IN_COPY]]
+// CHECK:         switch_enum [[IN_BORROW]]
+// CHECK:             case #Result.success!enumelt: [[INT:bb[0-9]+]]
+// CHECK:       [[INT]]
+// CHECK:         end_borrow [[IN_BORROW]]
+// CHECK:         end_lifetime [[IN_COPY]]
+// CHECK-LABEL: } // end sil function 'end_lifetime_at_trivial_case_dealloc_stack_2'
+sil [ossa] @end_lifetime_at_trivial_case_dealloc_stack_2 : $@convention(thin) (@in_guaranteed Result<Int, X>) -> () {
+entry(%in_addr : $*Result<Int, X>):
+  %main_addr = alloc_stack $Result<Int, X>
+  %temp = alloc_stack $Result<Int, X>
+  %in = load [copy] %in_addr
+  store %in to [init] %main_addr
+  %in_borrow = load_borrow %main_addr
+  switch_enum %in_borrow,
+      case #Result.success!enumelt: int,
+      case #Result.failure!enumelt: x
+
+int(%int : $Int):
+  end_borrow %in_borrow
+  %in_copy = load [copy] %main_addr
+  store %in_copy to [init] %temp
+  br int2
+
+int2:
+  br exit
+
+x(%13 : @guaranteed $X):
+  end_borrow %in_borrow
+  br exit
+
+exit:
+  destroy_addr %main_addr
+  %retval = tuple ()
+  return %retval
+}