[llvm] Implement missing PHINode functions (PR #101734)

[via llvm-commits]

https://github.com/Sterling-Augustine created https://github.com/llvm/llvm-project/pull/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.

>From 8ea3fbf24bd71583649cc98e8e088640a3fde8ff Mon Sep 17 00:00:00 2001
From: Sterling Augustine <saugustine at google.com>
Date: Fri, 2 Aug 2024 18:48:14 +0000
Subject: [PATCH] Implement missing PHINode functions

replaceIncomingBlockWith and removeIncomingValueIf are both
straightforward and done.

I'll defer copyIncomingBlocks until a couple of other changes that
also handle blocks go in.
---
 llvm/include/llvm/SandboxIR/SandboxIR.h    |  8 ++++----
 llvm/lib/SandboxIR/SandboxIR.cpp           | 16 ++++++++++++++++
 llvm/unittests/SandboxIR/SandboxIRTest.cpp | 13 ++++++++++++-
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/SandboxIR/SandboxIR.h b/llvm/include/llvm/SandboxIR/SandboxIR.h
index 0bd6cdac36f9a..d47bb4e87c219 100644
--- a/llvm/include/llvm/SandboxIR/SandboxIR.h
+++ b/llvm/include/llvm/SandboxIR/SandboxIR.h
@@ -1585,12 +1585,12 @@ 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,
+                             bool DeletePHIIfEmpty=true);
+  // 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!");
diff --git a/llvm/lib/SandboxIR/SandboxIR.cpp b/llvm/lib/SandboxIR/SandboxIR.cpp
index a67cb2a5bb25a..f34eeee2399ac 100644
--- a/llvm/lib/SandboxIR/SandboxIR.cpp
+++ b/llvm/lib/SandboxIR/SandboxIR.cpp
@@ -1150,6 +1150,22 @@ 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 Op = 0,
+            NumOps = cast<llvm::PHINode>(Val)->getNumOperands(); Op != NumOps; ++Op)
+    if (getIncomingBlock(Op) == Old)
+      setIncomingBlock(Op, New);
+}
+void PHINode::removeIncomingValueIf(function_ref< bool(unsigned)> Predicate,
+                                    bool DeletePHIIfEmpty) {
+  // 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.
+  for (unsigned Idx = 0; Idx < getNumIncomingValues(); ++Idx)
+    if (Predicate(Idx))
+      removeIncomingValue(Idx);
+}
 
 static llvm::Instruction::CastOps getLLVMCastOp(Instruction::Opcode Opc) {
   switch (Opc) {
diff --git a/llvm/unittests/SandboxIR/SandboxIRTest.cpp b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
index 5f858b47926ec..5d22f3b003247 100644
--- a/llvm/unittests/SandboxIR/SandboxIRTest.cpp
+++ b/llvm/unittests/SandboxIR/SandboxIRTest.cpp
@@ -1974,7 +1974,18 @@ define void @foo(i32 %arg) {
   EXPECT_EQ(PHI->hasConstantOrUndefValue(), LLVMPHI->hasConstantOrUndefValue());
   // Check isComplete().
   EXPECT_EQ(PHI->isComplete(), LLVMPHI->isComplete());
-
+  // Check replaceIncomingBlockWith
+  OrigBB = PHI->getIncomingBlock(0);
+  EXPECT_EQ(OrigBB, BB1);
+  EXPECT_NE(OrigBB, BB2);
+  PHI->replaceIncomingBlockWith(BB1, BB2);
+  EXPECT_EQ(PHI->getIncomingBlock(0), BB2);
+  // Check replaceIncomingValueIf
+  EXPECT_EQ(PHI->getNumIncomingValues(), 2u);
+  PHI->removeIncomingValueIf([&](unsigned Idx) {
+    return PHI->getIncomingBlock(Idx) == BB2;
+  });
+  EXPECT_EQ(PHI->getNumIncomingValues(), 1u);
   // Check create().
   auto *NewPHI = cast<sandboxir::PHINode>(
       sandboxir::PHINode::create(PHI->getType(), 0, Br, Ctx, "NewPHI"));