Context.h

Go to the documentation of this file.

//===- Context.h ------------------------------------------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
 
#ifndef LLVM_SANDBOXIR_CONTEXT_H
#define LLVM_SANDBOXIR_CONTEXT_H
 
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/SandboxIR/Tracker.h"
#include "llvm/SandboxIR/Type.h"
#include "llvm/Support/Compiler.h"
 
#include <cstdint>
 
namespace llvm {
namespace sandboxir {
 
class Argument;
class BBIterator;
class Constant;
class Module;
class Region;
class Value;
class Use;
 
class LLVM_ABI Context {
public:
  // A EraseInstrCallback receives the instruction about to be erased.
  using EraseInstrCallback = std::function<void(Instruction *)>;
  // A CreateInstrCallback receives the instruction about to be created.
  using CreateInstrCallback = std::function<void(Instruction *)>;
  // A MoveInstrCallback receives the instruction about to be moved, the
  // destination BB and an iterator pointing to the insertion position.
  using MoveInstrCallback =
      std::function<void(Instruction *, const BBIterator &)>;
  // A SetUseCallback receives the Use that is about to get its source set.
  using SetUseCallback = std::function<void(const Use &, Value *)>;
 
  /// An ID for a registered callback. Used for deregistration. A dedicated type
  /// is employed so as to keep IDs opaque to the end user; only Context should
  /// deal with its underlying representation.
  class CallbackID {
  public:
    // Uses a 64-bit integer so we don't have to worry about the unlikely case
    // of overflowing a 32-bit counter.
    using ValTy = uint64_t;
    static constexpr ValTy InvalidVal = 0;
 
  private:
    // Default initialization results in an invalid ID.
    ValTy Val = InvalidVal;
    explicit CallbackID(ValTy Val) : Val{Val} {
      assert(Val != InvalidVal && "newly-created ID is invalid!");
    }
 
  public:
    CallbackID() = default;
    bool operator==(const CallbackID &Other) const { return Val == Other.Val; }
    friend class Context;
    friend struct DenseMapInfo<CallbackID>;
  };
 
protected:
  LLVMContext &LLVMCtx;
  friend class Type;              // For LLVMCtx.
  friend class PointerType;       // For LLVMCtx.
  friend class IntegerType;       // For LLVMCtx.
  friend class ByteType;          // For LLVMCtx.
  friend class StructType;        // For LLVMCtx.
  friend class Region;            // For LLVMCtx.
  friend class IRSnapshotChecker; // To snapshot LLVMModuleToModuleMap.
 
  Tracker IRTracker;
 
  /// Maps LLVM Value to the corresponding sandboxir::Value. Owns all
  /// SandboxIR objects.
  DenseMap<llvm::Value *, std::unique_ptr<Value>> LLVMValueToValueMap;
 
  /// Maps an LLVM Module to the corresponding sandboxir::Module.
  DenseMap<llvm::Module *, std::unique_ptr<Module>> LLVMModuleToModuleMap;
 
  /// Type has a protected destructor to prohibit the user from managing the
  /// lifetime of the Type objects. Context is friend of Type, and this custom
  /// deleter can destroy Type.
  struct TypeDeleter {
    void operator()(Type *Ty) { delete Ty; }
  };
  /// Maps LLVM Type to the corresonding sandboxir::Type. Owns all Sandbox IR
  /// Type objects.
  DenseMap<llvm::Type *, std::unique_ptr<Type, TypeDeleter>> LLVMTypeToTypeMap;
 
  /// Callbacks called when an IR instruction is about to get erased. CallbackID
  /// is used as an identifier for deregistration.
  SmallVector<std::pair<CallbackID, EraseInstrCallback>> EraseInstrCallbacks;
  /// Callbacks called when an IR instruction is about to get created.
  /// CallbackID is used as an identifier for deregistration.
  SmallVector<std::pair<CallbackID, CreateInstrCallback>> CreateInstrCallbacks;
  /// Callbacks called when an IR instruction is about to get moved. CallbackID
  /// is used as an identifier for deregistration.
  SmallVector<std::pair<CallbackID, MoveInstrCallback>> MoveInstrCallbacks;
  /// Callbacks called when a Use gets its source set. CallbackID is used as an
  /// identifier for deregistration.
  SmallVector<std::pair<CallbackID, SetUseCallback>> SetUseCallbacks;
 
