Enable mem2reg of enum types

Previously it was disabled because we have incomplete address lifetimes for such types,
and transforming to value form in mem2reg would expose verification errors due to incompleteness.

Enable this case here and fixup the lifetimes using the new lifetime completion utility.

Author: meg-gupta

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -26,6 +26,7 @@
 #include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SIL/Dominance.h"
 #include "swift/SIL/Projection.h"
+#include "swift/SIL/OSSALifetimeCompletion.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
@@ -1784,13 +1785,42 @@ void StackAllocationPromoter::promoteAllocationToPhi() {
 }
 
 void StackAllocationPromoter::run() {
+  auto *function = asi->getFunction();
+
   // Reduce the number of load/stores in the function to minimum.
   // After this phase we are left with up to one load and store
   // per block and the last store is recorded.
   pruneAllocStackUsage();
 
   // Replace AllocStacks with Phi-nodes.
   promoteAllocationToPhi();
+
+  // Make sure that all of the allocations were promoted into registers.
+  assert(isWriteOnlyAllocation(asi) && "Non-write uses left behind");
+
+  SmallVector<SILValue> valuesToComplete;
+
+  // Enum types may have incomplete lifetimes in address form, when promoted to
+  // value form after mem2reg, they will end up with incomplete ossa lifetimes.
+  // Use the lifetime completion utility to complete such lifetimes.
+  // First, collect the stored values to complete.
+  if (asi->getType().isOrHasEnum()) {
+    for (auto it : initializationPoints) {
+      auto *si = it.second;
+      auto src = si->getOperand(0);
+      valuesToComplete.push_back(src);
+    }
+  }
+
+  // ... and erase the allocation.
+  deleter.forceDeleteWithUsers(asi);
+
+  // Now, complete lifetimes!
+  OSSALifetimeCompletion completion(function, domInfo);
+
+  for (auto it : valuesToComplete) {
+    completion.completeOSSALifetime(it);
+  }
 }
 
 //===----------------------------------------------------------------------===//
@@ -2091,12 +2121,16 @@ bool MemoryToRegisters::promoteSingleAllocation(AllocStackInst *alloc) {
     deleter.forceDeleteWithUsers(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())
+    // For lexical stack locs with enum type, 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.
+    // For non-lexical stack locs, we use the lifetime completion utility
+    // to fix this, but for lexical stack locs, we may create move_value
+    // [lexical] whose lifetime is not completed since the utility depends on
+    // the presence of unreachable blocks.
+    if (f.hasOwnership() && alloc->isLexical() && alloc->getType().isOrHasEnum())
       return false;
   }
 
@@ -2108,11 +2142,6 @@ bool MemoryToRegisters::promoteSingleAllocation(AllocStackInst *alloc) {
   StackAllocationPromoter(alloc, domInfo, domTreeLevels, ctx, deleter,
                           instructionsToDelete)
       .run();
-
-  // Make sure that all of the allocations were promoted into registers.
-  assert(isWriteOnlyAllocation(alloc) && "Non-write uses left behind");
-  // ... and erase the allocation.
-  deleter.forceDeleteWithUsers(alloc);
   return true;
 }
 

test/SILOptimizer/mem2reg_lifetime_nontrivial.sil

@@ -1061,7 +1061,6 @@ enum KlassOptional {
 sil @return_optional_or_error : $@convention(thin) () -> (@owned KlassOptional, @error Error)
 
 // CHECK-LABEL: sil [ossa] @test_deadphi4 :
-// 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) -> () {

test/SILOptimizer/mem2reg_ossa.sil

@@ -489,10 +489,10 @@ bb0:
   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>) -> () {
+// CHECK-LABEL: sil [ossa] @test_optional_in_multiple_blocks :
+// CHECK-NOT:         alloc_stack
+// CHECK:       } // end sil function 'test_optional_in_multiple_blocks'
+sil [ossa] @test_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>
@@ -504,8 +504,35 @@ bb5:
   br bb7
 
 bb6(%50 : @guaranteed $Klass):
-  %53 = load [take] %32 : $*Optional<Klass>
+  %53 = load [copy] %32 : $*Optional<Klass>
+  destroy_value %53 : $Optional<Klass>
+  destroy_addr %32 : $*Optional<Klass>
+  dealloc_stack %32 : $*Optional<Klass>
+  br bb7
+
+bb7:
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_optional_in_multiple_blocks_lexical :
+// CHECK:         alloc_stack
+// CHECK:       } // end sil function 'test_optional_in_multiple_blocks_lexical'
+sil [ossa] @test_optional_in_multiple_blocks_lexical : $@convention(method) (@guaranteed Optional<Klass>) -> () {
+bb0(%0 : @guaranteed $Optional<Klass>):
+  %1 = copy_value %0 : $Optional<Klass>
+  %32 = alloc_stack [lexical] $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 [copy] %32 : $*Optional<Klass>
   destroy_value %53 : $Optional<Klass>
+  destroy_addr %32 : $*Optional<Klass>
   dealloc_stack %32 : $*Optional<Klass>
   br bb7
 

test/SILOptimizer/mem2reg_ossa_nontrivial.sil

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