[llvm] [SandboxIR] Preserve the order of switch cases after revert. (PR #115577)

[via llvm-commits]

================
@@ -965,6 +965,88 @@ define void @foo(i32 %cond0, i32 %cond1) {
   EXPECT_EQ(Switch->findCaseDest(BB1), One);
 }
 
+TEST_F(TrackerTest, SwitchInstPreservesSuccesorOrder) {
+  parseIR(C, R"IR(
+define void @foo(i32 %cond0) {
+  entry:
+    switch i32 %cond0, label %default [ i32 0, label %bb0
+                                        i32 1, label %bb1
+                                        i32 2, label %bb2 ]
+  bb0:
+    ret void
+  bb1:
+    ret void
+  bb2:
+    ret void
+  default:
+    ret void
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  auto *LLVMEntry = getBasicBlockByName(LLVMF, "entry");
+
+  sandboxir::Context Ctx(C);
+  [[maybe_unused]] auto &F = *Ctx.createFunction(&LLVMF);
+  auto *Entry = cast<sandboxir::BasicBlock>(Ctx.getValue(LLVMEntry));
+  auto *BB0 = cast<sandboxir::BasicBlock>(
+      Ctx.getValue(getBasicBlockByName(LLVMF, "bb0")));
+  auto *BB1 = cast<sandboxir::BasicBlock>(
+      Ctx.getValue(getBasicBlockByName(LLVMF, "bb1")));
+  auto *BB2 = cast<sandboxir::BasicBlock>(
+      Ctx.getValue(getBasicBlockByName(LLVMF, "bb2")));
+  auto *Switch = cast<sandboxir::SwitchInst>(&*Entry->begin());
+
+  auto *DefaultDest = Switch->getDefaultDest();
+  auto *Zero = sandboxir::ConstantInt::get(sandboxir::Type::getInt32Ty(Ctx), 0);
+  auto *One = sandboxir::ConstantInt::get(sandboxir::Type::getInt32Ty(Ctx), 1);
+  auto *Two = sandboxir::ConstantInt::get(sandboxir::Type::getInt32Ty(Ctx), 2);
+
+  // Check that we can properly revert a removeCase multiple positions apart
+  // from the end of the operand list.
+  Ctx.save();
+  Switch->removeCase(Switch->findCaseValue(Zero));
+  EXPECT_EQ(Switch->getNumCases(), 2u);
+  Ctx.revert();
+  EXPECT_EQ(Switch->getNumCases(), 3u);
+  EXPECT_EQ(Switch->findCaseDest(BB0), Zero);
+  EXPECT_EQ(Switch->findCaseDest(BB1), One);
+  EXPECT_EQ(Switch->findCaseDest(BB2), Two);
+  EXPECT_EQ(Switch->getSuccessor(0), DefaultDest);
+  EXPECT_EQ(Switch->getSuccessor(1), BB0);
+  EXPECT_EQ(Switch->getSuccessor(2), BB1);
+  EXPECT_EQ(Switch->getSuccessor(3), BB2);
+
+  // Check that we can properly revert a removeCase of the last case.
+  Ctx.save();
+  Switch->removeCase(Switch->findCaseValue(Two));
+  EXPECT_EQ(Switch->getNumCases(), 2u);
+  Ctx.revert();
+  EXPECT_EQ(Switch->getNumCases(), 3u);
+  EXPECT_EQ(Switch->findCaseDest(BB0), Zero);
+  EXPECT_EQ(Switch->findCaseDest(BB1), One);
+  EXPECT_EQ(Switch->findCaseDest(BB2), Two);
+  EXPECT_EQ(Switch->getSuccessor(0), DefaultDest);
+  EXPECT_EQ(Switch->getSuccessor(1), BB0);
+  EXPECT_EQ(Switch->getSuccessor(2), BB1);
+  EXPECT_EQ(Switch->getSuccessor(3), BB2);
+
+  // Check order is preserved after reverting multiple removeCase invocations.
+  Ctx.save();
+  Switch->removeCase(Switch->findCaseValue(One));
+  Switch->removeCase(Switch->findCaseValue(Zero));
+  Switch->removeCase(Switch->findCaseValue(Two));
+  EXPECT_EQ(Switch->getNumCases(), 0u);
+  Ctx.revert();
+  EXPECT_EQ(Switch->getNumCases(), 3u);
+  EXPECT_EQ(Switch->findCaseDest(BB0), Zero);
----------------
vporpo wrote:

Perhaps this would be a good place to add a comment that if these checks fail, then it might be that the implementation in LLVM changed and that we should update the way revert() works.

https://github.com/llvm/llvm-project/pull/115577