[Mem2Reg] Complete single-block trivial enum cases

Add `end_lifetime` for values of non-trivial enum type which aren't
destroyed in non-dead-end blocks.

Author: nate-chandler

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -2097,14 +2097,33 @@ void MemoryToRegisters::removeSingleBlockAllocation(AllocStackInst *asi) {
     }
   }
 
-  if (!runningVals || !runningVals->isStorageValid) {
+  auto *function = asi->getFunction();
+  if (!function->hasOwnership() || !runningVals ||
+      !runningVals->isStorageValid ||
+      asi->getElementType().isTrivial(function)) {
     return;
   }
   // There is still valid storage after visiting all instructions in this
   // block which are the only instructions involving this alloc_stack.
   // That can happen if:
   // (1) this block is a dead-end. TODO: OSSACompleteLifetime: Complete such
   //                                     lifetimes.
+  // (2) a trivial case of a non-trivial enum was stored to the address
+
+  auto *deadEndBlocks = deadEndBlocksAnalysis->get(function);
+
+  if (!deadEndBlocks->isDeadEnd(parentBlock)) {
+    // We may have incomplete lifetimes for enum locations on trivial paths.
+    // After promoting them, complete lifetime here.
+    ASSERT(asi->getElementType().isOrHasEnum());
+    OSSACompleteLifetime completion(function, domInfo, *deadEndBlocks,
+                                    OSSACompleteLifetime::IgnoreTrivialVariable,
+                                    /*forceLivenessVerification=*/false,
+                                    /*nonDestroyingEnd=*/true);
+    completion.completeOSSALifetime(
+        runningVals->value.replacement(asi, nullptr),
+        OSSACompleteLifetime::Boundary::Liveness);
+  }
 }
 
 void MemoryToRegisters::collectStoredValues(AllocStackInst *asi,

test/SILOptimizer/mem2reg_ossa_nontrivial.sil

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