[SimplifyCFG] 'merge compatible invokes': support normal destination w/ uses

If the original invokes had uses, the uses must have been in PHI's,
but that immediately results in the incoming values being incompatible.
But we'll replace uses of the original invokes with the use of the
merged invoke, so as long as the incoming values become compatible
after that, we can merge.

Author: LebedevRI

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -296,17 +296,28 @@ class SimplifyCFGOpt {
 /// Return true if all the PHI nodes in the basic block \p BB
 /// receive compatible (identical) incoming values when coming from
 /// all of the predecessor blocks that are specified in \p IncomingBlocks.
-static bool IncomingValuesAreCompatible(BasicBlock *BB,
-                                        ArrayRef<BasicBlock *> IncomingBlocks) {
+///
+/// Note that if the values aren't exactly identical, but \p EquivalenceSet
+/// is provided, and *both* of the values are present in the set,
+/// then they are considered equal.
+static bool IncomingValuesAreCompatible(
+    BasicBlock *BB, ArrayRef<BasicBlock *> IncomingBlocks,
+    SmallPtrSetImpl<Value *> *EquivalenceSet = nullptr) {
   assert(IncomingBlocks.size() == 2 &&
          "Only for a pair of incoming blocks at the time!");
 
   // FIXME: it is okay if one of the incoming values is an `undef` value,
   //        iff the other incoming value is guaranteed to be a non-poison value.
   // FIXME: it is okay if one of the incoming values is a `poison` value.
-  return all_of(BB->phis(), [IncomingBlocks](PHINode &PN) {
-    return PN.getIncomingValueForBlock(IncomingBlocks[0]) ==
-           PN.getIncomingValueForBlock(IncomingBlocks[1]);
+  return all_of(BB->phis(), [IncomingBlocks, EquivalenceSet](PHINode &PN) {
+    Value *IV0 = PN.getIncomingValueForBlock(IncomingBlocks[0]);
+    Value *IV1 = PN.getIncomingValueForBlock(IncomingBlocks[1]);
+    if (IV0 == IV1)
+      return true;
+    if (EquivalenceSet && EquivalenceSet->contains(IV0) &&
+        EquivalenceSet->contains(IV1))
+      return true;
+    return false;
   });
 }
 
@@ -2309,13 +2320,10 @@ bool CompatibleSets::shouldBelongToSameSet(ArrayRef<InvokeInst *> Invokes) {
 
     // In the normal destination, the incoming values for these two `invoke`s
     // must be compatible.
+    SmallPtrSet<Value *, 16> EquivalenceSet(Invokes.begin(), Invokes.end());
     if (!IncomingValuesAreCompatible(
-            NormalBB, {Invokes[0]->getParent(), Invokes[1]->getParent()}))
-      return false;
-
-    // For now, simply don't deal with `invoke`s that have uses.
-    auto Unused = [](InvokeInst *II) { return II->use_empty(); };
-    if (!all_of(Invokes, Unused))
+            NormalBB, {Invokes[0]->getParent(), Invokes[1]->getParent()},
+            &EquivalenceSet))
       return false;
   }
 
@@ -2470,8 +2478,7 @@ static void MergeCompatibleInvokesImpl(ArrayRef<InvokeInst *> Invokes,
     for (BasicBlock *OrigSuccBB : successors(II->getParent()))
       OrigSuccBB->removePredecessor(II->getParent());
     BranchInst::Create(MergedInvoke->getParent(), II->getParent());
-    // Since the normal destination was unreachable, there are no live uses.
-    II->replaceAllUsesWith(UndefValue::get(II->getType()));
+    II->replaceAllUsesWith(MergedInvoke);
     II->eraseFromParent();
     ++NumInvokesMerged;
   }

llvm/test/Transforms/SimplifyCFG/X86/merge-compatible-invokes-of-landingpad.ll

@@ -1688,13 +1688,9 @@ define void @t28_invoke_ret_value_is_used_in_phi_node() personality i8* bitcast
 ; CHECK-LABEL: @t28_invoke_ret_value_is_used_in_phi_node(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[C0:%.*]] = call i1 @cond()
-; CHECK-NEXT:    br i1 [[C0]], label [[IF_THEN0:%.*]], label [[IF_ELSE:%.*]]
-; CHECK:       if.then0:
-; CHECK-NEXT:    [[V0:%.*]] = invoke i32 @returning_maybe_throw()
-; CHECK-NEXT:    to label [[INVOKE_CONT:%.*]] unwind label [[LPAD:%.*]]
+; CHECK-NEXT:    br i1 [[C0]], label [[IF_THEN1_INVOKE:%.*]], label [[IF_ELSE:%.*]]
 ; CHECK:       invoke.cont:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[V0]], [[IF_THEN0]] ], [ [[V1:%.*]], [[IF_THEN1:%.*]] ]
