diff --git a/llvm/lib/CodeGen/SafeStack.cpp b/llvm/lib/CodeGen/SafeStack.cpp index 7886bf341303d..63c4e94b49359 100644 --- a/llvm/lib/CodeGen/SafeStack.cpp +++ b/llvm/lib/CodeGen/SafeStack.cpp @@ -501,7 +501,14 @@ Value *SafeStack::moveStaticAllocasToUnsafeStack( static const StackColoring::LiveRange NoColoringRange(1, true); if (ClColoring) SSC.run(); - SSC.removeAllMarkers(); + + for (auto *I : SSC.getMarkers()) { + auto *Op = dyn_cast(I->getOperand(1)); + const_cast(I)->eraseFromParent(); + // Remove the operand bitcast, too, if it has no more uses left. + if (Op && Op->use_empty()) + Op->eraseFromParent(); + } // Unsafe stack always grows down. StackLayout SSL(StackAlignment); diff --git a/llvm/lib/CodeGen/SafeStackColoring.cpp b/llvm/lib/CodeGen/SafeStackColoring.cpp index 7009b7cf68b06..ad6d9bde3deab 100644 --- a/llvm/lib/CodeGen/SafeStackColoring.cpp +++ b/llvm/lib/CodeGen/SafeStackColoring.cpp @@ -42,14 +42,12 @@ static bool readMarker(const Instruction *I, bool *IsStart) { return true; } -void StackColoring::removeAllMarkers() { - for (auto *I : Markers) { - auto *Op = dyn_cast(I->getOperand(1)); - const_cast(I)->eraseFromParent(); - // Remove the operand bitcast, too, if it has no more uses left. - if (Op && Op->use_empty()) - Op->eraseFromParent(); - } +std::vector StackColoring::getMarkers() const { + std::vector Markers; + for (auto &M : InstructionNumbering) + if (M.getFirst()->isLifetimeStartOrEnd()) + Markers.push_back(M.getFirst()); + return Markers; } void StackColoring::collectMarkers() { @@ -78,7 +76,6 @@ void StackColoring::collectMarkers() { if (IsStart) InterestingAllocas.set(AllocaNo); BBMarkerSet[UI->getParent()][UI] = {AllocaNo, IsStart}; - Markers.push_back(UI); } } } diff --git a/llvm/lib/CodeGen/SafeStackColoring.h b/llvm/lib/CodeGen/SafeStackColoring.h index b58e9e47c07a1..b71b374182c00 100644 --- a/llvm/lib/CodeGen/SafeStackColoring.h +++ b/llvm/lib/CodeGen/SafeStackColoring.h @@ -102,7 +102,6 @@ class StackColoring { /// The set of allocas that have at least one lifetime.start. All other /// allocas get LiveRange that corresponds to the entire function. BitVector InterestingAllocas; - SmallVector Markers; struct Marker { unsigned AllocaNo; @@ -125,7 +124,7 @@ class StackColoring { StackColoring(const Function &F, ArrayRef Allocas); void run(); - void removeAllMarkers(); + std::vector getMarkers() const; /// Returns a set of "interesting" instructions where the given alloca is /// live. Not all instructions in a function are interesting: we pick a set