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

@@ -25,6 +25,7 @@
 #include "swift/Basic/TaggedUnion.h"
 #include "swift/SIL/BasicBlockDatastructures.h"
 #include "swift/SIL/Dominance.h"
+#include "swift/SIL/OSSALifetimeCompletion.h"
 #include "swift/SIL/Projection.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILFunction.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,31 @@ 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())
+    auto enableOptimizationForEnum = [](AllocStackInst *asi) {
+      if (asi->isLexical()) {
+        return false;
+      }
+      for (auto *use : asi->getUses()) {
+        auto *user = use->getUser();
+        if (!isa<StoreInst>(user) && !isa<StoreBorrowInst>(user)) {
+          continue;
+        }
+        auto stored = user->getOperand(CopyLikeInstruction::Src);
+        if (stored->isLexical()) {
+          return false;
+        }
+      }
+      return true;
+    };
+    // For stack locs of enum type that are lexical or with lexical stored
+    // values, we require that all uses are in the same block. This is because
+    // we can have incomplete lifetime of enum typed addresses, and on
+    // converting to value form this causes verification error. For all other
+    // stack locs of enum type, we use the lifetime completion utility to fix
+    // the lifetime. But when we have a lexical value, the utility can complete
+    // lifetimes on dead end blocks only.
+    if (f.hasOwnership() && alloc->getType().isOrHasEnum() &&
+        !enableOptimizationForEnum(alloc))
       return false;
   }
 
@@ -2108,11 +2157,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,62 @@ 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
+
+bb7:
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_optional_in_multiple_blocks_lexical_storedvalue :
+// CHECK:         alloc_stack
+// CHECK:       } // end sil function 'test_optional_in_multiple_blocks_lexical_storedvalue'
+sil [ossa] @test_optional_in_multiple_blocks_lexical_storedvalue : $@convention(method) (@owned Optional<Klass>) -> () {
+bb0(%0 : @owned $Optional<Klass>):
+  %1 = copy_value %0 : $Optional<Klass>
+  %32 = alloc_stack $Optional<Klass>
+  store %0 to [init] %32 : $*Optional<Klass>
+  switch_enum %1 : $Optional<Klass>, case #Optional.some!enumelt: bb6, case #Optional.none!enumelt: bb5
+
+bb5:
+  dealloc_stack %32 : $*Optional<Klass>
+  br bb7
+
+bb6(%50 : @owned $Klass):
+  %53 = load [copy] %32 : $*Optional<Klass>
   destroy_value %53 : $Optional<Klass>
+  destroy_value %50 : $Klass
+  destroy_addr %32 : $*Optional<Klass>
   dealloc_stack %32 : $*Optional<Klass>
   br bb7
 

test/SILOptimizer/mem2reg_ossa_nontrivial.sil

@@ -53,6 +53,11 @@ struct UInt8 {
   var _value : Builtin.Int8
 }
 
+enum KlassOptional {
+  case some(Klass)
+  case none
+}
+
 sil [ossa] @get_nontrivialstruct : $@convention(thin) () -> @owned NonTrivialStruct
 sil [ossa] @get_nontrivialenum : $@convention(thin) () -> @owned NonTrivialEnum
 sil [ossa] @get_optionalnontrivialstruct : $@convention(thin) () -> @owned FakeOptional<NonTrivialStruct>
@@ -1026,16 +1031,10 @@ bb12:
   return %39 : $()
 }
 
-enum KlassOptional {
-  case some(Klass)
-  case none
-}
-
 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):