Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Nov 12, 2024
Merged

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

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
b266ea7
[SandboxIR] Preserve the order of switch cases after revert.
slackito Nov 9, 2024
6bccc01
Added a separate "switch preserves successor order" test.
slackito Nov 11, 2024
120eb87
clang-format
slackito Nov 11, 2024
17ea6d8
Avoid depending on llvm::SwitchInst::removeCase impl details.
slackito Nov 12, 2024
8a980b2
clang-format again
slackito Nov 12, 2024
79bdeb8
address more review feedback
slackito Nov 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions llvm/include/llvm/SandboxIR/Tracker.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -315,12 +315,14 @@ class SwitchAddCase : public IRChangeBase {

class SwitchRemoveCase : public IRChangeBase {
SwitchInst *Switch;
unsigned Index;
ConstantInt *Val;
BasicBlock *Dest;

public:
SwitchRemoveCase(SwitchInst *Switch, ConstantInt *Val, BasicBlock *Dest)
: Switch(Switch), Val(Val), Dest(Dest) {}
SwitchRemoveCase(SwitchInst *Switch, unsigned Index, ConstantInt *Val,
BasicBlock *Dest)
: Switch(Switch), Index(Index), Val(Val), Dest(Dest) {}
void revert(Tracker &Tracker) final;
void accept() final {}
#ifndef NDEBUG
Expand Down
2 changes: 1 addition & 1 deletion llvm/lib/SandboxIR/Instruction.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1133,7 +1133,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());
this, Case.getCaseIndex(), Case.getCaseValue(), Case.getCaseSuccessor());

auto *LLVMSwitch = cast<llvm::SwitchInst>(Val);
unsigned CaseNum = It - case_begin();
Expand Down
14 changes: 13 additions & 1 deletion llvm/lib/SandboxIR/Tracker.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,19 @@ void CatchSwitchAddHandler::revert(Tracker &Tracker) {
LLVMCSI->removeHandler(LLVMCSI->handler_begin() + HandlerIdx);
}

void SwitchRemoveCase::revert(Tracker &Tracker) { Switch->addCase(Val, Dest); }
void SwitchRemoveCase::revert(Tracker &Tracker) {
// removeCase swaps the last case with the deleted one. To revert it, we use
// addCase (which adds the new case to the end), and swap the newly-added
// value and successor operands to the positions for the original case index.
Switch->addCase(Val, Dest);
auto ValUseA = Switch->getOperandUse(2 + Index * 2);
auto SucUseA = Switch->getOperandUse(2 + Index * 2 + 1);
unsigned NumOps = Switch->getNumOperands();
auto ValUseB = Switch->getOperandUse(NumOps - 2);
auto SucUseB = Switch->getOperandUse(NumOps - 2 + 1);
ValUseA.swap(ValUseB);
Copy link
Contributor

@vporpo vporpo Nov 9, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this would require a loop that swaps the last 'case' one position to the left at a time until we reach Index. Because if the switch contains 3 cases (Case0, Case1, Case2) and we remove case 0, then when we revert we are placing the removed one at the end (Case1, Case2, Case0). Then we would need to bring Case0 back to position 0 without changing the order of 1 and 2.

Also I think we could use swapOperands() instead of swap(Use), which may be a bit more compact.

Copy link
Collaborator Author

@slackito slackito Nov 11, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If I understand the implementation correctly, if you remove Case0 from (Case0, Case1, Case2), it swaps the element to be deleted with the last one, and then shrinks the list. So you actually get (Case2, Case1), which you can undo with a single swap after adding the case back at the end of the list.

Copy link
Contributor

@vporpo vporpo Nov 11, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh I hadn't realized that this is how it works, makes sense though. Hmm but should we rely on this behavior ? it seems too implementation specific, so if someone changes the implementation in LLVM IR, our tests will break. I think a comment that explains how this works would be really helpful.

The alternative would be to record the order of all cases when we create the change object. Then when we want to revert we remove them all and add them in again. Wdyt?

Copy link
Collaborator Author

@slackito slackito Nov 11, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think a comment that explains how this works would be really helpful.

That's what this comment a couple of lines above was supposed to be about:

  // removeCase swaps the last case with the deleted one. To revert it, we use
  // addCase (which adds the new case to the end), and swap the newly-added
  // value and successor operands to the positions for the original case index.

I can try to rephrase/expand it if it wasn't clear enough.

Hmm but should we rely on this behavior ? it seems too implementation specific, so if someone changes the implementation in LLVM IR, our tests will break.

Yeah, I was assuming the implementation wasn't very likely to change and that we would be fine updating this logic if it ever happens. If we don't feel comfortable with that option I can go with the alternative you suggested (recording all cases, then adding them again on revert) instead. Which one do you prefer?

Copy link
Contributor

@vporpo vporpo Nov 11, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can try to rephrase/expand it if it wasn't clear enough.

Sorry, what I wanted to say is that it would be nice to show the example with the ordering of cases by actually showing them too like with (case0, case1, case2).

Which one do you prefer?

I am fine with either, but if we stick with this one, we should probably be very specific about the fact that this is implementation dependent and that it could break.

SucUseA.swap(SucUseB);
}

#ifndef NDEBUG
void SwitchRemoveCase::dump() const {
Expand Down
82 changes: 82 additions & 0 deletions llvm/unittests/SandboxIR/TrackerTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Copy link
Contributor

@vporpo vporpo Nov 11, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator Author

@slackito slackito Nov 12, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changed the implementation so it doesn't rely on the specific ordering changes introduced by removeCase()

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);
Copy link
Contributor

@vporpo vporpo Nov 9, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add a test where we remove case 0 of a 3-case switch, to check that we revert correctly in that case too.

Copy link
Collaborator Author

@slackito slackito Nov 11, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

EXPECT_EQ(Switch->getSuccessor(3), BB2);
}

TEST_F(TrackerTest, SelectInst) {
parseIR(C, R"IR(
define void @foo(i1 %c0, i8 %v0, i8 %v1) {
Expand Down
Loading