namespace llvm {
+class DomTreeUpdater;
+
/// EscapeEnumerator - This is a little algorithm to find all escape points
/// from a function so that "finally"-style code can be inserted. In addition
/// to finding the existing return and unwind instructions, it also (if
bool Done;
bool HandleExceptions;
+ DomTreeUpdater *DTU;
+
public:
EscapeEnumerator(Function &F, const char *N = "cleanup",
- bool HandleExceptions = true)
+ bool HandleExceptions = true, DomTreeUpdater *DTU = nullptr)
: F(F), CleanupBBName(N), StateBB(F.begin()), StateE(F.end()),
Builder(F.getContext()), Done(false),
- HandleExceptions(HandleExceptions) {}
+ HandleExceptions(HandleExceptions), DTU(DTU) {}
IRBuilder<> *Next();
};
/// InvokeInst is a terminator instruction. Returns the newly split basic
/// block.
BasicBlock *changeToInvokeAndSplitBasicBlock(CallInst *CI,
- BasicBlock *UnwindEdge);
+ BasicBlock *UnwindEdge,
+ DomTreeUpdater *DTU = nullptr);
/// Replace 'BB's terminator with one that does not have an unwind successor
/// block. Rewrites `invoke` to `call`, etc. Updates any PHIs in unwind
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringExtras.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/CodeGen/Passes.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Constant.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/GlobalValue.h"
#include "llvm/IR/GlobalVariable.h"
ShadowStackGCLowering();
bool doInitialization(Module &M) override;
+ void getAnalysisUsage(AnalysisUsage &AU) const override;
bool runOnFunction(Function &F) override;
private:
INITIALIZE_PASS_BEGIN(ShadowStackGCLowering, DEBUG_TYPE,
"Shadow Stack GC Lowering", false, false)
INITIALIZE_PASS_DEPENDENCY(GCModuleInfo)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
INITIALIZE_PASS_END(ShadowStackGCLowering, DEBUG_TYPE,
"Shadow Stack GC Lowering", false, false)
return dyn_cast<GetElementPtrInst>(Val);
}
+void ShadowStackGCLowering::getAnalysisUsage(AnalysisUsage &AU) const {
+ AU.addPreserved<DominatorTreeWrapperPass>();
+}
+
/// runOnFunction - Insert code to maintain the shadow stack.
bool ShadowStackGCLowering::runOnFunction(Function &F) {
// Quick exit for functions that do not use the shadow stack GC.
if (Roots.empty())
return false;
+ Optional<DomTreeUpdater> DTU;
+ if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
+ DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);
+
// Build the constant map and figure the type of the shadow stack entry.
Value *FrameMap = GetFrameMap(F);
Type *ConcreteStackEntryTy = GetConcreteStackEntryType(F);
AtEntry.CreateStore(NewHeadVal, Head);
// For each instruction that escapes...
- EscapeEnumerator EE(F, "gc_cleanup");
+ EscapeEnumerator EE(F, "gc_cleanup", /*HandleExceptions=*/true,
+ DTU.hasValue() ? DTU.getPointer() : nullptr);
while (IRBuilder<> *AtExit = EE.Next()) {
// Pop the entry from the shadow stack. Don't reuse CurrentHead from
// AtEntry, since that would make the value live for the entire function.
SmallVector<Value *, 16> Args;
for (unsigned I = Calls.size(); I != 0;) {
CallInst *CI = cast<CallInst>(Calls[--I]);
- changeToInvokeAndSplitBasicBlock(CI, CleanupBB);
+ changeToInvokeAndSplitBasicBlock(CI, CleanupBB, DTU);
}
Builder.SetInsertPoint(RI);
}
BasicBlock *llvm::changeToInvokeAndSplitBasicBlock(CallInst *CI,
- BasicBlock *UnwindEdge) {
+ BasicBlock *UnwindEdge,
+ DomTreeUpdater *DTU) {
BasicBlock *BB = CI->getParent();
// Convert this function call into an invoke instruction. First, split the
// basic block.
- BasicBlock *Split =
- BB->splitBasicBlock(CI->getIterator(), CI->getName() + ".noexc");
+ BasicBlock *Split = SplitBlock(BB, CI, DTU, /*LI=*/nullptr, /*MSSAU*/ nullptr,
+ CI->getName() + ".noexc");
- // Delete the unconditional branch inserted by splitBasicBlock
+ // Delete the unconditional branch inserted by SplitBlock
BB->getInstList().pop_back();
// Create the new invoke instruction.
II->setCallingConv(CI->getCallingConv());
II->setAttributes(CI->getAttributes());
+ if (DTU)
+ DTU->applyUpdates({{DominatorTree::Insert, BB, UnwindEdge}});
+
// Make sure that anything using the call now uses the invoke! This also
// updates the CallGraph if present, because it uses a WeakTrackingVH.
CI->replaceAllUsesWith(II);