[SandboxIR] Implement missing PHINode functions (#101734)

replaceIncomingBlockWith and removeIncomingValueIf are both
straightforward and done.

I'll defer copyIncomingBlocks until a couple of other changes that also
handle blocks go in.

Author: Sterling-Augustine

llvm/include/llvm/SandboxIR/SandboxIR.h

@@ -1607,12 +1607,11 @@ class PHINode final : public Instruction {
     return cast<llvm::PHINode>(Val)->hasConstantOrUndefValue();
   }
   bool isComplete() const { return cast<llvm::PHINode>(Val)->isComplete(); }
-  // TODO: Implement the below functions:
-  // void replaceIncomingBlockWith (const BasicBlock *Old, BasicBlock *New);
+  void replaceIncomingBlockWith(const BasicBlock *Old, BasicBlock *New);
+  void removeIncomingValueIf(function_ref<bool(unsigned)> Predicate);
+  // TODO: Implement
   // void copyIncomingBlocks(iterator_range<const_block_iterator> BBRange,
   //                         uint32_t ToIdx = 0)
-  // void removeIncomingValueIf(function_ref< bool(unsigned)> Predicate,
-  //                            bool DeletePHIIfEmpty=true)
 #ifndef NDEBUG
   void verify() const final {
     assert(isa<llvm::PHINode>(Val) && "Expected PHINode!");

llvm/lib/SandboxIR/SandboxIR.cpp

@@ -1145,7 +1145,6 @@ Value *PHINode::removeIncomingValue(unsigned Idx) {
   auto &Tracker = Ctx.getTracker();
   if (Tracker.isTracking())
     Tracker.track(std::make_unique<PHIRemoveIncoming>(*this, Idx, Tracker));
-
   llvm::Value *LLVMV =
       cast<llvm::PHINode>(Val)->removeIncomingValue(Idx,
                                                     /*DeletePHIIfEmpty=*/false);
@@ -1177,6 +1176,27 @@ Value *PHINode::hasConstantValue() const {
   llvm::Value *LLVMV = cast<llvm::PHINode>(Val)->hasConstantValue();
   return LLVMV != nullptr ? Ctx.getValue(LLVMV) : nullptr;
 }
+void PHINode::replaceIncomingBlockWith(const BasicBlock *Old, BasicBlock *New) {
+  assert(New && Old && "Sandbox IR PHI node got a null basic block!");
+  for (unsigned Idx = 0, NumOps = cast<llvm::PHINode>(Val)->getNumOperands();
+       Idx != NumOps; ++Idx)
+    if (getIncomingBlock(Idx) == Old)
+      setIncomingBlock(Idx, New);
+}
+void PHINode::removeIncomingValueIf(function_ref<bool(unsigned)> Predicate) {
+  // Avoid duplicate tracking by going through this->removeIncomingValue here at
+  // the expense of some performance. Copy PHI::removeIncomingValueIf more
+  // directly if performance becomes an issue.
+
+  // Removing the element at index X, moves the element previously at X + 1
+  // to X. Working from the end avoids complications from that.
+  unsigned Idx = getNumIncomingValues();
+  while (Idx > 0) {
+    if (Predicate(Idx - 1))
+      removeIncomingValue(Idx - 1);
+    --Idx;
+  }
+}
 
 static llvm::Instruction::CastOps getLLVMCastOp(Instruction::Opcode Opc) {
   switch (Opc) {

llvm/unittests/SandboxIR/SandboxIRTest.cpp

@@ -1925,10 +1925,17 @@ define void @foo(i32 %arg) {
   br label %bb2
 
 bb2:
-  %phi = phi i32 [ %arg, %bb1 ], [ 0, %bb2 ]
+  %phi = phi i32 [ %arg, %bb1 ], [ 0, %bb2 ], [ 1, %bb3 ], [ 2, %bb4 ], [ 3, %bb5 ]
   br label %bb2
 
 bb3:
+  br label %bb2
+
+bb4:
+  br label %bb2
+
+bb5:
+  br label %bb2
   ret void
 }
 )IR");
@@ -2023,7 +2030,29 @@ define void @foo(i32 %arg) {
   EXPECT_EQ(PHI->hasConstantOrUndefValue(), LLVMPHI->hasConstantOrUndefValue());
   // Check isComplete().
   EXPECT_EQ(PHI->isComplete(), LLVMPHI->isComplete());
-
+  // Check replaceIncomingValueIf
+  EXPECT_EQ(PHI->getNumIncomingValues(), 5u);
+  auto *RemainBB0 = PHI->getIncomingBlock(0);
+  auto *RemoveBB0 = PHI->getIncomingBlock(1);
+  auto *RemainBB1 = PHI->getIncomingBlock(2);
+  auto *RemoveBB1 = PHI->getIncomingBlock(3);
+  auto *RemainBB2 = PHI->getIncomingBlock(4);
+  PHI->removeIncomingValueIf([&](unsigned Idx) {
+    return PHI->getIncomingBlock(Idx) == RemoveBB0 ||
+           PHI->getIncomingBlock(Idx) == RemoveBB1;
+  });
+  EXPECT_EQ(PHI->getNumIncomingValues(), 3u);
+  EXPECT_EQ(PHI->getIncomingBlock(0), RemainBB0);
+  EXPECT_EQ(PHI->getIncomingBlock(1), RemainBB1);
+  EXPECT_EQ(PHI->getIncomingBlock(2), RemainBB2);
+  // Check replaceIncomingBlockWith
+  OrigBB = RemainBB0;
+  auto *NewBB = RemainBB1;
+  EXPECT_NE(NewBB, OrigBB);
+  PHI->replaceIncomingBlockWith(OrigBB, NewBB);
+  EXPECT_EQ(PHI->getIncomingBlock(0), NewBB);
+  EXPECT_EQ(PHI->getIncomingBlock(1), RemainBB1);
+  EXPECT_EQ(PHI->getIncomingBlock(2), RemainBB2);
   // Check create().
   auto *NewPHI = cast<sandboxir::PHINode>(
       sandboxir::PHINode::create(PHI->getType(), 0, Br, Ctx, "NewPHI"));

llvm/unittests/SandboxIR/TrackerTest.cpp

@@ -779,6 +779,15 @@ define void @foo(i8 %arg0, i8 %arg1, i8 %arg2) {
   EXPECT_EQ(PHI->getIncomingBlock(1), BB1);
   EXPECT_EQ(PHI->getIncomingValue(1), Arg1);
 
+  // Check removeIncomingValueIf(FromBB1).
+  Ctx.save();
+  PHI->removeIncomingValueIf(
+      [&](unsigned Idx) { return PHI->getIncomingBlock(Idx) == BB1; });
+  EXPECT_EQ(PHI->getNumIncomingValues(), 1u);
+  Ctx.revert();
+  EXPECT_EQ(PHI->getNumIncomingValues(), 2u);
+  EXPECT_EQ(PHI->getIncomingBlock(0), BB0);
+  EXPECT_EQ(PHI->getIncomingBlock(1), BB1);
   // Check removeIncomingValue() remove all.
   Ctx.save();
   PHI->removeIncomingValue(0u);