[Bitcode] Fix incomplete deduplication of PHI entries

Author: FlStamerGS

llvm/include/llvm/IR/BasicBlock.h

@@ -703,6 +703,10 @@ class BasicBlock final : public Value, // Basic blocks are data objects also
   /// block \p New instead of to it.
   LLVM_ABI void replaceSuccessorsPhiUsesWith(BasicBlock *New);
 
+  /// Update all phi nodes in this basic block by deduplicating the references
+  /// to \p BB.
+  LLVM_ABI void deduplicatePhiUsesOf(BasicBlock *BB);
+
   /// Return true if this basic block is an exception handling block.
   bool isEHPad() const { return getFirstNonPHIIt()->isEHPad(); }
 

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -6105,18 +6105,14 @@ Error BitcodeReader::parseFunctionBody(Function *F) {
         // seen value here, to avoid expanding a constant expression multiple
         // times.
         auto It = Args.find(BB);
-        BasicBlock *EdgeBB = ConstExprEdgeBBs.lookup({BB, CurBB});
         if (It != Args.end()) {
-          // If this predecessor was also replaced with a constexpr basic
-          // block, it must be de-duplicated.
-          if (!EdgeBB) {
-            PN->addIncoming(It->second, BB);
-          }
+          PN->addIncoming(It->second, BB);
           continue;
         }
 
         // If there already is a block for this edge (from a different phi),
         // use it.
+        BasicBlock *EdgeBB = ConstExprEdgeBBs.lookup({BB, CurBB});
         if (!EdgeBB) {
           // Otherwise, use a temporary block (that we will discard if it
           // turns out to be unnecessary).
@@ -6944,6 +6940,7 @@ Error BitcodeReader::parseFunctionBody(Function *F) {
     BranchInst::Create(To, EdgeBB);
     From->getTerminator()->replaceSuccessorWith(To, EdgeBB);
     To->replacePhiUsesWith(From, EdgeBB);
+    To->deduplicatePhiUsesOf(EdgeBB);
     EdgeBB->moveBefore(To);
   }
 

llvm/lib/IR/BasicBlock.cpp

@@ -658,6 +658,26 @@ void BasicBlock::replaceSuccessorsPhiUsesWith(BasicBlock *New) {
   this->replaceSuccessorsPhiUsesWith(this, New);
 }
 
+void BasicBlock::deduplicatePhiUsesOf(BasicBlock *BB) {
+  // N.B. This might not be a complete BasicBlock, so don't assume
+  // that it ends with a non-phi instruction.
+  for (Instruction &I : *this) {
+    PHINode *PN = dyn_cast<PHINode>(&I);
+    if (!PN)
+      break;
+    // Since the order of basic blocks in a PHINode are sorted we can
+    // use the index of \p BB + 1 as the first index to check for duplicates.
+    unsigned Idx = PN->getBasicBlockIndex(BB) + 1;
+    while (Idx < PN->getNumIncomingValues()) {
+      BasicBlock *DuplicateBB = PN->getIncomingBlock(Idx);
+      if (DuplicateBB != BB) {
+        break;
+      }
+      PN->removeIncomingValue(Idx);
+    }
+  }
+}
+
 bool BasicBlock::isLandingPad() const {
   return isa<LandingPadInst>(getFirstNonPHIIt());
 }

llvm/test/Bitcode/constexpr-to-instr-dups.ll

@@ -28,3 +28,30 @@ cont:
                  [1, %nonconst]
   ret i32 %res
 }
+
+define i32 @test_multiple_phis(i32 %arg) {
+entry:
+  switch i32 %arg, label %cont [
+    i32 1, label %cont
+    i32 2, label %nonconst
+  ]
+
+nonconst:                                         ; preds = %entry
+  %cmp = icmp ne i32 %arg, 2
+  br i1 %cmp, label %cont, label %cont
+
+; CHECK-LABEL: phi.constexpr:
+; CHECK-NEXT:    %constexpr = ptrtoint ptr @foo to i32
+; CHECK-NEXT:    %constexpr1 = or i32 %constexpr, 5
+; CHECK-NEXT:    br label %cont
+
+
+; CHECK-LABEL: cont:
+; CHECK-NEXT:    %phi1 = phi i32 [ 0, %phi.constexpr ], [ 1, %nonconst ], [ 1, %nonconst ]
+; CHECK-NEXT:    %phi2 = phi i32 [ %constexpr1, %phi.constexpr ], [ 1, %nonconst ], [ 1, %nonconst ]
+; CHECK-NEXT:    ret i32 %phi2
+cont:                                             ; preds = %nonconst, %nonconst, %entry, %entry
+  %phi1 = phi i32 [ 0, %entry ], [ 0, %entry ], [ 1, %nonconst ], [ 1, %nonconst ]
+  %phi2 = phi i32 [ or (i32 ptrtoint (ptr @foo to i32), i32 5), %entry ], [ or (i32 ptrtoint (ptr @foo to i32), i32 5), %entry ], [ 1, %nonconst ], [ 1, %nonconst ]
+  ret i32 %phi2
+}

llvm/test/Bitcode/constexpr-to-instr-dups.ll.bc

@@ -94,6 +94,64 @@ TEST(BasicBlockTest, PhiRange) {
   }
 }
 