  /// A counter used for assigning callback IDs during registration. The same
  /// counter is used for all kinds of callbacks so we can detect mismatched
  /// registration/deregistration.
  CallbackID::ValTy NextCallbackID = 1;
 
  /// Remove \p V from the maps and returns the unique_ptr.
  std::unique_ptr<Value> detachLLVMValue(llvm::Value *V);
  /// Remove \p SBV from all SandboxIR maps and stop owning it. This effectively
  /// detaches \p V from the underlying IR.
  std::unique_ptr<Value> detach(Value *V);
  friend class Instruction; // For detach().
  /// Take ownership of VPtr and store it in `LLVMValueToValueMap`.
  Value *registerValue(std::unique_ptr<Value> &&VPtr);
  friend class EraseFromParent; // For registerValue().
  /// This is the actual function that creates sandboxir values for \p V,
  /// and among others handles all instruction types.
  Value *getOrCreateValueInternal(llvm::Value *V, llvm::User *U = nullptr);
  /// Get or create a sandboxir::Argument for an existing LLVM IR \p LLVMArg.
  Argument *getOrCreateArgument(llvm::Argument *LLVMArg);
  /// Get or create a sandboxir::Value for an existing LLVM IR \p LLVMV.
  Value *getOrCreateValue(llvm::Value *LLVMV) {
    return getOrCreateValueInternal(LLVMV, 0);
  }
  /// Get or create a sandboxir::Constant from an existing LLVM IR \p LLVMC.
  Constant *getOrCreateConstant(llvm::Constant *LLVMC);
  friend class ConstantDataSequential; // For getOrCreateConstant().
  friend class Utils; // For getMemoryBase
 
  void runEraseInstrCallbacks(Instruction *I);
  void runCreateInstrCallbacks(Instruction *I);
  void runMoveInstrCallbacks(Instruction *I, const BBIterator &Where);
  void runSetUseCallbacks(const Use &U, Value *NewSrc);
 
  friend class User;  // For runSetUseCallbacks().
  friend class Value; // For runSetUseCallbacks().
 
  // Friends for getOrCreateConstant().
#define DEF_CONST(ID, CLASS) friend class CLASS;
#include "llvm/SandboxIR/Values.def"
 
  /// Create a sandboxir::BasicBlock for an existing LLVM IR \p BB. This will
  /// also create all contents of the block.
  BasicBlock *createBasicBlock(llvm::BasicBlock *BB);
  friend class BasicBlock; // For getOrCreateValue().
 
