diff --git a/llvm/include/llvm/CodeGen/StackProtector.h b/llvm/include/llvm/CodeGen/StackProtector.h index 1c79218a6062..89992d5d3892 100644 --- a/llvm/include/llvm/CodeGen/StackProtector.h +++ b/llvm/include/llvm/CodeGen/StackProtector.h @@ -119,6 +119,7 @@ public: } SSPLayoutKind getSSPLayout(const AllocaInst *AI) const; + void adjustForColoring(const AllocaInst *From, const AllocaInst *To); virtual bool runOnFunction(Function &Fn); }; diff --git a/llvm/lib/CodeGen/StackColoring.cpp b/llvm/lib/CodeGen/StackColoring.cpp index 83def79b2fd3..4932644ea0e0 100644 --- a/llvm/lib/CodeGen/StackColoring.cpp +++ b/llvm/lib/CodeGen/StackColoring.cpp @@ -43,6 +43,7 @@ #include "llvm/CodeGen/MachineRegisterInfo.h" #include "llvm/CodeGen/PseudoSourceValue.h" #include "llvm/CodeGen/SlotIndexes.h" +#include "llvm/CodeGen/StackProtector.h" #include "llvm/DebugInfo.h" #include "llvm/IR/Dominators.h" #include "llvm/IR/Function.h" @@ -117,6 +118,8 @@ class StackColoring : public MachineFunctionPass { VNInfo::Allocator VNInfoAllocator; /// SlotIndex analysis object. SlotIndexes *Indexes; + /// The stack protector object. + StackProtector *SP; /// The list of lifetime markers found. These markers are to be removed /// once the coloring is done. @@ -191,6 +194,7 @@ INITIALIZE_PASS_BEGIN(StackColoring, "stack-coloring", "Merge disjoint stack slots", false, false) INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree) INITIALIZE_PASS_DEPENDENCY(SlotIndexes) +INITIALIZE_PASS_DEPENDENCY(StackProtector) INITIALIZE_PASS_END(StackColoring, "stack-coloring", "Merge disjoint stack slots", false, false) @@ -198,6 +202,7 @@ void StackColoring::getAnalysisUsage(AnalysisUsage &AU) const { AU.addRequired(); AU.addPreserved(); AU.addRequired(); + AU.addRequired(); MachineFunctionPass::getAnalysisUsage(AU); } @@ -517,6 +522,10 @@ void StackColoring::remapInstructions(DenseMap &SlotRemap) { Inst = Cast; } + // Allow the stack protector to adjust its value map to account for the + // upcoming replacement. + SP->adjustForColoring(From, To); + // Note that this will not replace uses in MMOs (which we'll update below), // or anywhere else (which is why we won't delete the original // instruction). @@ -674,6 +683,7 @@ bool StackColoring::runOnMachineFunction(MachineFunction &Func) { MF = &Func; MFI = MF->getFrameInfo(); Indexes = &getAnalysis(); + SP = &getAnalysis(); BlockLiveness.clear(); BasicBlocks.clear(); BasicBlockNumbering.clear(); diff --git a/llvm/lib/CodeGen/StackProtector.cpp b/llvm/lib/CodeGen/StackProtector.cpp index b350237f9b61..f2925ef7712f 100644 --- a/llvm/lib/CodeGen/StackProtector.cpp +++ b/llvm/lib/CodeGen/StackProtector.cpp @@ -57,6 +57,27 @@ StackProtector::getSSPLayout(const AllocaInst *AI) const { return AI ? Layout.lookup(AI) : SSPLK_None; } +void StackProtector::adjustForColoring(const AllocaInst *From, + const AllocaInst *To) { + // When coloring replaces one alloca with another, transfer the SSPLayoutKind + // tag from the remapped to the target alloca. The remapped alloca should + // have a size smaller than or equal to the replacement alloca. + SSPLayoutMap::iterator I = Layout.find(From); + if (I != Layout.end()) { + SSPLayoutKind Kind = I->second; + Layout.erase(I); + + // Transfer the tag, but make sure that SSPLK_AddrOf does not overwrite + // SSPLK_SmallArray or SSPLK_LargeArray, and make sure that + // SSPLK_SmallArray does not overwrite SSPLK_LargeArray. + I = Layout.find(To); + if (I == Layout.end()) + Layout.insert(std::make_pair(To, Kind)); + else if (I->second != SSPLK_LargeArray && Kind != SSPLK_AddrOf) + I->second = Kind; + } +} + bool StackProtector::runOnFunction(Function &Fn) { F = &Fn; M = F->getParent();