+TEST(BasicBlockTest, DeduplicatePhiEntries) {
+  LLVMContext Context;
+
+  // Create the main block.
+  std::unique_ptr<BasicBlock> BB(BasicBlock::Create(Context));
+
+  // Create some predecessors of it.
+  std::unique_ptr<BasicBlock> BB1(BasicBlock::Create(Context));
+  BranchInst::Create(BB.get(), BB1.get());
+
+  // Make sure this doesn't crash if there are no phis.
+  int PhiCount = 0;
+  for (auto &PN : BB->phis()) {
+    (void)PN;
+    PhiCount++;
+  }
+  ASSERT_EQ(PhiCount, 0) << "empty block should have no phis";
+
+  // Make it a cycle.
+  auto *BI = BranchInst::Create(BB.get(), BB.get());
+
+  // Now insert some PHI nodes.
+  auto *Int32Ty = Type::getInt32Ty(Context);
+  auto *P1 = PHINode::Create(Int32Ty, /*NumReservedValues*/ 3, "phi.1", BI);
+  auto *P2 = PHINode::Create(Int32Ty, /*NumReservedValues*/ 3, "phi.2", BI);
+
+  // Some non-PHI nodes.
+  auto *Sum = BinaryOperator::CreateAdd(P1, P2, "sum", BI);
+
+  // Now wire up the incoming values that are interesting.
+  P1->addIncoming(P2, BB.get());
+  P2->addIncoming(Sum, BB.get());
+
+  // Wire up the rest of the incoming values, we want duplicate entries.
+  for (auto &PN : BB->phis()) {
+    auto *SomeValue = UndefValue::get(Int32Ty);
+    PN.addIncoming(SomeValue, BB1.get());
+    PN.addIncoming(SomeValue, BB1.get());
+  }
+
+  // Check that the phi nodes have 3 incoming values.
+  EXPECT_EQ(P1->getNumIncomingValues(), 3u);
+  EXPECT_EQ(P2->getNumIncomingValues(), 3u);
+
+  // Deduplicate the incoming values for BB1.
+  BB->deduplicatePhiUsesOf(BB1.get());
+
+  // Check that the phi nodes have 2 incoming values.
+  EXPECT_EQ(P1->getNumIncomingValues(), 2u);
+  EXPECT_EQ(P2->getNumIncomingValues(), 2u);
+
+  // Check that the incoming values are unique.
+  for (auto &PN : BB->phis()) {
+    EXPECT_EQ(BB.get(), PN.getIncomingBlock(0));
+    EXPECT_EQ(BB1.get(), PN.getIncomingBlock(1));
+  }
+}
+
 #define CHECK_ITERATORS(Range1, Range2)                                        \
   EXPECT_EQ(std::distance(Range1.begin(), Range1.end()),                       \
             std::distance(Range2.begin(), Range2.end()));                      \