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

Preserving the case order is not strictly necessary to preserve
semantics (for example, operations like SwitchInst::removeCase will
happily swap cases around). However, I'm planning to introduce an
optional verification step for SandboxIR that will use StructuralHash to
compare IR after a revert to the original IR to help catch tracker bugs,
and the order difference triggers a difference there.
This commit is contained in:
Jorge Gorbe Moya 2024-11-12 14:10:46 -08:00 committed by GitHub
parent a2042521a0
commit 9d85ba5724
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 108 additions and 8 deletions

View File

@ -315,12 +315,15 @@ public:
class SwitchRemoveCase : public IRChangeBase {
SwitchInst *Switch;
ConstantInt *Val;
BasicBlock *Dest;
struct Case {
ConstantInt *Val;
BasicBlock *Dest;
};
SmallVector<Case> Cases;
public:
SwitchRemoveCase(SwitchInst *Switch, ConstantInt *Val, BasicBlock *Dest)
: Switch(Switch), Val(Val), Dest(Dest) {}
SwitchRemoveCase(SwitchInst *Switch);
void revert(Tracker &Tracker) final;
void accept() final {}
#ifndef NDEBUG

View File

@ -1131,9 +1131,7 @@ void SwitchInst::addCase(ConstantInt *OnVal, BasicBlock *Dest) {
}
SwitchInst::CaseIt SwitchInst::removeCase(CaseIt It) {
auto &Case = *It;
Ctx.getTracker().emplaceIfTracking<SwitchRemoveCase>(
this, Case.getCaseValue(), Case.getCaseSuccessor());
Ctx.getTracker().emplaceIfTracking<SwitchRemoveCase>(this);
auto *LLVMSwitch = cast<llvm::SwitchInst>(Val);
unsigned CaseNum = It - case_begin();

View File

@ -170,7 +170,24 @@ void CatchSwitchAddHandler::revert(Tracker &Tracker) {
LLVMCSI->removeHandler(LLVMCSI->handler_begin() + HandlerIdx);
}
void SwitchRemoveCase::revert(Tracker &Tracker) { Switch->addCase(Val, Dest); }
SwitchRemoveCase::SwitchRemoveCase(SwitchInst *Switch) : Switch(Switch) {
for (const auto &C : Switch->cases())
Cases.push_back({C.getCaseValue(), C.getCaseSuccessor()});
}
void SwitchRemoveCase::revert(Tracker &Tracker) {
// SwitchInst::removeCase doesn't provide any guarantees about the order of
// cases after removal. In order to preserve the original ordering, we save
// all of them and, when reverting, clear them all then insert them in the
// desired order. This still relies on the fact that `addCase` will insert
// them at the end, but it is documented to invalidate `case_end()` so it's
// probably okay.
unsigned NumCases = Switch->getNumCases();
for (unsigned I = 0; I < NumCases; ++I)
Switch->removeCase(Switch->case_begin());
for (auto &Case : Cases)
Switch->addCase(Case.Val, Case.Dest);
}
#ifndef NDEBUG
void SwitchRemoveCase::dump() const {

View File

@ -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);
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);
}
TEST_F(TrackerTest, SelectInst) {
parseIR(C, R"IR(
define void @foo(i1 %c0, i8 %v0, i8 %v1) {