-; CHECK-NEXT:    call void @consume(i32 [[PHI]])
+; CHECK-NEXT:    call void @consume(i32 [[TMP0:%.*]])
 ; CHECK-NEXT:    call void @sideeffect()
 ; CHECK-NEXT:    unreachable
 ; CHECK:       lpad:
@@ -1704,10 +1700,10 @@ define void @t28_invoke_ret_value_is_used_in_phi_node() personality i8* bitcast
 ; CHECK-NEXT:    resume { i8*, i32 } [[EH]]
 ; CHECK:       if.else:
 ; CHECK-NEXT:    [[C1:%.*]] = call i1 @cond()
-; CHECK-NEXT:    br i1 [[C1]], label [[IF_THEN1]], label [[IF_END:%.*]]
-; CHECK:       if.then1:
-; CHECK-NEXT:    [[V1]] = invoke i32 @returning_maybe_throw()
-; CHECK-NEXT:    to label [[INVOKE_CONT]] unwind label [[LPAD]]
+; CHECK-NEXT:    br i1 [[C1]], label [[IF_THEN1_INVOKE]], label [[IF_END:%.*]]
+; CHECK:       if.then1.invoke:
+; CHECK-NEXT:    [[TMP0]] = invoke i32 @returning_maybe_throw()
+; CHECK-NEXT:    to label [[INVOKE_CONT:%.*]] unwind label [[LPAD:%.*]]
 ; CHECK:       if.end:
 ; CHECK-NEXT:    call void @sideeffect()
 ; CHECK-NEXT:    ret void
@@ -1933,15 +1929,10 @@ define void @t32_invoke_ret_value_is_used_in_phi_node_other_phi_is_fine() person
 ; CHECK-LABEL: @t32_invoke_ret_value_is_used_in_phi_node_other_phi_is_fine(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[C0:%.*]] = call i1 @cond()
-; CHECK-NEXT:    br i1 [[C0]], label [[IF_THEN0:%.*]], label [[IF_ELSE:%.*]]
-; CHECK:       if.then0:
-; CHECK-NEXT:    [[V0:%.*]] = invoke i32 @returning_maybe_throw()
-; CHECK-NEXT:    to label [[INVOKE_CONT:%.*]] unwind label [[LPAD:%.*]]
+; CHECK-NEXT:    br i1 [[C0]], label [[IF_THEN1_INVOKE:%.*]], label [[IF_ELSE:%.*]]
 ; CHECK:       invoke.cont:
-; CHECK-NEXT:    [[PHI0:%.*]] = phi i32 [ [[V0]], [[IF_THEN0]] ], [ [[V1:%.*]], [[IF_THEN1:%.*]] ]
-; CHECK-NEXT:    [[PHI1:%.*]] = phi i32 [ 0, [[IF_THEN0]] ], [ 0, [[IF_THEN1]] ]
-; CHECK-NEXT:    call void @consume(i32 [[PHI0]])
-; CHECK-NEXT:    call void @consume(i32 [[PHI1]])
+; CHECK-NEXT:    call void @consume(i32 [[TMP0:%.*]])
+; CHECK-NEXT:    call void @consume(i32 0)
 ; CHECK-NEXT:    call void @sideeffect()
 ; CHECK-NEXT:    unreachable
 ; CHECK:       lpad:
@@ -1951,10 +1942,10 @@ define void @t32_invoke_ret_value_is_used_in_phi_node_other_phi_is_fine() person
 ; CHECK-NEXT:    resume { i8*, i32 } [[EH]]
 ; CHECK:       if.else:
 ; CHECK-NEXT:    [[C1:%.*]] = call i1 @cond()
-; CHECK-NEXT:    br i1 [[C1]], label [[IF_THEN1]], label [[IF_END:%.*]]
-; CHECK:       if.then1:
-; CHECK-NEXT:    [[V1]] = invoke i32 @returning_maybe_throw()
-; CHECK-NEXT:    to label [[INVOKE_CONT]] unwind label [[LPAD]]
+; CHECK-NEXT:    br i1 [[C1]], label [[IF_THEN1_INVOKE]], label [[IF_END:%.*]]
+; CHECK:       if.then1.invoke:
+; CHECK-NEXT:    [[TMP0]] = invoke i32 @returning_maybe_throw()
+; CHECK-NEXT:    to label [[INVOKE_CONT:%.*]] unwind label [[LPAD:%.*]]
 ; CHECK:       if.end:
 ; CHECK-NEXT:    call void @sideeffect()
 ; CHECK-NEXT:    ret void