[IR] (NFC) [BasicBlock] updatePHIEdges, duplicatePHIEdges and fixPHIEdges

[Marandil]

Summary

[IR] Add three public functions to BasicBlock to handle bulk changes to PHINodes caused by changing control flow graph.
updatePHIEdges changes all usages of one BB into another BB.
duplicatePHIEdges duplicates all incoming edges from one BB into incoming edges from another BB.
fixPHIEdges checks if a given BB is still a valid predecessor and based on this uses either of update or duplicate.

This diff is a follow up after https://reviews.llvm.org/D33787. I am submitting a new diff and closing the old one, since the inline comments in the latter would be unintelligible after the API changes.

include/llvm/IR/BasicBlock.h

                     static_cast<const BasicBlock *>(this)->getLandingPadInst());
   }
 
+  /// \brief Update all PHI nodes by changing IncomingOld to IncomingNew
+  void updatePHIEdges(BasicBlock* IncomingOld,
+                      BasicBlock* IncomingNew);
+
+  /// \brief Update all PHI nodes by duplicating incoming edges from IncomingOld
+  /// into incoming edges from IncomingNew
+  void duplicatePHIEdges(BasicBlock* IncomingOld,
+                         BasicBlock* IncomingNew);
+
+  /// \brief Update or duplicate PHI Edges to include IncomingNew depending on
+  /// whether or not IncomingOld is still a valid predecessor.
+  void fixPHIEdges(BasicBlock* IncomingOld,
+                   BasicBlock* IncomingNew);
+
 private:
   /// \brief Increment the internal refcount of the number of BlockAddresses
   /// referencing this BasicBlock by \p Amt.

lib/IR/BasicBlock.cpp

   // know that incoming branches will be from New, not from Old.
   //
   for (succ_iterator I = succ_begin(New), E = succ_end(New); I != E; ++I) {
-    // Loop over any phi nodes in the basic block, updating the BB field of
-    // incoming values...
-    BasicBlock *Successor = *I;
-    for (auto &PN : Successor->phis()) {
-      int Idx = PN.getBasicBlockIndex(this);
-      while (Idx != -1) {
-        PN.setIncomingBlock((unsigned)Idx, New);
-        Idx = PN.getBasicBlockIndex(this);
-      }
-    }
+    I->updatePHIEdges(this, New);
   }
   return New;
 }

 const LandingPadInst *BasicBlock::getLandingPadInst() const {
   return dyn_cast<LandingPadInst>(getFirstNonPHI());
 }
+
+
+/// Update all PHI nodes by changing IncomingOld to IncomingNew
+void BasicBlock::updatePHIEdges(BasicBlock* IncomingOld,
+                                BasicBlock* IncomingNew) {
+  for (auto &PN : this->phis()) {
+    // Coping with possibly unfinished PHI Nodes
+    int Idx;
+    while ((Idx = PN.getBasicBlockIndex(IncomingOld)) >= 0)
+      PN.setIncomingBlock((unsigned)Idx, IncomingNew);
+  }
+}
+
+/// Update all PHI nodes by duplicating incoming edges from IncomingOld
+/// into incoming edges from IncomingNew
+void BasicBlock::duplicatePHIEdges(BasicBlock* IncomingOld,
+                                   BasicBlock* IncomingNew) {
+  for (auto &PN : this->phis()) {
+    int Idx = PN.getBasicBlockIndex(IncomingOld);
+    // Old predecessor shall exist in PHI.
+    assert (Idx >= 0 && "Broken PHINode!");
+    PN.addIncoming(PN.getIncomingValue((unsigned)Idx), IncomingNew);
+  }
+}
+
+/// Find all PHI nodes and depending on whether OldIncoming is a valid
+/// predecessor, add NewIncoming with the same value as OldIncoming, or change
+/// the incoming links on them.
+void BasicBlock::fixPHIEdges(BasicBlock* IncomingOld,
+                             BasicBlock* IncomingNew) {
+  // Iterate over predecessors of BB; mark if any of them is IncomingOld or
+  // IncomingNew.
+  bool oldPredecessor = false;
+  bool newPredecessor = false;
+  for (BasicBlock *Pred : predecessors(this)) {
+    if (Pred == IncomingOld)
+      oldPredecessor = true;
+    if (Pred == IncomingNew)
+      newPredecessor = true;
+  }
+
+  assert (newPredecessor && "Attempting to update PHI predecessor for block "
+                            "with no such predecessor!");
+
+  if (oldPredecessor) {
+    // If IncomingOld is still valid, need to duplicate the edge.
+    duplicatePHIEdges(IncomingOld, IncomingNew);
+  } else {
+    updatePHIEdges(IncomingOld, IncomingNew);
+  }
+}

