Fix DCE for ownership forwarding instructions

lib/SILOptimizer/Transforms/DeadCodeElimination.cpp

@@ -22,7 +22,9 @@
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILUndef.h"
+#include "swift/SIL/OSSALifetimeCompletion.h"
 #include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
+#include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/BasicBlockOptUtils.h"
@@ -126,7 +128,10 @@ class DCE {
   BasicBlockSet LiveBlocks;
   llvm::SmallVector<SILInstruction *, 64> Worklist;
   PostDominanceInfo *PDT;
+  DominanceInfo *DT;
+  DeadEndBlocks *deadEndBlocks;
   llvm::DenseMap<SILBasicBlock *, ControllingInfo> ControllingInfoMap;
+  llvm::SmallVector<SILValue> valuesToComplete;
 
   // Maps instructions which produce a failing condition (like overflow
   // builtins) to the actual cond_fail instructions which handle the failure.
@@ -194,8 +199,10 @@ class DCE {
   void endLifetimeOfLiveValue(Operand *op, SILInstruction *insertPt);
 
 public:
-  DCE(SILFunction *F, PostDominanceInfo *PDT)
-    : F(F), LiveArguments(F), LiveInstructions(F), LiveBlocks(F), PDT(PDT) {}
+  DCE(SILFunction *F, PostDominanceInfo *PDT, DominanceInfo *DT,
+      DeadEndBlocks *deadEndBlocks)
+      : F(F), LiveArguments(F), LiveInstructions(F), LiveBlocks(F), PDT(PDT),
+        DT(DT), deadEndBlocks(deadEndBlocks) {}
 
   /// The entry point to the transformation.
   bool run() {
@@ -626,6 +633,14 @@ void DCE::endLifetimeOfLiveValue(Operand *op, SILInstruction *insertPt) {
   }
 
   assert(op->isLifetimeEnding());
+
+  // If DCE is going to delete the block in which we have to insert a
+  // compensating lifetime end, let complete lifetimes utility handle it.
+  if (!LiveBlocks.contains(insertPt->getParent())) {
+    valuesToComplete.push_back(value);
+    return;
+  }
+
   SILBuilderWithScope builder(insertPt);
   if (value->getOwnershipKind() == OwnershipKind::Owned) {
     builder.createDestroyValue(RegularLocation::getAutoGeneratedLocation(),
@@ -772,6 +787,12 @@ bool DCE::removeDead() {
     }
   }
 
+  OSSALifetimeCompletion completion(F, DT, *deadEndBlocks);
+  for (auto value : valuesToComplete) {
+    completion.completeOSSALifetime(value,
+                                    OSSALifetimeCompletion::Boundary::Liveness);
+  }
+
   return Changed;
 }
 
@@ -977,8 +998,11 @@ class DCEPass : public SILFunctionTransform {
     LLVM_DEBUG(llvm::dbgs() << "*** DCE on function: " << F->getName()
                             << " ***\n");
 
-    auto *DA = PM->getAnalysis<PostDominanceAnalysis>();
-    PostDominanceInfo *PDT = DA->get(F);
+    auto *PDA = PM->getAnalysis<PostDominanceAnalysis>();
+    PostDominanceInfo *PDT = PDA->get(F);
+
+    auto *DA = PM->getAnalysis<DominanceAnalysis>();
+    auto *DEA = getAnalysis<DeadEndBlocksAnalysis>();
 
     // If we have a function that consists of nothing but a
     // structurally infinite loop like:
@@ -987,7 +1011,7 @@ class DCEPass : public SILFunctionTransform {
     if (!PDT->getRootNode())
       return;
 
-    DCE dce(F, PDT);
+    DCE dce(F, PDT, DA->get(F), DEA->get(F));
     if (dce.run()) {
       using InvalidationKind = SILAnalysis::InvalidationKind;
       unsigned Inv = InvalidationKind::Instructions;

test/SILOptimizer/dead_code_elimination_nontrivial_ossa.sil

@@ -1280,3 +1280,99 @@ bb2(%phi1 : @guaranteed $NonTrivialStruct, %phi2 : @guaranteed $NonTrivialStruct
   return %9999 : $()
 }
 
+// CHECK-LABEL: sil [ossa] @dce_deadterm4 :
+// CHECK-NOT: switch_enum
+// CHECK-LABEL: } // end sil function 'dce_deadterm4'
+sil [ossa] @dce_deadterm4 : $@convention(thin) (@owned FakeOptional<Builtin.NativeObject>) -> () {
+bb0(%0 : @owned $FakeOptional<Builtin.NativeObject>):
+  %0a = alloc_stack $FakeOptional<Builtin.NativeObject>
+  store %0 to [init] %0a : $*FakeOptional<Builtin.NativeObject>
+  %1 = load [take] %0a : $*FakeOptional<Builtin.NativeObject>
+  %2 = copy_value %1 : $FakeOptional<Builtin.NativeObject>
+  apply undef(%2) : $@convention(thin) (@guaranteed FakeOptional<Builtin.NativeObject>) -> ()
+  switch_enum %1 : $FakeOptional<Builtin.NativeObject>, case #FakeOptional.some!enumelt: bb1, case #FakeOptional.none!enumelt: bb2
+
+bb1(%3 : @owned $Builtin.NativeObject):
+  destroy_value %3 : $Builtin.NativeObject
+  br bb3
+
+bb2:
+  br bb3
+
+bb3:
+  switch_enum %2 : $FakeOptional<Builtin.NativeObject>, case #FakeOptional.some!enumelt: bb4, case #FakeOptional.none!enumelt: bb5
+
+bb4(%4 : @owned $Builtin.NativeObject):
+  destroy_value %4 : $Builtin.NativeObject
+  br bb6
+
+bb5:
+  br bb6
+
+bb6:
+  dealloc_stack %0a : $*FakeOptional<Builtin.NativeObject>
+  %9999 = tuple()
+  return %9999 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @dce_deadterm5 :
+// CHECK: switch_enum
+// CHECK-NOT: switch_enum
+// CHECK-LABEL: } // end sil function 'dce_deadterm5'
+sil [ossa] @dce_deadterm5 : $@convention(thin) (@owned FakeOptional<Builtin.NativeObject>) -> () {
+bb0(%0 : @owned $FakeOptional<Builtin.NativeObject>):
+  %0a = alloc_stack $FakeOptional<Builtin.NativeObject>
+  store %0 to [init] %0a : $*FakeOptional<Builtin.NativeObject>
+  %1 = load [take] %0a : $*FakeOptional<Builtin.NativeObject>
+  %2 = copy_value %1 : $FakeOptional<Builtin.NativeObject>
+  apply undef(%2) : $@convention(thin) (@guaranteed FakeOptional<Builtin.NativeObject>) -> ()
+  switch_enum %1 : $FakeOptional<Builtin.NativeObject>, case #FakeOptional.some!enumelt: bb1, case #FakeOptional.none!enumelt: bb2
+
+bb1(%3 : @owned $Builtin.NativeObject):
+  destroy_value %3 : $Builtin.NativeObject
+  br bb3
+
+bb2:
+  switch_enum %2 : $FakeOptional<Builtin.NativeObject>, case #FakeOptional.some!enumelt: bb4, case #FakeOptional.none!enumelt: bb5
+
+bb3:
+  destroy_value %2 : $FakeOptional<Builtin.NativeObject>
+  br bb6
+
+bb4(%4 : @owned $Builtin.NativeObject):
+  destroy_value %4 : $Builtin.NativeObject
+  br bb6
+
+bb5:
+  br bb6
+
+bb6:
+  dealloc_stack %0a : $*FakeOptional<Builtin.NativeObject>
+  %9999 = tuple()
+  return %9999 : $()
+}
+
+// CHECK-LABEL: sil [ossa] @dce_deadforwarding :
+// CHECK-NOT: cond_br
+// CHECK-LABEL: } // end sil function 'dce_deadforwarding'
+sil [ossa] @dce_deadforwarding : $@convention(thin) (@guaranteed Klass) -> () {
+bb0(%0 : @guaranteed $Klass):
+  %1 = copy_value %0 : $Klass
+  %func = function_ref @$use_klass2 : $@convention(thin) (@guaranteed Klass) -> ()
+  %funcres = apply %func(%1) : $@convention(thin) (@guaranteed Klass) -> ()
+  cond_br undef, bb1, bb2
+
+bb1:
+  %3 = struct $NonTrivialStruct(%1 : $Klass)
+  destroy_value %3 : $NonTrivialStruct
+  br bb3
+
+bb2:
+  %4 = struct $NonTrivialStruct(%1 : $Klass)
+  destroy_value %4 : $NonTrivialStruct
+  br bb3
+
+bb3:
+  %res = tuple ()
+  return %res : $()
+}