Skip to content

Commit cfb92be

Browse files
vporpovporpo
authored
[SandboxIR] Implement CallBrInst (#100823)
This patch implements sandboxir::CallBrInst which mirrors llvm::CallBrInst. LLVM IR does not expose the Uses to DefaultDest and IndirectDest so we need special Tracker objects for both of setters.
1 parent 62bd08a commit cfb92be

File tree

7 files changed

+406
-5
lines changed

7 files changed

+406
-5
lines changed

‎llvm/include/llvm/SandboxIR/SandboxIR.h

Lines changed: 74 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,11 @@
3030
// |
3131
// +- CastInst
3232
// |
33-
// +- CallBase ------+- CallInst
34-
// | |
35-
// +- CmpInst +- InvokeInst
36-
// |
37-
// +- ExtractElementInst
33+
// +- CallBase -----------+- CallBrInst
34+
// | |
35+
// +- CmpInst +- CallInst
36+
// | |
37+
// +- ExtractElementInst +- InvokeInst
3838
// |
3939
// +- GetElementPtrInst
4040
// |
@@ -92,6 +92,7 @@ class Value;
9292
class CallBase;
9393
class CallInst;
9494
class InvokeInst;
95+
class CallBrInst;
9596

9697
/// Iterator for the `Use` edges of a User's operands.
9798
/// \Returns the operand `Use` when dereferenced.
@@ -206,6 +207,7 @@ class Value {
206207
friend class CallBase; // For getting `Val`.
207208
friend class CallInst; // For getting `Val`.
208209
friend class InvokeInst; // For getting `Val`.
210+
friend class CallBrInst; // For getting `Val`.
209211

210212
/// All values point to the context.
211213
Context &Ctx;
@@ -545,6 +547,7 @@ class Instruction : public sandboxir::User {
545547
friend class ReturnInst; // For getTopmostLLVMInstruction().
546548
friend class CallInst; // For getTopmostLLVMInstruction().
547549
friend class InvokeInst; // For getTopmostLLVMInstruction().
550+
friend class CallBrInst; // For getTopmostLLVMInstruction().
548551

549552
/// \Returns the LLVM IR Instructions that this SandboxIR maps to in program
550553
/// order.
@@ -877,6 +880,7 @@ class CallBase : public Instruction {
877880
: Instruction(ID, Opc, I, Ctx) {}
878881
friend class CallInst; // For constructor.
879882
friend class InvokeInst; // For constructor.
883+
friend class CallBrInst; // For constructor.
880884

881885
public:
882886
static bool classof(const Value *From) {
@@ -1108,6 +1112,69 @@ class InvokeInst final : public CallBase {
11081112
#endif
11091113
};
11101114

1115+
class CallBrInst final : public CallBase {
1116+
/// Use Context::createCallBrInst(). Don't call the
1117+
/// constructor directly.
1118+
CallBrInst(llvm::Instruction *I, Context &Ctx)
1119+
: CallBase(ClassID::CallBr, Opcode::CallBr, I, Ctx) {}
1120+
friend class Context; // For accessing the constructor in
1121+
// create*()
1122+
Use getOperandUseInternal(unsigned OpIdx, bool Verify) const final {
1123+
return getOperandUseDefault(OpIdx, Verify);
1124+
}
1125+
SmallVector<llvm::Instruction *, 1> getLLVMInstrs() const final {
1126+
return {cast<llvm::Instruction>(Val)};
1127+
}
1128+
1129+
public:
1130+
static CallBrInst *create(FunctionType *FTy, Value *Func,
1131+
BasicBlock *DefaultDest,
1132+
ArrayRef<BasicBlock *> IndirectDests,
1133+
ArrayRef<Value *> Args, BBIterator WhereIt,
1134+
BasicBlock *WhereBB, Context &Ctx,
1135+
const Twine &NameStr = "");
1136+
static CallBrInst *create(FunctionType *FTy, Value *Func,
1137+
BasicBlock *DefaultDest,
1138+
ArrayRef<BasicBlock *> IndirectDests,
1139+
ArrayRef<Value *> Args, Instruction *InsertBefore,
1140+
Context &Ctx, const Twine &NameStr = "");
1141+
static CallBrInst *create(FunctionType *FTy, Value *Func,
1142+
BasicBlock *DefaultDest,
1143+
ArrayRef<BasicBlock *> IndirectDests,
1144+
ArrayRef<Value *> Args, BasicBlock *InsertAtEnd,
1145+
Context &Ctx, const Twine &NameStr = "");
1146+
static bool classof(const Value *From) {
1147+
return From->getSubclassID() == ClassID::CallBr;
1148+
}
1149+
unsigned getUseOperandNo(const Use &Use) const final {
1150+
return getUseOperandNoDefault(Use);
1151+
}
1152+
unsigned getNumOfIRInstrs() const final { return 1u; }
1153+
unsigned getNumIndirectDests() const {
1154+
return cast<llvm::CallBrInst>(Val)->getNumIndirectDests();
1155+
}
1156+
Value *getIndirectDestLabel(unsigned Idx) const;
1157+
Value *getIndirectDestLabelUse(unsigned Idx) const;
1158+
BasicBlock *getDefaultDest() const;
1159+
BasicBlock *getIndirectDest(unsigned Idx) const;
1160+
SmallVector<BasicBlock *, 16> getIndirectDests() const;
1161+
void setDefaultDest(BasicBlock *BB);
1162+
void setIndirectDest(unsigned Idx, BasicBlock *BB);
1163+
BasicBlock *getSuccessor(unsigned Idx) const;
1164+
unsigned getNumSuccessors() const {
1165+
return cast<llvm::CallBrInst>(Val)->getNumSuccessors();
1166+
}
1167+
#ifndef NDEBUG
1168+
void verify() const final {}
1169+
friend raw_ostream &operator<<(raw_ostream &OS, const CallBrInst &I) {
1170+
I.dump(OS);
1171+
return OS;
1172+
}
1173+
void dump(raw_ostream &OS) const override;
1174+
LLVM_DUMP_METHOD void dump() const override;
1175+
#endif
1176+
};
1177+
11111178
/// An LLLVM Instruction that has no SandboxIR equivalent class gets mapped to
11121179
/// an OpaqueInstr.
11131180
class OpaqueInst : public sandboxir::Instruction {
@@ -1260,6 +1327,8 @@ class Context {
12601327
friend CallInst; // For createCallInst()
12611328
InvokeInst *createInvokeInst(llvm::InvokeInst *I);
12621329
friend InvokeInst; // For createInvokeInst()
1330+
CallBrInst *createCallBrInst(llvm::CallBrInst *I);
1331+
friend CallBrInst; // For createCallBrInst()
12631332

12641333
public:
12651334
Context(LLVMContext &LLVMCtx)

‎llvm/include/llvm/SandboxIR/Tracker.h

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
namespace llvm::sandboxir {
5454

5555
class BasicBlock;
56+
class CallBrInst;
5657
class Instruction;
5758
class Tracker;
5859

@@ -177,6 +178,41 @@ class RemoveFromParent : public IRChangeBase {
177178
#endif // NDEBUG
178179
};
179180

181+
class CallBrInstSetDefaultDest : public IRChangeBase {
182+
CallBrInst *CallBr;
183+
BasicBlock *OrigDefaultDest;
184+
185+
public:
186+
CallBrInstSetDefaultDest(CallBrInst *CallBr, Tracker &Tracker);
187+
void revert() final;
188+
void accept() final {}
189+
#ifndef NDEBUG
190+
void dump(raw_ostream &OS) const final {
191+
dumpCommon(OS);
192+
OS << "CallBrInstSetDefaultDest";
193+
}
194+
LLVM_DUMP_METHOD void dump() const final;
195+
#endif
196+
};
197+
198+
class CallBrInstSetIndirectDest : public IRChangeBase {
199+
CallBrInst *CallBr;
200+
unsigned Idx;
201+
BasicBlock *OrigIndirectDest;
202+
203+
public:
204+
CallBrInstSetIndirectDest(CallBrInst *CallBr, unsigned Idx, Tracker &Tracker);
205+
void revert() final;
206+
void accept() final {}
207+
#ifndef NDEBUG
208+
void dump(raw_ostream &OS) const final {
209+
dumpCommon(OS);
210+
OS << "CallBrInstSetIndirectDest";
211+
}
212+
LLVM_DUMP_METHOD void dump() const final;
213+
#endif
214+
};
215+
180216
class MoveInstr : public IRChangeBase {
181217
/// The instruction that moved.
182218
Instruction *MovedI;

‎llvm/include/llvm/SandboxIR/Use.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ class Use {
4242
friend class OperandUseIterator; // For constructor
4343
friend class UserUseIterator; // For accessing members
4444
friend class CallBase; // For LLVMUse
45+
friend class CallBrInst; // For constructor
4546

4647
public:
4748
operator Value *() const { return get(); }

‎llvm/lib/SandboxIR/SandboxIR.cpp

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
#include "llvm/SandboxIR/SandboxIR.h"
1010
#include "llvm/ADT/SmallPtrSet.h"
11+
#include "llvm/ADT/SmallVector.h"
1112
#include "llvm/IR/Constants.h"
1213
#include "llvm/Support/Debug.h"
1314
#include <sstream>
@@ -896,6 +897,105 @@ void InvokeInst::dump() const {
896897
dump(dbgs());
897898
dbgs() << "\n";
898899
}
900+
#endif // NDEBUG
901+
902+
CallBrInst *CallBrInst::create(FunctionType *FTy, Value *Func,
903+
BasicBlock *DefaultDest,
904+
ArrayRef<BasicBlock *> IndirectDests,
905+
ArrayRef<Value *> Args, BBIterator WhereIt,
906+
BasicBlock *WhereBB, Context &Ctx,
907+
const Twine &NameStr) {
908+
auto &Builder = Ctx.getLLVMIRBuilder();
909+
if (WhereIt != WhereBB->end())
910+
Builder.SetInsertPoint((*WhereIt).getTopmostLLVMInstruction());
911+
else
912+
Builder.SetInsertPoint(cast<llvm::BasicBlock>(WhereBB->Val));
913+
914+
SmallVector<llvm::BasicBlock *> LLVMIndirectDests;
915+
LLVMIndirectDests.reserve(IndirectDests.size());
916+
for (BasicBlock *IndDest : IndirectDests)
917+
LLVMIndirectDests.push_back(cast<llvm::BasicBlock>(IndDest->Val));
918+
919+
SmallVector<llvm::Value *> LLVMArgs;
920+
LLVMArgs.reserve(Args.size());
921+
for (Value *Arg : Args)
922+
LLVMArgs.push_back(Arg->Val);
923+
924+
llvm::CallBrInst *CallBr = Builder.CreateCallBr(
925+
FTy, Func->Val, cast<llvm::BasicBlock>(DefaultDest->Val),
926+
LLVMIndirectDests, LLVMArgs, NameStr);
927+
return Ctx.createCallBrInst(CallBr);
928+
}
929+
930+
CallBrInst *CallBrInst::create(FunctionType *FTy, Value *Func,
931+
BasicBlock *DefaultDest,
932+
ArrayRef<BasicBlock *> IndirectDests,
933+
ArrayRef<Value *> Args,
934+
Instruction *InsertBefore, Context &Ctx,
935+
const Twine &NameStr) {
936+
return create(FTy, Func, DefaultDest, IndirectDests, Args,
937+
InsertBefore->getIterator(), InsertBefore->getParent(), Ctx,
938+
NameStr);
939+
}
940+
CallBrInst *CallBrInst::create(FunctionType *FTy, Value *Func,
941+
BasicBlock *DefaultDest,
942+
ArrayRef<BasicBlock *> IndirectDests,
943+
ArrayRef<Value *> Args, BasicBlock *InsertAtEnd,
944+
Context &Ctx, const Twine &NameStr) {
945+
return create(FTy, Func, DefaultDest, IndirectDests, Args, InsertAtEnd->end(),
946+
InsertAtEnd, Ctx, NameStr);
947+
}
948+
949+
Value *CallBrInst::getIndirectDestLabel(unsigned Idx) const {
950+
return Ctx.getValue(cast<llvm::CallBrInst>(Val)->getIndirectDestLabel(Idx));
951+
}
952+
Value *CallBrInst::getIndirectDestLabelUse(unsigned Idx) const {
953+
return Ctx.getValue(
954+
cast<llvm::CallBrInst>(Val)->getIndirectDestLabelUse(Idx));
955+
}
956+
BasicBlock *CallBrInst::getDefaultDest() const {
957+
return cast<BasicBlock>(
958+
Ctx.getValue(cast<llvm::CallBrInst>(Val)->getDefaultDest()));
959+
}
960+
BasicBlock *CallBrInst::getIndirectDest(unsigned Idx) const {
961+
return cast<BasicBlock>(
962+
Ctx.getValue(cast<llvm::CallBrInst>(Val)->getIndirectDest(Idx)));
963+
}
964+
llvm::SmallVector<BasicBlock *, 16> CallBrInst::getIndirectDests() const {
965+
SmallVector<BasicBlock *, 16> BBs;
966+
for (llvm::BasicBlock *LLVMBB :
967+
cast<llvm::CallBrInst>(Val)->getIndirectDests())
968+
BBs.push_back(cast<BasicBlock>(Ctx.getValue(LLVMBB)));
969+
return BBs;
970+
}
971+
void CallBrInst::setDefaultDest(BasicBlock *BB) {
972+
auto &Tracker = Ctx.getTracker();
973+
if (Tracker.isTracking())
974+
Tracker.track(std::make_unique<CallBrInstSetDefaultDest>(this, Tracker));
975+
cast<llvm::CallBrInst>(Val)->setDefaultDest(cast<llvm::BasicBlock>(BB->Val));
976+
}
977+
void CallBrInst::setIndirectDest(unsigned Idx, BasicBlock *BB) {
978+
auto &Tracker = Ctx.getTracker();
979+
if (Tracker.isTracking())
980+
Tracker.track(
981+
std::make_unique<CallBrInstSetIndirectDest>(this, Idx, Tracker));
982+
cast<llvm::CallBrInst>(Val)->setIndirectDest(Idx,
983+
cast<llvm::BasicBlock>(BB->Val));
984+
}
985+
BasicBlock *CallBrInst::getSuccessor(unsigned Idx) const {
986+
return cast<BasicBlock>(
987+
Ctx.getValue(cast<llvm::CallBrInst>(Val)->getSuccessor(Idx)));
988+
}
989+
990+
#ifndef NDEBUG
991+
void CallBrInst::dump(raw_ostream &OS) const {
992+
dumpCommonPrefix(OS);
993+
dumpCommonSuffix(OS);
994+
}
995+
void CallBrInst::dump() const {
996+
dump(dbgs());
997+
dbgs() << "\n";
998+
}
899999

9001000
void OpaqueInst::dump(raw_ostream &OS) const {
9011001
dumpCommonPrefix(OS);
@@ -1060,6 +1160,11 @@ Value *Context::getOrCreateValueInternal(llvm::Value *LLVMV, llvm::User *U) {
10601160
It->second = std::unique_ptr<InvokeInst>(new InvokeInst(LLVMInvoke, *this));
10611161
return It->second.get();
10621162
}
1163+
case llvm::Instruction::CallBr: {
1164+
auto *LLVMCallBr = cast<llvm::CallBrInst>(LLVMV);
1165+
It->second = std::unique_ptr<CallBrInst>(new CallBrInst(LLVMCallBr, *this));
1166+
return It->second.get();
1167+
}
10631168
default:
10641169
break;
10651170
}
@@ -1113,6 +1218,11 @@ InvokeInst *Context::createInvokeInst(llvm::InvokeInst *I) {
11131218
return cast<InvokeInst>(registerValue(std::move(NewPtr)));
11141219
}
11151220

1221+
CallBrInst *Context::createCallBrInst(llvm::CallBrInst *I) {
1222+
auto NewPtr = std::unique_ptr<CallBrInst>(new CallBrInst(I, *this));
1223+
return cast<CallBrInst>(registerValue(std::move(NewPtr)));
1224+
}
1225+
11161226
Value *Context::getValue(llvm::Value *V) const {
11171227
auto It = LLVMValueToValueMap.find(V);
11181228
if (It != LLVMValueToValueMap.end())

‎llvm/lib/SandboxIR/Tracker.cpp

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,37 @@ void RemoveFromParent::dump() const {
129129
}
130130
#endif
131131

132+
CallBrInstSetDefaultDest::CallBrInstSetDefaultDest(CallBrInst *CallBr,
133+
Tracker &Tracker)
134+
: IRChangeBase(Tracker), CallBr(CallBr) {
135+
OrigDefaultDest = CallBr->getDefaultDest();
136+
}
137+
void CallBrInstSetDefaultDest::revert() {
138+
CallBr->setDefaultDest(OrigDefaultDest);
139+
}
140+
#ifndef NDEBUG
141+
void CallBrInstSetDefaultDest::dump() const {
142+
dump(dbgs());
143+
dbgs() << "\n";
144+
}
145+
#endif
146+
147+
CallBrInstSetIndirectDest::CallBrInstSetIndirectDest(CallBrInst *CallBr,
148+
unsigned Idx,
149+
Tracker &Tracker)
150+
: IRChangeBase(Tracker), CallBr(CallBr), Idx(Idx) {
151+
OrigIndirectDest = CallBr->getIndirectDest(Idx);
152+
}
153+
void CallBrInstSetIndirectDest::revert() {
154+
CallBr->setIndirectDest(Idx, OrigIndirectDest);
155+
}
156+
#ifndef NDEBUG
157+
void CallBrInstSetIndirectDest::dump() const {
158+
dump(dbgs());
159+
dbgs() << "\n";
160+
}
161+
#endif
162+
132163
MoveInstr::MoveInstr(Instruction *MovedI, Tracker &Tracker)
133164
: IRChangeBase(Tracker), MovedI(MovedI) {
134165
if (auto *NextI = MovedI->getNextNode())

0 commit comments

Comments
 (0)