Revert "[StackSafety] Skip ambiguous lifetime analysis"

This reverts commit 0b2616a.

Crashes with safe-stack.

Author: vitalybuka

llvm/include/llvm/Analysis/StackLifetime.h

@@ -121,8 +121,6 @@ class StackLifetime {
   DenseMap<const BasicBlock *, SmallVector<std::pair<unsigned, Marker>, 4>>
       BBMarkers;
 
-  bool HasUnknownLifetimeStartOrEnd = false;
-
   void dumpAllocas() const;
   void dumpBlockLiveness() const;
   void dumpLiveRanges() const;

llvm/lib/Analysis/StackLifetime.cpp

@@ -11,7 +11,6 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
-#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/AssemblyAnnotationWriter.h"
 #include "llvm/IR/BasicBlock.h"
@@ -64,27 +63,42 @@ bool StackLifetime::isAliveAfter(const AllocaInst *AI,
   return getLiveRange(AI).test(InstNum);
 }
 
+static bool readMarker(const Instruction *I, bool *IsStart) {
+  if (!I->isLifetimeStartOrEnd())
+    return false;
+
+  auto *II = cast<IntrinsicInst>(I);
+  *IsStart = II->getIntrinsicID() == Intrinsic::lifetime_start;
+  return true;
+}
+
 void StackLifetime::collectMarkers() {
   InterestingAllocas.resize(NumAllocas);
   DenseMap<const BasicBlock *, SmallDenseMap<const IntrinsicInst *, Marker>>
       BBMarkerSet;
 
   // Compute the set of start/end markers per basic block.
-  for (const BasicBlock *BB : depth_first(&F)) {
-    for (const Instruction &I : *BB) {
-      const IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
-      if (!II || !II->isLifetimeStartOrEnd())
-        continue;
-      const AllocaInst *AI = llvm::findAllocaForValue(II->getArgOperand(1));
-      if (!AI) {
-        HasUnknownLifetimeStartOrEnd = true;
-        continue;
+  for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
+    const AllocaInst *AI = Allocas[AllocaNo];
+    SmallVector<const Instruction *, 8> WorkList;
+    WorkList.push_back(AI);
+    while (!WorkList.empty()) {
+      const Instruction *I = WorkList.pop_back_val();
+      for (const User *U : I->users()) {
+        if (auto *BI = dyn_cast<BitCastInst>(U)) {
+          WorkList.push_back(BI);
+          continue;
+        }
+        auto *UI = dyn_cast<IntrinsicInst>(U);
+        if (!UI)
+          continue;
+        bool IsStart;
+        if (!readMarker(UI, &IsStart))
+          continue;
+        if (IsStart)
+          InterestingAllocas.set(AllocaNo);
+        BBMarkerSet[UI->getParent()][UI] = {AllocaNo, IsStart};
       }
-      bool IsStart = II->getIntrinsicID() == Intrinsic::lifetime_start;
-      unsigned AllocaNo = AllocaNumbering[AI];
-      if (IsStart)
-        InterestingAllocas.set(AllocaNo);
-      BBMarkerSet[BB][II] = {AllocaNo, IsStart};
     }
   }
 
@@ -290,20 +304,6 @@ StackLifetime::StackLifetime(const Function &F,
 }
 
 void StackLifetime::run() {
-  if (HasUnknownLifetimeStartOrEnd) {
-    // There is marker which we can't assign to a specific alloca, so we
-    // fallback to the most conservative results for the type.
-    switch (Type) {
-    case LivenessType::May:
-      LiveRanges.resize(NumAllocas, getFullLiveRange());
-      break;
-    case LivenessType::Must:
-      LiveRanges.resize(NumAllocas, LiveRange(Instructions.size()));
-      break;
-    }
-    return;
-  }
-
   LiveRanges.resize(NumAllocas, LiveRange(Instructions.size()));
   for (unsigned I = 0; I < NumAllocas; ++I)
     if (!InterestingAllocas.test(I))

llvm/test/Analysis/StackSafetyAnalysis/lifetime.ll

@@ -742,7 +742,7 @@ if.end:
 ; MAY-NEXT: Alive: <x y>
 ; MUST-NEXT: Alive: <y>
 
-  ret void
+ret void
 }
 
 define void @unreachable() {
@@ -778,62 +778,7 @@ end:
 ; CHECK: end:
 ; CHECK-NEXT: Alive: <x y>
 
-  ret void
-}
-
-define void @non_alloca(i8* %p) {
-; CHECK-LABEL: define void @non_alloca
-entry:
-; CHECK: entry:
-; MAY-NEXT: Alive: <x y>
-; MUST-NEXT: Alive: <>
-  %x = alloca i8, align 4
-  %y = alloca i8, align 4
-
-  call void @llvm.lifetime.start.p0i8(i64 4, i8* %p)
-; CHECK: call void @llvm.lifetime.start.p0i8(i64 4, i8* %p)
-; MAY-NEXT: Alive: <x y>
-; MUST-NEXT: Alive: <>
-
-  call void @llvm.lifetime.start.p0i8(i64 4, i8* %x)
-; CHECK: call void @llvm.lifetime.start.p0i8(i64 4, i8* %x)
-; MAY-NEXT: Alive: <x y>
-; MUST-NEXT: Alive: <>
-
-  call void @llvm.lifetime.end.p0i8(i64 4, i8* %p)
-; CHECK: call void @llvm.lifetime.end.p0i8(i64 4, i8* %p)
-; MAY-NEXT: Alive: <x y>
-; MUST-NEXT: Alive: <>
-
-  ret void
-}
-
-define void @select_alloca(i1 %v) {
-; CHECK-LABEL: define void @select_alloca
-entry:
-; CHECK: entry:
-; MAY-NEXT: Alive: <x y>
-; MUST-NEXT: Alive: <>
-  %x = alloca i8, align 4
-  %y = alloca i8, align 4
-  %cxcy = select i1 %v, i8* %x, i8* %y
-
-  call void @llvm.lifetime.start.p0i8(i64 1, i8* %cxcy)
-; CHECK: call void @llvm.lifetime.start.p0i8(i64 1, i8* %cxcy)
-; MAY-NEXT: Alive: <x y>
-; MUST-NEXT: Alive: <>
-
-  call void @llvm.lifetime.start.p0i8(i64 1, i8* %x)
-; CHECK: call void @llvm.lifetime.start.p0i8(i64 1, i8* %x)
-; MAY-NEXT: Alive: <x y>
-; MUST-NEXT: Alive: <>
-
-  call void @llvm.lifetime.end.p0i8(i64 1, i8* %x)
-; CHECK: call void @llvm.lifetime.end.p0i8(i64 1, i8* %x)
-; MAY-NEXT: Alive: <x y>
-; MUST-NEXT: Alive: <>
-
-  ret void
+ret void
 }
 
 declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture)

llvm/test/CodeGen/AArch64/stack-tagging.ll

@@ -192,17 +192,18 @@ another_bb:
 ; CHECK: alloca { i32, [12 x i8] }, align 16
 ; CHECK: call { i32, [12 x i8] }* @llvm.aarch64.tagp
 ; CHECK: call void @llvm.aarch64.settag(
-; CHECK: alloca { i32, [12 x i8] }, align 16
-; CHECK: call { i32, [12 x i8] }* @llvm.aarch64.tagp
-; CHECK: call void @llvm.aarch64.settag(
+; SSI: alloca i32, align 4
+; NOSSI: alloca { i32, [12 x i8] }, align 16
+; NOSSI: call { i32, [12 x i8] }* @llvm.aarch64.tagp
+; NOSSI: call void @llvm.aarch64.settag(
 ; CHECK: store i32
 ; CHECK: call void @noUse32(i32*
 ; CHECK: store i32
 ; CHECK: store i32
 ; CHECK: call void @noUse32(i32*
 ; CHECK: call void @llvm.aarch64.settag(
 ; CHECK: call void @llvm.aarch64.settag(
-; CHECK: call void @llvm.aarch64.settag(
+; NOSSI: call void @llvm.aarch64.settag(
 ; CHECK: ret void
 
-!0 = !{}
+!0 = !{}