[safestack] Fix stack guard live range.

Stack guard slot is live throughout the function.

llvm-svn: 276712

3 files changed, 38 insertions, 2 deletions

diff --git a/llvm/lib/CodeGen/SafeStack.cpp b/llvm/lib/CodeGen/SafeStack.cpp
index 19cd59b9dba..4a1b9958a5b 100644
--- a/llvm/lib/CodeGen/SafeStack.cpp
+++ b/llvm/lib/CodeGen/SafeStack.cpp
@@ -530,7 +530,7 @@ Value *SafeStack::moveStaticAllocasToUnsafeStack(
     unsigned Align =
         std::max(DL->getPrefTypeAlignment(Ty), StackGuardSlot->getAlignment());
     SSL.addObject(StackGuardSlot, getStaticAllocaAllocationSize(StackGuardSlot),
-                  Align, SSC.getLiveRange(StackGuardSlot));
+                  Align, SSC.getFullLiveRange());
   }
 
   for (Argument *Arg : ByValArguments) {
diff --git a/llvm/lib/CodeGen/SafeStackColoring.cpp b/llvm/lib/CodeGen/SafeStackColoring.cpp
index 709614f57e7..795eb8d2719 100644
--- a/llvm/lib/CodeGen/SafeStackColoring.cpp
+++ b/llvm/lib/CodeGen/SafeStackColoring.cpp
@@ -25,7 +25,9 @@ static cl::opt<bool> ClColoring("safe-stack-coloring",
                                 cl::Hidden, cl::init(true));
 
 const StackColoring::LiveRange &StackColoring::getLiveRange(AllocaInst *AI) {
-  return LiveRanges[AllocaNumbering[AI]];
+  const auto IT = AllocaNumbering.find(AI);
+  assert(IT != AllocaNumbering.end());
+  return LiveRanges[IT->second];
 }
 
 bool StackColoring::readMarker(Instruction *I, bool *IsStart) {
diff --git a/llvm/test/Transforms/SafeStack/coloring-ssp.ll b/llvm/test/Transforms/SafeStack/coloring-ssp.ll
new file mode 100644
index 00000000000..d71babe200d
--- /dev/null
+++ b/llvm/test/Transforms/SafeStack/coloring-ssp.ll
@@ -0,0 +1,34 @@
+; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
+
+; %x and %y share a stack slot between them, but not with the stack guard.
+define void @f() safestack sspreq {
+; CHECK-LABEL: define void @f
+entry:
+; CHECK:  %[[USP:.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr
+; CHECK:   getelementptr i8, i8* %[[USP]], i32 -16
+
+; CHECK:  %[[A:.*]] = getelementptr i8, i8* %[[USP]], i32 -8
+; CHECK:  %[[StackGuardSlot:.*]] = bitcast i8* %[[A]] to i8**
+; CHECK:  store i8* %{{.*}}, i8** %[[StackGuardSlot]]
+
+  %x = alloca i64, align 8
+  %y = alloca i64, align 8
+  %x0 = bitcast i64* %x to i8*
+  %y0 = bitcast i64* %y to i8*
+
+  call void @llvm.lifetime.start(i64 -1, i8* %x0)
+; CHECK:  getelementptr i8, i8* %[[USP]], i32 -16
+  call void @capture64(i64* %x)
+  call void @llvm.lifetime.end(i64 -1, i8* %x0)
+
+  call void @llvm.lifetime.start(i64 -1, i8* %y0)
+; CHECK:  getelementptr i8, i8* %[[USP]], i32 -16
+  call void @capture64(i64* %y)
+  call void @llvm.lifetime.end(i64 -1, i8* %y0)
+
+  ret void
+}
+
+declare void @llvm.lifetime.start(i64, i8* nocapture)
+declare void @llvm.lifetime.end(i64, i8* nocapture)
+declare void @capture64(i64*)