xref: /freebsd/contrib/llvm-project/llvm/include/llvm/SandboxIR/Context.h (revision 700637cbb5e582861067a11aaca4d053546871d2)
1 //===- Context.h ------------------------------------------------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #ifndef LLVM_SANDBOXIR_CONTEXT_H
10 #define LLVM_SANDBOXIR_CONTEXT_H
11 
12 #include "llvm/ADT/DenseMap.h"
13 #include "llvm/ADT/MapVector.h"
14 #include "llvm/ADT/SmallVector.h"
15 #include "llvm/IR/LLVMContext.h"
16 #include "llvm/SandboxIR/Tracker.h"
17 #include "llvm/SandboxIR/Type.h"
18 #include "llvm/Support/Compiler.h"
19 
20 #include <cstdint>
21 
22 namespace llvm {
23 namespace sandboxir {
24 
25 class Argument;
26 class BBIterator;
27 class Constant;
28 class Module;
29 class Region;
30 class Value;
31 class Use;
32 
33 class Context {
34 public:
35   // A EraseInstrCallback receives the instruction about to be erased.
36   using EraseInstrCallback = std::function<void(Instruction *)>;
37   // A CreateInstrCallback receives the instruction about to be created.
38   using CreateInstrCallback = std::function<void(Instruction *)>;
39   // A MoveInstrCallback receives the instruction about to be moved, the
40   // destination BB and an iterator pointing to the insertion position.
41   using MoveInstrCallback =
42       std::function<void(Instruction *, const BBIterator &)>;
43   // A SetUseCallback receives the Use that is about to get its source set.
44   using SetUseCallback = std::function<void(const Use &, Value *)>;
45 
46   /// An ID for a registered callback. Used for deregistration. A dedicated type
47   /// is employed so as to keep IDs opaque to the end user; only Context should
48   /// deal with its underlying representation.
49   class CallbackID {
50   public:
51     // Uses a 64-bit integer so we don't have to worry about the unlikely case
52     // of overflowing a 32-bit counter.
53     using ValTy = uint64_t;
54     static constexpr const ValTy InvalidVal = 0;
55 
56   private:
57     // Default initialization results in an invalid ID.
58     ValTy Val = InvalidVal;
CallbackID(ValTy Val)59     explicit CallbackID(ValTy Val) : Val{Val} {
60       assert(Val != InvalidVal && "newly-created ID is invalid!");
61     }
62 
63   public:
64     CallbackID() = default;
65     friend class Context;
66     friend struct DenseMapInfo<CallbackID>;
67   };
68 
69 protected:
70   LLVMContext &LLVMCtx;
71   friend class Type;              // For LLVMCtx.
72   friend class PointerType;       // For LLVMCtx.
73   friend class IntegerType;       // For LLVMCtx.
74   friend class StructType;        // For LLVMCtx.
75   friend class Region;            // For LLVMCtx.
76   friend class IRSnapshotChecker; // To snapshot LLVMModuleToModuleMap.
77 
78   Tracker IRTracker;
79 
80   /// Maps LLVM Value to the corresponding sandboxir::Value. Owns all
81   /// SandboxIR objects.
82   DenseMap<llvm::Value *, std::unique_ptr<Value>> LLVMValueToValueMap;
83 
84   /// Maps an LLVM Module to the corresponding sandboxir::Module.
85   DenseMap<llvm::Module *, std::unique_ptr<Module>> LLVMModuleToModuleMap;
86 
87   /// Type has a protected destructor to prohibit the user from managing the
88   /// lifetime of the Type objects. Context is friend of Type, and this custom
89   /// deleter can destroy Type.
90   struct TypeDeleter {
91     void operator()(Type *Ty) { delete Ty; }
92   };
93   /// Maps LLVM Type to the corresonding sandboxir::Type. Owns all Sandbox IR
94   /// Type objects.
95   DenseMap<llvm::Type *, std::unique_ptr<Type, TypeDeleter>> LLVMTypeToTypeMap;
96 
97   /// Callbacks called when an IR instruction is about to get erased. Keys are
98   /// used as IDs for deregistration.
99   MapVector<CallbackID, EraseInstrCallback> EraseInstrCallbacks;
100   /// Callbacks called when an IR instruction is about to get created. Keys are
101   /// used as IDs for deregistration.
102   MapVector<CallbackID, CreateInstrCallback> CreateInstrCallbacks;
103   /// Callbacks called when an IR instruction is about to get moved. Keys are
104   /// used as IDs for deregistration.
105   MapVector<CallbackID, MoveInstrCallback> MoveInstrCallbacks;
106   /// Callbacks called when a Use gets its source set. Keys are used as IDs for
107   /// deregistration.
108   MapVector<CallbackID, SetUseCallback> SetUseCallbacks;
109 
110   /// A counter used for assigning callback IDs during registration. The same
111   /// counter is used for all kinds of callbacks so we can detect mismatched
112   /// registration/deregistration.
113   CallbackID::ValTy NextCallbackID = 1;
114 
115   /// Remove \p V from the maps and returns the unique_ptr.
116   LLVM_ABI std::unique_ptr<Value> detachLLVMValue(llvm::Value *V);
117   /// Remove \p SBV from all SandboxIR maps and stop owning it. This effectively
118   /// detaches \p V from the underlying IR.
119   LLVM_ABI std::unique_ptr<Value> detach(Value *V);
120   friend class Instruction; // For detach().
121   /// Take ownership of VPtr and store it in `LLVMValueToValueMap`.
122   LLVM_ABI Value *registerValue(std::unique_ptr<Value> &&VPtr);
123   friend class EraseFromParent; // For registerValue().
124   /// This is the actual function that creates sandboxir values for \p V,
125   /// and among others handles all instruction types.
126   LLVM_ABI Value *getOrCreateValueInternal(llvm::Value *V,
127                                            llvm::User *U = nullptr);
128   /// Get or create a sandboxir::Argument for an existing LLVM IR \p LLVMArg.
129   LLVM_ABI Argument *getOrCreateArgument(llvm::Argument *LLVMArg);
130   /// Get or create a sandboxir::Value for an existing LLVM IR \p LLVMV.
131   Value *getOrCreateValue(llvm::Value *LLVMV) {
132     return getOrCreateValueInternal(LLVMV, 0);
133   }
134   /// Get or create a sandboxir::Constant from an existing LLVM IR \p LLVMC.
135   LLVM_ABI Constant *getOrCreateConstant(llvm::Constant *LLVMC);
136   friend class ConstantDataSequential; // For getOrCreateConstant().
137   friend class Utils; // For getMemoryBase
138 
139   LLVM_ABI void runEraseInstrCallbacks(Instruction *I);
140   LLVM_ABI void runCreateInstrCallbacks(Instruction *I);
141   LLVM_ABI void runMoveInstrCallbacks(Instruction *I, const BBIterator &Where);
142   LLVM_ABI void runSetUseCallbacks(const Use &U, Value *NewSrc);
143 
144   friend class User;  // For runSetUseCallbacks().
145   friend class Value; // For runSetUseCallbacks().
146 
147   // Friends for getOrCreateConstant().
148 #define DEF_CONST(ID, CLASS) friend class CLASS;
149 #include "llvm/SandboxIR/Values.def"
150 
151   /// Create a sandboxir::BasicBlock for an existing LLVM IR \p BB. This will
152   /// also create all contents of the block.
153   LLVM_ABI BasicBlock *createBasicBlock(llvm::BasicBlock *BB);
154   friend class BasicBlock; // For getOrCreateValue().
155 
156   IRBuilder<ConstantFolder> LLVMIRBuilder;
157   auto &getLLVMIRBuilder() { return LLVMIRBuilder; }
158 
159   LLVM_ABI VAArgInst *createVAArgInst(llvm::VAArgInst *SI);
160   friend VAArgInst; // For createVAArgInst()
161   LLVM_ABI FreezeInst *createFreezeInst(llvm::FreezeInst *SI);
162   friend FreezeInst; // For createFreezeInst()
163   LLVM_ABI FenceInst *createFenceInst(llvm::FenceInst *SI);
164   friend FenceInst; // For createFenceInst()
165   LLVM_ABI SelectInst *createSelectInst(llvm::SelectInst *SI);
166   friend SelectInst; // For createSelectInst()
167   LLVM_ABI InsertElementInst *
168   createInsertElementInst(llvm::InsertElementInst *IEI);
169   friend InsertElementInst; // For createInsertElementInst()
170   LLVM_ABI ExtractElementInst *
171   createExtractElementInst(llvm::ExtractElementInst *EEI);
172   friend ExtractElementInst; // For createExtractElementInst()
173   LLVM_ABI ShuffleVectorInst *
174   createShuffleVectorInst(llvm::ShuffleVectorInst *SVI);
175   friend ShuffleVectorInst; // For createShuffleVectorInst()
176   LLVM_ABI ExtractValueInst *
177   createExtractValueInst(llvm::ExtractValueInst *IVI);
178   friend ExtractValueInst; // For createExtractValueInst()
179   LLVM_ABI InsertValueInst *createInsertValueInst(llvm::InsertValueInst *IVI);
180   friend InsertValueInst; // For createInsertValueInst()
181   LLVM_ABI BranchInst *createBranchInst(llvm::BranchInst *I);
182   friend BranchInst; // For createBranchInst()
183   LLVM_ABI LoadInst *createLoadInst(llvm::LoadInst *LI);
184   friend LoadInst; // For createLoadInst()
185   LLVM_ABI StoreInst *createStoreInst(llvm::StoreInst *SI);
186   friend StoreInst; // For createStoreInst()
187   LLVM_ABI ReturnInst *createReturnInst(llvm::ReturnInst *I);
188   friend ReturnInst; // For createReturnInst()
189   LLVM_ABI CallInst *createCallInst(llvm::CallInst *I);
190   friend CallInst; // For createCallInst()
191   LLVM_ABI InvokeInst *createInvokeInst(llvm::InvokeInst *I);
192   friend InvokeInst; // For createInvokeInst()
193   LLVM_ABI CallBrInst *createCallBrInst(llvm::CallBrInst *I);
194   friend CallBrInst; // For createCallBrInst()
195   LLVM_ABI LandingPadInst *createLandingPadInst(llvm::LandingPadInst *I);
196   friend LandingPadInst; // For createLandingPadInst()
197   LLVM_ABI CatchPadInst *createCatchPadInst(llvm::CatchPadInst *I);
198   friend CatchPadInst; // For createCatchPadInst()
199   LLVM_ABI CleanupPadInst *createCleanupPadInst(llvm::CleanupPadInst *I);
200   friend CleanupPadInst; // For createCleanupPadInst()
201   LLVM_ABI CatchReturnInst *createCatchReturnInst(llvm::CatchReturnInst *I);
202   friend CatchReturnInst; // For createCatchReturnInst()
203   LLVM_ABI CleanupReturnInst *
204   createCleanupReturnInst(llvm::CleanupReturnInst *I);
205   friend CleanupReturnInst; // For createCleanupReturnInst()
206   LLVM_ABI GetElementPtrInst *
207   createGetElementPtrInst(llvm::GetElementPtrInst *I);
208   friend GetElementPtrInst; // For createGetElementPtrInst()
209   LLVM_ABI CatchSwitchInst *createCatchSwitchInst(llvm::CatchSwitchInst *I);
210   friend CatchSwitchInst; // For createCatchSwitchInst()
211   LLVM_ABI ResumeInst *createResumeInst(llvm::ResumeInst *I);
212   friend ResumeInst; // For createResumeInst()
213   LLVM_ABI SwitchInst *createSwitchInst(llvm::SwitchInst *I);
214   friend SwitchInst; // For createSwitchInst()
215   LLVM_ABI UnaryOperator *createUnaryOperator(llvm::UnaryOperator *I);
216   friend UnaryOperator; // For createUnaryOperator()
217   LLVM_ABI BinaryOperator *createBinaryOperator(llvm::BinaryOperator *I);
218   friend BinaryOperator; // For createBinaryOperator()
219   LLVM_ABI AtomicRMWInst *createAtomicRMWInst(llvm::AtomicRMWInst *I);
220   friend AtomicRMWInst; // For createAtomicRMWInst()
221   LLVM_ABI AtomicCmpXchgInst *
222   createAtomicCmpXchgInst(llvm::AtomicCmpXchgInst *I);
223   friend AtomicCmpXchgInst; // For createAtomicCmpXchgInst()
224   LLVM_ABI AllocaInst *createAllocaInst(llvm::AllocaInst *I);
225   friend AllocaInst; // For createAllocaInst()
226   LLVM_ABI CastInst *createCastInst(llvm::CastInst *I);
227   friend CastInst; // For createCastInst()
228   LLVM_ABI PHINode *createPHINode(llvm::PHINode *I);
229   friend PHINode; // For createPHINode()
230   LLVM_ABI UnreachableInst *createUnreachableInst(llvm::UnreachableInst *UI);
231   friend UnreachableInst; // For createUnreachableInst()
232   LLVM_ABI CmpInst *createCmpInst(llvm::CmpInst *I);
233   friend CmpInst; // For createCmpInst()
234   LLVM_ABI ICmpInst *createICmpInst(llvm::ICmpInst *I);
235   friend ICmpInst; // For createICmpInst()
236   LLVM_ABI FCmpInst *createFCmpInst(llvm::FCmpInst *I);
237   friend FCmpInst; // For createFCmpInst()
238 
239 public:
240   LLVM_ABI Context(LLVMContext &LLVMCtx);
241   LLVM_ABI ~Context();
242   /// Clears function-level state.
243   LLVM_ABI void clear();
244 
245   Tracker &getTracker() { return IRTracker; }
246   /// Convenience function for `getTracker().save()`
247   void save() { IRTracker.save(); }
248   /// Convenience function for `getTracker().revert()`
249   void revert() { IRTracker.revert(); }
250   /// Convenience function for `getTracker().accept()`
251   void accept() { IRTracker.accept(); }
252 
253   LLVM_ABI sandboxir::Value *getValue(llvm::Value *V) const;
254   const sandboxir::Value *getValue(const llvm::Value *V) const {
255     return getValue(const_cast<llvm::Value *>(V));
256   }
257 
258   LLVM_ABI Module *getModule(llvm::Module *LLVMM) const;
259 
260   LLVM_ABI Module *getOrCreateModule(llvm::Module *LLVMM);
261 
262   Type *getType(llvm::Type *LLVMTy) {
263     if (LLVMTy == nullptr)
264       return nullptr;
265     auto Pair = LLVMTypeToTypeMap.try_emplace(LLVMTy);
266     auto It = Pair.first;
267     if (Pair.second)
268       It->second = std::unique_ptr<Type, TypeDeleter>(new Type(LLVMTy, *this));
269     return It->second.get();
270   }
271 
272   /// Create a sandboxir::Function for an existing LLVM IR \p F, including all
273   /// blocks and instructions.
274   /// This is the main API function for creating Sandbox IR.
275   /// Note: this will not fully populate its parent module. The only globals
276   /// that will be available are those used within the function.
277   LLVM_ABI Function *createFunction(llvm::Function *F);
278 
279   /// Create a sandboxir::Module corresponding to \p LLVMM.
280   LLVM_ABI Module *createModule(llvm::Module *LLVMM);
281 
282   /// \Returns the number of values registered with Context.
283   size_t getNumValues() const { return LLVMValueToValueMap.size(); }
284 
285   /// Register a callback that gets called when a SandboxIR instruction is about
286   /// to be removed from its parent. Note that this will also be called when
287   /// reverting the creation of an instruction.
288   /// \Returns a callback ID for later deregistration.
289   LLVM_ABI CallbackID registerEraseInstrCallback(EraseInstrCallback CB);
290   LLVM_ABI void unregisterEraseInstrCallback(CallbackID ID);
291 
292   /// Register a callback that gets called right after a SandboxIR instruction
293   /// is created. Note that this will also be called when reverting the removal
294   /// of an instruction.
295   /// \Returns a callback ID for later deregistration.
296   LLVM_ABI CallbackID registerCreateInstrCallback(CreateInstrCallback CB);
297   LLVM_ABI void unregisterCreateInstrCallback(CallbackID ID);
298 
299   /// Register a callback that gets called when a SandboxIR instruction is about
300   /// to be moved. Note that this will also be called when reverting a move.
301   /// \Returns a callback ID for later deregistration.
302   LLVM_ABI CallbackID registerMoveInstrCallback(MoveInstrCallback CB);
303   LLVM_ABI void unregisterMoveInstrCallback(CallbackID ID);
304 
305   /// Register a callback that gets called when a Use gets set.
306   /// \Returns a callback ID for later deregistration.
307   LLVM_ABI CallbackID registerSetUseCallback(SetUseCallback CB);
308   LLVM_ABI void unregisterSetUseCallback(CallbackID ID);
309 };
310 
311 } // namespace sandboxir
312 
313 // DenseMap info for CallbackIDs
314 template <> struct DenseMapInfo<sandboxir::Context::CallbackID> {
315   using CallbackID = sandboxir::Context::CallbackID;
316   using ReprInfo = DenseMapInfo<CallbackID::ValTy>;
317 
318   static CallbackID getEmptyKey() {
319     return CallbackID{ReprInfo::getEmptyKey()};
320   }
321   static CallbackID getTombstoneKey() {
322     return CallbackID{ReprInfo::getTombstoneKey()};
323   }
324   static unsigned getHashValue(const CallbackID &ID) {
325     return ReprInfo::getHashValue(ID.Val);
326   }
327   static bool isEqual(const CallbackID &LHS, const CallbackID &RHS) {
328     return ReprInfo::isEqual(LHS.Val, RHS.Val);
329   }
330 };
331 
332 } // namespace llvm
333 
334 #endif // LLVM_SANDBOXIR_CONTEXT_H
335