From 5d964e262f0c446dd1e25179d342a953a3611d76 Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka@google.com>
Date: Mon, 22 Jun 2020 01:28:25 -0700
Subject: [PATCH] [StackSafety] Check variable lifetime

We can't consider variable safe if out-of-lifetime access is possible.
So if StackLifetime can't prove that the instruction always uses
the variable when it's still alive, we consider it unsafe.
---
 llvm/lib/Analysis/StackSafetyAnalysis.cpp       | 14 ++++
 llvm/test/Analysis/StackSafetyAnalysis/local.ll | 87 +++++++++++++++++++++++--
 2 files changed, 97 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Analysis/StackSafetyAnalysis.cpp b/llvm/lib/Analysis/StackSafetyAnalysis.cpp
index 88aa5ce..8ec69fc 100644
--- a/llvm/lib/Analysis/StackSafetyAnalysis.cpp
+++ b/llvm/lib/Analysis/StackSafetyAnalysis.cpp
@@ -297,6 +297,7 @@ bool StackSafetyLocalAnalysis::analyzeAllUses(Value *Ptr,
   SmallPtrSet<const Value *, 16> Visited;
   SmallVector<const Value *, 8> WorkList;
   WorkList.push_back(Ptr);
+  const AllocaInst *AI = dyn_cast<AllocaInst>(Ptr);
 
   // A DFS search through all uses of the alloca in bitcasts/PHI/GEPs/etc.
   while (!WorkList.empty()) {
@@ -310,6 +311,10 @@ bool StackSafetyLocalAnalysis::analyzeAllUses(Value *Ptr,
 
       switch (I->getOpcode()) {
       case Instruction::Load: {
+        if (AI && !SL.isAliveAfter(AI, I)) {
+          US.updateRange(UnknownRange);
+          return false;
+        }
         US.updateRange(
             getAccessRange(UI, Ptr, DL.getTypeStoreSize(I->getType())));
         break;
@@ -324,6 +329,10 @@ bool StackSafetyLocalAnalysis::analyzeAllUses(Value *Ptr,
           US.updateRange(UnknownRange);
           return false;
         }
+        if (AI && !SL.isAliveAfter(AI, I)) {
+          US.updateRange(UnknownRange);
+          return false;
+        }
         US.updateRange(getAccessRange(
             UI, Ptr, DL.getTypeStoreSize(I->getOperand(0)->getType())));
         break;
@@ -341,6 +350,11 @@ bool StackSafetyLocalAnalysis::analyzeAllUses(Value *Ptr,
         if (I->isLifetimeStartOrEnd())
           break;
 
+        if (AI && !SL.isAliveAfter(AI, I)) {
+          US.updateRange(UnknownRange);
+          return false;
+        }
+
         if (const MemIntrinsic *MI = dyn_cast<MemIntrinsic>(I)) {
           US.updateRange(getMemIntrinsicAccessRange(MI, UI, Ptr));
           break;
diff --git a/llvm/test/Analysis/StackSafetyAnalysis/local.ll b/llvm/test/Analysis/StackSafetyAnalysis/local.ll
index ad4745c..890d4dc 100644
--- a/llvm/test/Analysis/StackSafetyAnalysis/local.ll
+++ b/llvm/test/Analysis/StackSafetyAnalysis/local.ll
@@ -476,8 +476,8 @@ define void @Overflow() {
 ; CHECK-LABEL: @Overflow dso_preemptable{{$}}
 ; CHECK-NEXT: args uses:
 ; CHECK-NEXT: allocas uses:
-; LOCAL: x[1]: empty-set, @LoadMinInt64(arg0, [-9223372036854775808,-9223372036854775807)){{$}}
-; GLOBAL: x[1]: full-set, @LoadMinInt64(arg0, [-9223372036854775808,-9223372036854775807)){{$}}
+; LOCAL-NEXT: x[1]: empty-set, @LoadMinInt64(arg0, [-9223372036854775808,-9223372036854775807)){{$}}
+; GLOBAL-NEXT: x[1]: full-set, @LoadMinInt64(arg0, [-9223372036854775808,-9223372036854775807)){{$}}
 ; CHECK-NOT: ]:
 entry:
   %x = alloca i8, align 4
@@ -491,7 +491,7 @@ define void @DeadBlock(i64* %p) {
 ; CHECK-NEXT: args uses:
 ; CHECK-NEXT: p[]: empty-set{{$}}
 ; CHECK-NEXT: allocas uses:
-; CHECK: x[1]: empty-set{{$}}
+; CHECK-NEXT: x[1]: empty-set{{$}}
 ; CHECK-NOT: ]:
 entry:
   %x = alloca i8, align 4
@@ -504,4 +504,83 @@ dead:
 
 end:
   ret void
-}
\ No newline at end of file
+}
+
+define void @LifeNotStarted() {
+; CHECK-LABEL: @LifeNotStarted dso_preemptable{{$}}
+; CHECK-NEXT: args uses:
+; CHECK-NEXT: allocas uses:
+; CHECK: x[1]: full-set{{$}}
+; CHECK: y[1]: full-set{{$}}
+; CHECK: z[1]: full-set{{$}}
+; CHECK-NOT: ]:
+entry:
+  %x = alloca i8, align 4
+  %y = alloca i8, align 4
+  %z = alloca i8, align 4
+
+  store i8 5, i8* %x
+  %n = load i8, i8* %y
+  call void @llvm.memset.p0i8.i32(i8* nonnull %z, i8 0, i32 1, i1 false)
+
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %x)
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %y)
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %z)
+
+  ret void
+}
+
+define void @LifeOK() {
+; CHECK-LABEL: @LifeOK dso_preemptable{{$}}
+; CHECK-NEXT: args uses:
+; CHECK-NEXT: allocas uses:
+; CHECK: x[1]: [0,1){{$}}
+; CHECK: y[1]: [0,1){{$}}
+; CHECK: z[1]: [0,1){{$}}
+; CHECK-NOT: ]:
+entry:
+  %x = alloca i8, align 4
+  %y = alloca i8, align 4
+  %z = alloca i8, align 4
+
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %x)
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %y)
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %z)
+
+  store i8 5, i8* %x
+  %n = load i8, i8* %y
+  call void @llvm.memset.p0i8.i32(i8* nonnull %z, i8 0, i32 1, i1 false)
+
+  ret void
+}
+
+define void @LifeEnded() {
+; CHECK-LABEL: @LifeEnded dso_preemptable{{$}}
+; CHECK-NEXT: args uses:
+; CHECK-NEXT: allocas uses:
+; CHECK: x[1]: full-set{{$}}
+; CHECK: y[1]: full-set{{$}}
+; CHECK: z[1]: full-set{{$}}
+; CHECK-NOT: ]:
+entry:
+  %x = alloca i8, align 4
+  %y = alloca i8, align 4
+  %z = alloca i8, align 4
+
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %x)
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %y)
+  call void @llvm.lifetime.start.p0i8(i64 1, i8* %z)
+
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %x)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %y)
+  call void @llvm.lifetime.end.p0i8(i64 1, i8* %z)
+
+  store i8 5, i8* %x
+  %n = load i8, i8* %y
+  call void @llvm.memset.p0i8.i32(i8* nonnull %z, i8 0, i32 1, i1 false)
+
+  ret void
+}
+
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture)
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture)
\ No newline at end of file
-- 
2.7.4