unittests/IR/BasicBlockTest.cpp

 
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/NoFolder.h"
+#include "llvm/Support/SourceMgr.h"
 #include "gmock/gmock-matchers.h"
 #include "gtest/gtest.h"
 #include <memory>

   }
 }
 
+class BasicBlockPhiTest : public ::testing::Test
+{
+protected:
+  LLVMContext C;
+
+  // Create a somewhat complex function with a couple of basic blocks and a
+  // single phi node.
+  const char *ModuleString = "define i32 @f(i32 %x) {\n"
+                             "bb0:\n"
+                             "  switch i32 %x, label %bb1 [\n"
+                             "         i32 2,  label %bb2\n"
+                             "         i32 3,  label %bb3\n"
+                             "         i32 4,  label %bb4]\n"
+                             "bb1:\n"
+                             "  br label %bb5\n"
+                             "bb2:\n"
+                             "  br label %bb5\n"
+                             "bb3:\n"
+                             "  br label %bb5\n"
+                             "bb4:\n"
+                             "  br label %bb5\n"
+                             "bb5:\n"
+                             "  %r = phi i32 [ 0, %bb1 ],\n"
+                             "               [ 1, %bb2 ],\n"
+                             "               [ 2, %bb3 ],\n"
+                             "               [ 3, %bb4 ]\n"
+                             "  ret i32 %r\n"
+                             "}\n";
+
+  std::unique_ptr<Module> M;
+  Function *F;
+
+  BasicBlock* BB0;
+  BasicBlock* BB1;
+  BasicBlock* BB2;
+  BasicBlock* BB3;
+  BasicBlock* BB4;
+  BasicBlock* BB5;
+
+  virtual void SetUp()
+  {
+    SMDiagnostic Err;
+
+    M = parseAssemblyString(ModuleString, Err, C);
+    F = M->getFunction("f");
+
+    auto It = F->begin();
+
+    BB0 = &*It++;
+    BB1 = &*It++;
+    BB2 = &*It++;
+    BB3 = &*It++;
+    BB4 = &*It++;
+    BB5 = &*It++;
+
+    EXPECT_TRUE(It == F->end());
+  }
+};
+
+TEST_F(BasicBlockPhiTest, SwitchReplaceFixPHI) {
+
+  SwitchInst* SI = dyn_cast<SwitchInst>(&BB0->front());
+  PHINode*    PN = dyn_cast<PHINode>(&BB5->front());
+
+  EXPECT_TRUE(SI != nullptr);
+  EXPECT_TRUE(PN != nullptr);
+
+  // Insert new label bb6.
+  IRBuilder<> Builder(C);
+  BasicBlock* BB6 = BasicBlock::Create(C, "bb6", F);
+
+  // Make bb1 jump to bb5.
+  BranchInst* BB1_Br = dyn_cast<BranchInst>(&BB1->front());
+  BB1_Br->setSuccessor(0, BB6);
+
+  // Add immediate branch to bb1.
+  Builder.SetInsertPoint(BB6);
+  Builder.CreateBr(BB5);
+
+  BB5->fixPHIEdges(BB1, BB6);
+
+  outs() << *PN << "\n";
+  EXPECT_TRUE(PN->getIncomingBlock(0) == BB6);
+}
+
+TEST_F(BasicBlockPhiTest, SwitchDefaultFixPHI) {
+
+  SwitchInst* SI = dyn_cast<SwitchInst>(&BB0->front());
+  PHINode*    PN = dyn_cast<PHINode>(&BB5->front());
+
+  EXPECT_TRUE(SI != nullptr);
+  EXPECT_TRUE(PN != nullptr);
+
+  // Make the switch directly branch to BB5 on default.
+  SI->setDefaultDest(BB5);
+  BB1->eraseFromParent();
+
+  BB5->fixPHIEdges(BB1, BB0);
+
+  outs() << *PN << "\n";
+  EXPECT_TRUE(PN->getIncomingBlock(0) == BB0);
+}
+
 } // End anonymous namespace.
 } // End llvm namespace.