Merge pull request #68325 from meg-gupta/fixnonestoremem2reg

Fix handling of none value stores to non-trivial enums in Mem2Reg

Author: meg-gupta

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -203,13 +203,6 @@ class LiveValues {
     return LiveValues::forOwned({stored, move});
   }
 
-  static LiveValues forValues(SILValue stored, SILValue lexical) {
-    if (stored->getOwnershipKind() == OwnershipKind::Guaranteed) {
-      return LiveValues::forGuaranteed({stored, lexical});
-    }
-    return LiveValues::forOwned({stored, lexical});
-  }
-
   static LiveValues toReplace(AllocStackInst *asi, SILValue replacement) {
     if (replacement->getOwnershipKind() == OwnershipKind::Guaranteed) {
       return LiveValues::forGuaranteed(Guaranteed::toReplace(asi, replacement));
@@ -1290,7 +1283,9 @@ StackAllocationPromoter::getLiveOutValues(BasicBlockSetVector &phiBlocks,
       auto *inst = it->second;
       auto stored = inst->getOperand(CopyLikeInstruction::Src);
       auto lexical = getLexicalValueForStore(inst, asi);
-      return LiveValues::forValues(stored, lexical);
+      return isa<StoreBorrowInst>(inst)
+                 ? LiveValues::forGuaranteed(stored, lexical)
+                 : LiveValues::forOwned(stored, lexical);
     }
 
     // If there is a Phi definition in this block:

test/SILOptimizer/mem2reg_ossa_nontrivial.sil

@@ -1227,3 +1227,69 @@ entry(%instance : @owned $AnyObject):
   %82 = tuple ()
   return %82 : $()
 }
+
+// CHECK-LABEL: sil [ossa] @test_enum_store_borrow_none :
+// CHECK:       alloc_stack
+// CHECK-NOT:   alloc_stack
+// CHECK:       dealloc_stack
+// CHECK-LABEL: } // end sil function 'test_enum_store_borrow_none'
+sil [ossa] @test_enum_store_borrow_none : $@convention(thin) () -> () {
+bb0:
+  %1 = alloc_stack [lexical] $FakeOptional<Klass>
+  %none = enum $FakeOptional<Klass>, #FakeOptional.none!enumelt
+  %2 = store_borrow %none to %1 : $*FakeOptional<Klass>
+  %3 = alloc_stack $FakeOptional<Klass>
+  %4 = load [copy] %2 : $*FakeOptional<Klass>
+  store %4 to [init] %3 : $*FakeOptional<Klass>
+  switch_enum_addr %3 : $*FakeOptional<Klass>, case #FakeOptional.some!enumelt: bb1, case #FakeOptional.none!enumelt: bb2
+
+bb1:
+  %7 = unchecked_take_enum_data_addr %3 : $*FakeOptional<Klass>, #FakeOptional.some!enumelt
+  %8 = load [take] %7 : $*Klass
+  destroy_value %8 : $Klass
+  dealloc_stack %3 : $*FakeOptional<Klass>
+  br bb3
+
+bb2:
+  dealloc_stack %3 : $*FakeOptional<Klass>
+  br bb3
+
+bb3:
+  end_borrow %2 : $*FakeOptional<Klass>
+  dealloc_stack %1 : $*FakeOptional<Klass>
+  %t = tuple ()
+  return %t : $()
+}
+
+// CHECK-LABEL: sil [ossa] @test_enum_store_none :
+// CHECK:       alloc_stack
+// CHECK-NOT:   alloc_stack
+// CHECK:       dealloc_stack
+// CHECK-LABEL: } // end sil function 'test_enum_store_none'
+sil [ossa] @test_enum_store_none : $@convention(thin) () -> () {
+bb0:
+  %1 = alloc_stack [lexical] $FakeOptional<Klass>
+  %none = enum $FakeOptional<Klass>, #FakeOptional.none!enumelt
+  store %none to [init] %1 : $*FakeOptional<Klass>
+  %3 = alloc_stack $FakeOptional<Klass>
+  %4 = load [copy] %1 : $*FakeOptional<Klass>
+  store %4 to [init] %3 : $*FakeOptional<Klass>
+  switch_enum_addr %3 : $*FakeOptional<Klass>, case #FakeOptional.some!enumelt: bb1, case #FakeOptional.none!enumelt: bb2
+
+bb1:
+  %7 = unchecked_take_enum_data_addr %3 : $*FakeOptional<Klass>, #FakeOptional.some!enumelt
+  %8 = load [take] %7 : $*Klass
+  destroy_value %8 : $Klass
+  dealloc_stack %3 : $*FakeOptional<Klass>
+  br bb3
+
+bb2:
+  dealloc_stack %3 : $*FakeOptional<Klass>
+  br bb3
+
+bb3:
+  destroy_addr %1 : $*FakeOptional<Klass>
+  dealloc_stack %1 : $*FakeOptional<Klass>
+  %t = tuple ()
+  return %t : $()
+}