llvm/SandboxIR/Context.h

0001 //===- Context.h ------------------------------------------------*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008
0009 #ifndef LLVM_SANDBOXIR_CONTEXT_H
0010 #define LLVM_SANDBOXIR_CONTEXT_H
0011
0012 #include "llvm/ADT/DenseMap.h"
0013 #include "llvm/ADT/MapVector.h"
0014 #include "llvm/ADT/SmallVector.h"
0015 #include "llvm/IR/LLVMContext.h"
0016 #include "llvm/SandboxIR/Tracker.h"
0017 #include "llvm/SandboxIR/Type.h"
0018
0019 #include <cstdint>
0020
0021 namespace llvm {
0022 namespace sandboxir {
0023
0024 class Argument;
0025 class BBIterator;
0026 class Constant;
0027 class Module;
0028 class Value;
0029
0030 class Context {
0031 public:
0032   // A EraseInstrCallback receives the instruction about to be erased.
0033   using EraseInstrCallback = std::function<void(Instruction *)>;
0034   // A CreateInstrCallback receives the instruction about to be created.
0035   using CreateInstrCallback = std::function<void(Instruction *)>;
0036   // A MoveInstrCallback receives the instruction about to be moved, the
0037   // destination BB and an iterator pointing to the insertion position.
0038   using MoveInstrCallback =
0039       std::function<void(Instruction *, const BBIterator &)>;
0040
0041   /// An ID for a registered callback. Used for deregistration. A dedicated type
0042   /// is employed so as to keep IDs opaque to the end user; only Context should
0043   /// deal with its underlying representation.
0044   class CallbackID {
0045   public:
0046     // Uses a 64-bit integer so we don't have to worry about the unlikely case
0047     // of overflowing a 32-bit counter.
0048     using ValTy = uint64_t;
0049     static constexpr const ValTy InvalidVal = 0;
0050
0051   private:
0052     // Default initialization results in an invalid ID.
0053     ValTy Val = InvalidVal;
0054     explicit CallbackID(ValTy Val) : Val{Val} {
0055       assert(Val != InvalidVal && "newly-created ID is invalid!");
0056     }
0057
0058   public:
0059     CallbackID() = default;
0060     friend class Context;
0061     friend struct DenseMapInfo<CallbackID>;
0062   };
0063
0064 protected:
0065   LLVMContext &LLVMCtx;
0066   friend class Type;              // For LLVMCtx.
0067   friend class PointerType;       // For LLVMCtx.
0068   friend class IntegerType;       // For LLVMCtx.
0069   friend class StructType;        // For LLVMCtx.
0070   friend class Region;            // For LLVMCtx.
0071   friend class IRSnapshotChecker; // To snapshot LLVMModuleToModuleMap.
0072
0073   Tracker IRTracker;
0074
0075   /// Maps LLVM Value to the corresponding sandboxir::Value. Owns all
0076   /// SandboxIR objects.
0077   DenseMap<llvm::Value *, std::unique_ptr<Value>> LLVMValueToValueMap;
0078
0079   /// Maps an LLVM Module to the corresponding sandboxir::Module.
0080   DenseMap<llvm::Module *, std::unique_ptr<Module>> LLVMModuleToModuleMap;
0081
0082   /// Type has a protected destructor to prohibit the user from managing the
0083   /// lifetime of the Type objects. Context is friend of Type, and this custom
0084   /// deleter can destroy Type.
0085   struct TypeDeleter {
0086     void operator()(Type *Ty) { delete Ty; }
0087   };
0088   /// Maps LLVM Type to the corresonding sandboxir::Type. Owns all Sandbox IR
0089   /// Type objects.
0090   DenseMap<llvm::Type *, std::unique_ptr<Type, TypeDeleter>> LLVMTypeToTypeMap;
0091
0092   /// Callbacks called when an IR instruction is about to get erased. Keys are
0093   /// used as IDs for deregistration.
0094   MapVector<CallbackID, EraseInstrCallback> EraseInstrCallbacks;
0095   /// Callbacks called when an IR instruction is about to get created. Keys are
0096   /// used as IDs for deregistration.
0097   MapVector<CallbackID, CreateInstrCallback> CreateInstrCallbacks;
0098   /// Callbacks called when an IR instruction is about to get moved. Keys are
0099   /// used as IDs for deregistration.
0100   MapVector<CallbackID, MoveInstrCallback> MoveInstrCallbacks;
0101
0102   /// A counter used for assigning callback IDs during registration. The same
0103   /// counter is used for all kinds of callbacks so we can detect mismatched
0104   /// registration/deregistration.
0105   CallbackID::ValTy NextCallbackID = 1;
0106
0107   /// Remove \p V from the maps and returns the unique_ptr.
0108   std::unique_ptr<Value> detachLLVMValue(llvm::Value *V);
0109   /// Remove \p SBV from all SandboxIR maps and stop owning it. This effectively
0110   /// detaches \p V from the underlying IR.
0111   std::unique_ptr<Value> detach(Value *V);
0112   friend class Instruction; // For detach().
0113   /// Take ownership of VPtr and store it in `LLVMValueToValueMap`.
0114   Value *registerValue(std::unique_ptr<Value> &&VPtr);
0115   friend class EraseFromParent; // For registerValue().
0116   /// This is the actual function that creates sandboxir values for \p V,
0117   /// and among others handles all instruction types.
0118   Value *getOrCreateValueInternal(llvm::Value *V, llvm::User *U = nullptr);
0119   /// Get or create a sandboxir::Argument for an existing LLVM IR \p LLVMArg.
0120   Argument *getOrCreateArgument(llvm::Argument *LLVMArg);
0121   /// Get or create a sandboxir::Value for an existing LLVM IR \p LLVMV.
0122   Value *getOrCreateValue(llvm::Value *LLVMV) {
0123     return getOrCreateValueInternal(LLVMV, 0);
0124   }
0125   /// Get or create a sandboxir::Constant from an existing LLVM IR \p LLVMC.
0126   Constant *getOrCreateConstant(llvm::Constant *LLVMC);
0127   friend class Utils; // For getMemoryBase
0128
0129   void runEraseInstrCallbacks(Instruction *I);
0130   void runCreateInstrCallbacks(Instruction *I);
0131   void runMoveInstrCallbacks(Instruction *I, const BBIterator &Where);
0132
0133   // Friends for getOrCreateConstant().
0134 #define DEF_CONST(ID, CLASS) friend class CLASS;
0135 #include "llvm/SandboxIR/Values.def"
0136
0137   /// Create a sandboxir::BasicBlock for an existing LLVM IR \p BB. This will
0138   /// also create all contents of the block.
0139   BasicBlock *createBasicBlock(llvm::BasicBlock *BB);
0140   friend class BasicBlock; // For getOrCreateValue().
0141
0142   IRBuilder<ConstantFolder> LLVMIRBuilder;
0143   auto &getLLVMIRBuilder() { return LLVMIRBuilder; }
0144
0145   VAArgInst *createVAArgInst(llvm::VAArgInst *SI);
0146   friend VAArgInst; // For createVAArgInst()
0147   FreezeInst *createFreezeInst(llvm::FreezeInst *SI);
0148   friend FreezeInst; // For createFreezeInst()
0149   FenceInst *createFenceInst(llvm::FenceInst *SI);
0150   friend FenceInst; // For createFenceInst()
0151   SelectInst *createSelectInst(llvm::SelectInst *SI);
0152   friend SelectInst; // For createSelectInst()
0153   InsertElementInst *createInsertElementInst(llvm::InsertElementInst *IEI);
0154   friend InsertElementInst; // For createInsertElementInst()
0155   ExtractElementInst *createExtractElementInst(llvm::ExtractElementInst *EEI);
0156   friend ExtractElementInst; // For createExtractElementInst()
0157   ShuffleVectorInst *createShuffleVectorInst(llvm::ShuffleVectorInst *SVI);
0158   friend ShuffleVectorInst; // For createShuffleVectorInst()
0159   ExtractValueInst *createExtractValueInst(llvm::ExtractValueInst *IVI);
0160   friend ExtractValueInst; // For createExtractValueInst()
0161   InsertValueInst *createInsertValueInst(llvm::InsertValueInst *IVI);
0162   friend InsertValueInst; // For createInsertValueInst()
0163   BranchInst *createBranchInst(llvm::BranchInst *I);
0164   friend BranchInst; // For createBranchInst()
0165   LoadInst *createLoadInst(llvm::LoadInst *LI);
0166   friend LoadInst; // For createLoadInst()
0167   StoreInst *createStoreInst(llvm::StoreInst *SI);
0168   friend StoreInst; // For createStoreInst()
0169   ReturnInst *createReturnInst(llvm::ReturnInst *I);
0170   friend ReturnInst; // For createReturnInst()
0171   CallInst *createCallInst(llvm::CallInst *I);
0172   friend CallInst; // For createCallInst()
0173   InvokeInst *createInvokeInst(llvm::InvokeInst *I);
0174   friend InvokeInst; // For createInvokeInst()
0175   CallBrInst *createCallBrInst(llvm::CallBrInst *I);
0176   friend CallBrInst; // For createCallBrInst()
0177   LandingPadInst *createLandingPadInst(llvm::LandingPadInst *I);
0178   friend LandingPadInst; // For createLandingPadInst()
0179   CatchPadInst *createCatchPadInst(llvm::CatchPadInst *I);
0180   friend CatchPadInst; // For createCatchPadInst()
0181   CleanupPadInst *createCleanupPadInst(llvm::CleanupPadInst *I);
0182   friend CleanupPadInst; // For createCleanupPadInst()
0183   CatchReturnInst *createCatchReturnInst(llvm::CatchReturnInst *I);
0184   friend CatchReturnInst; // For createCatchReturnInst()
0185   CleanupReturnInst *createCleanupReturnInst(llvm::CleanupReturnInst *I);
0186   friend CleanupReturnInst; // For createCleanupReturnInst()
0187   GetElementPtrInst *createGetElementPtrInst(llvm::GetElementPtrInst *I);
0188   friend GetElementPtrInst; // For createGetElementPtrInst()
0189   CatchSwitchInst *createCatchSwitchInst(llvm::CatchSwitchInst *I);
0190   friend CatchSwitchInst; // For createCatchSwitchInst()
0191   ResumeInst *createResumeInst(llvm::ResumeInst *I);
0192   friend ResumeInst; // For createResumeInst()
0193   SwitchInst *createSwitchInst(llvm::SwitchInst *I);
0194   friend SwitchInst; // For createSwitchInst()
0195   UnaryOperator *createUnaryOperator(llvm::UnaryOperator *I);
0196   friend UnaryOperator; // For createUnaryOperator()
0197   BinaryOperator *createBinaryOperator(llvm::BinaryOperator *I);
0198   friend BinaryOperator; // For createBinaryOperator()
0199   AtomicRMWInst *createAtomicRMWInst(llvm::AtomicRMWInst *I);
0200   friend AtomicRMWInst; // For createAtomicRMWInst()
0201   AtomicCmpXchgInst *createAtomicCmpXchgInst(llvm::AtomicCmpXchgInst *I);
0202   friend AtomicCmpXchgInst; // For createAtomicCmpXchgInst()
0203   AllocaInst *createAllocaInst(llvm::AllocaInst *I);
0204   friend AllocaInst; // For createAllocaInst()
0205   CastInst *createCastInst(llvm::CastInst *I);
0206   friend CastInst; // For createCastInst()
0207   PHINode *createPHINode(llvm::PHINode *I);
0208   friend PHINode; // For createPHINode()
0209   UnreachableInst *createUnreachableInst(llvm::UnreachableInst *UI);
0210   friend UnreachableInst; // For createUnreachableInst()
0211   CmpInst *createCmpInst(llvm::CmpInst *I);
0212   friend CmpInst; // For createCmpInst()
0213   ICmpInst *createICmpInst(llvm::ICmpInst *I);
0214   friend ICmpInst; // For createICmpInst()
0215   FCmpInst *createFCmpInst(llvm::FCmpInst *I);
0216   friend FCmpInst; // For createFCmpInst()
0217
0218 public:
0219   Context(LLVMContext &LLVMCtx);
0220   ~Context();
0221   /// Clears function-level state.
0222   void clear();
0223
0224   Tracker &getTracker() { return IRTracker; }
0225   /// Convenience function for `getTracker().save()`
0226   void save() { IRTracker.save(); }
0227   /// Convenience function for `getTracker().revert()`
0228   void revert() { IRTracker.revert(); }
0229   /// Convenience function for `getTracker().accept()`
0230   void accept() { IRTracker.accept(); }
0231
0232   sandboxir::Value *getValue(llvm::Value *V) const;
0233   const sandboxir::Value *getValue(const llvm::Value *V) const {
0234     return getValue(const_cast<llvm::Value *>(V));
0235   }
0236
0237   Module *getModule(llvm::Module *LLVMM) const;
0238
0239   Module *getOrCreateModule(llvm::Module *LLVMM);
0240
0241   Type *getType(llvm::Type *LLVMTy) {
0242     if (LLVMTy == nullptr)
0243       return nullptr;
0244     auto Pair = LLVMTypeToTypeMap.insert({LLVMTy, nullptr});
0245     auto It = Pair.first;
0246     if (Pair.second)
0247       It->second = std::unique_ptr<Type, TypeDeleter>(new Type(LLVMTy, *this));
0248     return It->second.get();
0249   }
0250
0251   /// Create a sandboxir::Function for an existing LLVM IR \p F, including all
0252   /// blocks and instructions.
0253   /// This is the main API function for creating Sandbox IR.
0254   /// Note: this will not fully populate its parent module. The only globals
0255   /// that will be available are those used within the function.
0256   Function *createFunction(llvm::Function *F);
0257
0258   /// Create a sandboxir::Module corresponding to \p LLVMM.
0259   Module *createModule(llvm::Module *LLVMM);
0260
0261   /// \Returns the number of values registered with Context.
0262   size_t getNumValues() const { return LLVMValueToValueMap.size(); }
0263
0264   /// Register a callback that gets called when a SandboxIR instruction is about
0265   /// to be removed from its parent. Note that this will also be called when
0266   /// reverting the creation of an instruction.
0267   /// \Returns a callback ID for later deregistration.
0268   CallbackID registerEraseInstrCallback(EraseInstrCallback CB);
0269   void unregisterEraseInstrCallback(CallbackID ID);
0270
0271   /// Register a callback that gets called right after a SandboxIR instruction
0272   /// is created. Note that this will also be called when reverting the removal
0273   /// of an instruction.
0274   /// \Returns a callback ID for later deregistration.
0275   CallbackID registerCreateInstrCallback(CreateInstrCallback CB);
0276   void unregisterCreateInstrCallback(CallbackID ID);
0277
0278   /// Register a callback that gets called when a SandboxIR instruction is about
0279   /// to be moved. Note that this will also be called when reverting a move.
0280   /// \Returns a callback ID for later deregistration.
0281   CallbackID registerMoveInstrCallback(MoveInstrCallback CB);
0282   void unregisterMoveInstrCallback(CallbackID ID);
0283
0284   // TODO: Add callbacks for instructions inserted/removed if needed.
0285 };
0286
0287 } // namespace sandboxir
0288
0289 // DenseMap info for CallbackIDs
0290 template <> struct DenseMapInfo<sandboxir::Context::CallbackID> {
0291   using CallbackID = sandboxir::Context::CallbackID;
0292   using ReprInfo = DenseMapInfo<CallbackID::ValTy>;
0293
0294   static CallbackID getEmptyKey() {
0295     return CallbackID{ReprInfo::getEmptyKey()};
0296   }
0297   static CallbackID getTombstoneKey() {
0298     return CallbackID{ReprInfo::getTombstoneKey()};
0299   }
0300   static unsigned getHashValue(const CallbackID &ID) {
0301     return ReprInfo::getHashValue(ID.Val);
0302   }
0303   static bool isEqual(const CallbackID &LHS, const CallbackID &RHS) {
0304     return ReprInfo::isEqual(LHS.Val, RHS.Val);
0305   }
0306 };
0307
0308 } // namespace llvm
0309
0310 #endif // LLVM_SANDBOXIR_CONTEXT_H