  VAArgInst *createVAArgInst(llvm::VAArgInst *SI);
  friend VAArgInst; // For createVAArgInst()
  FreezeInst *createFreezeInst(llvm::FreezeInst *SI);
  friend FreezeInst; // For createFreezeInst()
  FenceInst *createFenceInst(llvm::FenceInst *SI);
  friend FenceInst; // For createFenceInst()
  SelectInst *createSelectInst(llvm::SelectInst *SI);
  friend SelectInst; // For createSelectInst()
  InsertElementInst *createInsertElementInst(llvm::InsertElementInst *IEI);
  friend InsertElementInst; // For createInsertElementInst()
  ExtractElementInst *createExtractElementInst(llvm::ExtractElementInst *EEI);
  friend ExtractElementInst; // For createExtractElementInst()
  ShuffleVectorInst *createShuffleVectorInst(llvm::ShuffleVectorInst *SVI);
  friend ShuffleVectorInst; // For createShuffleVectorInst()
  ExtractValueInst *createExtractValueInst(llvm::ExtractValueInst *IVI);
  friend ExtractValueInst; // For createExtractValueInst()
  InsertValueInst *createInsertValueInst(llvm::InsertValueInst *IVI);
  friend InsertValueInst; // For createInsertValueInst()
  UncondBrInst *createUncondBrInst(llvm::UncondBrInst *UBI);
  friend UncondBrInst; // For createUncondBrInst()
  CondBrInst *createCondBrInst(llvm::CondBrInst *CBI);
  friend CondBrInst; // For createCondBrInst()
  LoadInst *createLoadInst(llvm::LoadInst *LI);
  friend LoadInst; // For createLoadInst()
  StoreInst *createStoreInst(llvm::StoreInst *SI);
  friend StoreInst; // For createStoreInst()
  ReturnInst *createReturnInst(llvm::ReturnInst *I);
  friend ReturnInst; // For createReturnInst()
  CallInst *createCallInst(llvm::CallInst *I);
  friend CallInst; // For createCallInst()
  InvokeInst *createInvokeInst(llvm::InvokeInst *I);
  friend InvokeInst; // For createInvokeInst()
  CallBrInst *createCallBrInst(llvm::CallBrInst *I);
  friend CallBrInst; // For createCallBrInst()
  LandingPadInst *createLandingPadInst(llvm::LandingPadInst *I);
  friend LandingPadInst; // For createLandingPadInst()
  CatchPadInst *createCatchPadInst(llvm::CatchPadInst *I);
  friend CatchPadInst; // For createCatchPadInst()
  CleanupPadInst *createCleanupPadInst(llvm::CleanupPadInst *I);
  friend CleanupPadInst; // For createCleanupPadInst()
  CatchReturnInst *createCatchReturnInst(llvm::CatchReturnInst *I);
  friend CatchReturnInst; // For createCatchReturnInst()
  CleanupReturnInst *createCleanupReturnInst(llvm::CleanupReturnInst *I);
  friend CleanupReturnInst; // For createCleanupReturnInst()
  GetElementPtrInst *createGetElementPtrInst(llvm::GetElementPtrInst *I);
  friend GetElementPtrInst; // For createGetElementPtrInst()
  CatchSwitchInst *createCatchSwitchInst(llvm::CatchSwitchInst *I);
  friend CatchSwitchInst; // For createCatchSwitchInst()
  ResumeInst *createResumeInst(llvm::ResumeInst *I);
  friend ResumeInst; // For createResumeInst()
  SwitchInst *createSwitchInst(llvm::SwitchInst *I);
  friend SwitchInst; // For createSwitchInst()
  UnaryOperator *createUnaryOperator(llvm::UnaryOperator *I);
  friend UnaryOperator; // For createUnaryOperator()
  BinaryOperator *createBinaryOperator(llvm::BinaryOperator *I);
  friend BinaryOperator; // For createBinaryOperator()
  AtomicRMWInst *createAtomicRMWInst(llvm::AtomicRMWInst *I);
  friend AtomicRMWInst; // For createAtomicRMWInst()
  AtomicCmpXchgInst *createAtomicCmpXchgInst(llvm::AtomicCmpXchgInst *I);
  friend AtomicCmpXchgInst; // For createAtomicCmpXchgInst()
  AllocaInst *createAllocaInst(llvm::AllocaInst *I);
  friend AllocaInst; // For createAllocaInst()
  CastInst *createCastInst(llvm::CastInst *I);
  friend CastInst; // For createCastInst()
  PHINode *createPHINode(llvm::PHINode *I);
  friend PHINode; // For createPHINode()
  UnreachableInst *createUnreachableInst(llvm::UnreachableInst *UI);
  friend UnreachableInst; // For createUnreachableInst()
  CmpInst *createCmpInst(llvm::CmpInst *I);
  friend CmpInst; // For createCmpInst()
  ICmpInst *createICmpInst(llvm::ICmpInst *I);
  friend ICmpInst; // For createICmpInst()
  FCmpInst *createFCmpInst(llvm::FCmpInst *I);
  friend FCmpInst; // For createFCmpInst()
 
public:
  Context(LLVMContext &LLVMCtx);
  Context(const Context &) = delete;
  Context &operator=(const Context &) = delete;
  virtual ~Context();
  /// Clears function-level state.
  void clear();
 
  Tracker &getTracker() { return IRTracker; }
  /// Convenience function for `getTracker().save()`
  void save() { IRTracker.save(); }
  /// Convenience function for `getTracker().revert()`
  void revert(bool RevertAll = false) { IRTracker.revert(RevertAll); }
  /// Convenience function for `getTracker().accept()`
  void accept(bool AcceptAll = false) { IRTracker.accept(AcceptAll); }
 
  sandboxir::Value *getValue(llvm::Value *V) const;
  const sandboxir::Value *getValue(const llvm::Value *V) const {
    return getValue(const_cast<llvm::Value *>(V));
  }
 
  Module *getModule(llvm::Module *LLVMM) const;
 
  Module *getOrCreateModule(llvm::Module *LLVMM);
 
  Type *getType(llvm::Type *LLVMTy) {
    if (LLVMTy == nullptr)
      return nullptr;
    auto Pair = LLVMTypeToTypeMap.try_emplace(LLVMTy);
    auto It = Pair.first;
    if (Pair.second)
      It->second = std::unique_ptr<Type, TypeDeleter>(new Type(LLVMTy, *this));
    return It->second.get();
  }
 
  /// Create a sandboxir::Function for an existing LLVM IR \p F, including all
  /// blocks and instructions.
  /// This is the main API function for creating Sandbox IR.
  /// Note: this will not fully populate its parent module. The only globals
  /// that will be available are those used within the function.
  Function *createFunction(llvm::Function *F);
 
  /// Create a sandboxir::Module corresponding to \p LLVMM.
  Module *createModule(llvm::Module *LLVMM);
 
  /// \Returns the number of values registered with Context.
  size_t getNumValues() const { return LLVMValueToValueMap.size(); }
 
private:
  // An arbitrary limit, to check for accidental misuse. We expect a small
  // number of callbacks to be registered at a time, but we can increase this
  // number if we discover we needed more.
  [[maybe_unused]] static constexpr int MaxRegisteredCallbacks = 16;
 
  template <typename CBT, typename CBVecT>
  CallbackID registerCallbackCommon(CBT CB, CBVecT &CBVec,
                                    std::optional<CallbackID> BeforeID) {
    assert(CBVec.size() <= MaxRegisteredCallbacks &&
           "EraseInstrCallbacks size limit exceeded");
    auto BeforeIt = CBVec.end();
    if (BeforeID) {
      BeforeIt = find_if(CBVec, [BeforeID](const auto &Pair) {
        return Pair.first == *BeforeID;
      });
      assert(BeforeIt != CBVec.end() && "Not found!");
    }
    CallbackID ID{NextCallbackID++};
    CBVec.insert(BeforeIt, {ID, std::move(CB)});
    return ID;
  }
 
public:
  /// Register a callback that gets called when a SandboxIR instruction is about
  /// to be removed from its parent. Note that this will also be called when
  /// reverting the creation of an instruction.
  /// If \p BeforeID is specified, CB will be ordered just before \p BeforeID,
  /// so it will be called first once the callbacks get executed.
  /// \Returns a callback ID for later deregistration.
  CallbackID
  registerEraseInstrCallback(EraseInstrCallback CB,
                             std::optional<CallbackID> BeforeID = std::nullopt);
  void unregisterEraseInstrCallback(CallbackID ID);
 
  /// Register a callback that gets called right after a SandboxIR instruction
  /// is created. Note that this will also be called when reverting the removal
  /// of an instruction.
  /// If \p BeforeID is specified, CB will be ordered just before \p BeforeID,
  /// so it will be called first once the callbacks get executed.
  /// \Returns a callback ID for later deregistration.
  CallbackID registerCreateInstrCallback(
      CreateInstrCallback CB,
      std::optional<CallbackID> BeforeID = std::nullopt);
  void unregisterCreateInstrCallback(CallbackID ID);
 
  /// Register a callback that gets called when a SandboxIR instruction is about
  /// to be moved. Note that this will also be called when reverting a move.
  /// If \p BeforeID is specified, CB will be ordered just before \p BeforeID,
  /// so it will be called first once the callbacks get executed.
  /// \Returns a callback ID for later deregistration.
  CallbackID
  registerMoveInstrCallback(MoveInstrCallback CB,
                            std::optional<CallbackID> BeforeID = std::nullopt);
  void unregisterMoveInstrCallback(CallbackID ID);
 
  /// Register a callback that gets called when a Use gets set.
  /// If \p BeforeID is specified, CB will be ordered just before \p BeforeID,
  /// so it will be called first once the callbacks get executed.
  /// \Returns a callback ID for later deregistration.
  CallbackID
  registerSetUseCallback(SetUseCallback CB,
                         std::optional<CallbackID> BeforeID = std::nullopt);
  void unregisterSetUseCallback(CallbackID ID);
};
 
} // namespace sandboxir
 
// DenseMap info for CallbackIDs
template <> struct DenseMapInfo<sandboxir::Context::CallbackID> {
  using CallbackID = sandboxir::Context::CallbackID;
  using ReprInfo = DenseMapInfo<CallbackID::ValTy>;
 
  static unsigned getHashValue(const CallbackID &ID) {
    return ReprInfo::getHashValue(ID.Val);
  }
  static bool isEqual(const CallbackID &LHS, const CallbackID &RHS) {
    return ReprInfo::isEqual(LHS.Val, RHS.Val);
  }
};
 
} // namespace llvm
 
#endif // LLVM_SANDBOXIR_CONTEXT_H