[InstCombine] Use CaptureTracking in foldAllocaCmp()

foldAllocaCmp() checks whether the alloca is not captured (ignoring
the icmp). Replace the manual implementation of escape analysis
with CaptureTracking.

The primary practical difference is that CaptureTracking handles
nocapture arguments, while foldAllocaCmp() was using a hardcoded
list.

This is basically just the CaptureTracking refactoring from D120371
without the other changes.

Author: nikic

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

@@ -14,6 +14,7 @@
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/Analysis/CaptureTracking.h"
 #include "llvm/Analysis/CmpInstAnalysis.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/InstructionSimplify.h"
@@ -915,57 +916,27 @@ Instruction *InstCombinerImpl::foldAllocaCmp(ICmpInst &ICI,
   // comparisons against the alloca consistently, and avoids the risk of
   // erroneously folding a comparison of the pointer with itself.
 
-  unsigned MaxIter = 32; // Break cycles and bound to constant-time.
+  struct CmpCaptureTracker : public CaptureTracker {
+    bool Captured = false;
+    unsigned NumCmps = 0;
 
-  SmallVector<const Use *, 32> Worklist;
-  for (const Use &U : Alloca->uses()) {
-    if (Worklist.size() >= MaxIter)
-      return nullptr;
-    Worklist.push_back(&U);
-  }
-
-  unsigned NumCmps = 0;
-  while (!Worklist.empty()) {
-    assert(Worklist.size() <= MaxIter);
-    const Use *U = Worklist.pop_back_val();
-    const Value *V = U->getUser();
-    --MaxIter;
+    void tooManyUses() override { Captured = true; }
 
-    if (isa<BitCastInst>(V) || isa<GetElementPtrInst>(V) || isa<PHINode>(V) ||
-        isa<SelectInst>(V)) {
-      // Track the uses.
-    } else if (isa<LoadInst>(V)) {
-      // Loading from the pointer doesn't escape it.
-      continue;
-    } else if (const auto *SI = dyn_cast<StoreInst>(V)) {
-      // Storing *to* the pointer is fine, but storing the pointer escapes it.
-      if (SI->getValueOperand() == U->get())
-        return nullptr;
-      continue;
-    } else if (isa<ICmpInst>(V)) {
-      if (NumCmps++)
-        return nullptr; // Found more than one cmp.
-      continue;
-    } else if (const auto *Intrin = dyn_cast<IntrinsicInst>(V)) {
-      switch (Intrin->getIntrinsicID()) {
-        // These intrinsics don't escape or compare the pointer. Memset is safe
-        // because we don't allow ptrtoint. Memcpy and memmove are safe because
-        // we don't allow stores, so src cannot point to V.
-        case Intrinsic::lifetime_start: case Intrinsic::lifetime_end:
-        case Intrinsic::memcpy: case Intrinsic::memmove: case Intrinsic::memset:
-          continue;
-        default:
-          return nullptr;
+    bool captured(const Use *U) override {
+      if (isa<ICmpInst>(U->getUser()) && ++NumCmps == 1) {
+        // Ignore one icmp capture.
+        return false;
       }
-    } else {
-      return nullptr;
-    }
-    for (const Use &U : V->uses()) {
-      if (Worklist.size() >= MaxIter)
-        return nullptr;
-      Worklist.push_back(&U);
+
+      Captured = true;
+      return true;
     }
-  }
+  };
+
+  CmpCaptureTracker Tracker;
+  PointerMayBeCaptured(Alloca, &Tracker);
+  if (Tracker.Captured)
+    return nullptr;
 
   auto *Res = ConstantInt::get(ICI.getType(),
                                !CmpInst::isTrueWhenEqual(ICI.getPredicate()));

llvm/test/Transforms/InstCombine/compare-alloca.ll

@@ -193,14 +193,12 @@ define void @neg_consistent_fold2() {
   ret void
 }
 
-define void @neg_consistent_fold3() {
-; CHECK-LABEL: @neg_consistent_fold3(
-; CHECK-NEXT:    [[M1:%.*]] = alloca [4 x i8], align 1
-; CHECK-NEXT:    [[LGP:%.*]] = load ptr, ptr @gp, align 8
+; FIXME: The end result is correct, but the fold happens for the wrong
+; reason (incorrect icmp GlobalValue special case in CaptureTracking).
+define void @consistent_fold3() {
+; CHECK-LABEL: @consistent_fold3(
 ; CHECK-NEXT:    [[RHS2:%.*]] = call ptr @hidden_inttoptr()
-; CHECK-NEXT:    [[CMP1:%.*]] = icmp eq ptr [[M1]], [[LGP]]
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp eq ptr [[M1]], [[RHS2]]
-; CHECK-NEXT:    call void @witness(i1 [[CMP1]], i1 [[CMP2]])
+; CHECK-NEXT:    call void @witness(i1 false, i1 false)
 ; CHECK-NEXT:    ret void
 ;
   %m = alloca i8, i32 4
@@ -231,13 +229,11 @@ define void @neg_consistent_fold4() {
 
 declare void @unknown(ptr)
 
-; TODO: Missing optimization
 define i1 @consistent_nocapture_inttoptr() {
 ; CHECK-LABEL: @consistent_nocapture_inttoptr(
 ; CHECK-NEXT:    [[M1:%.*]] = alloca [4 x i8], align 1
 ; CHECK-NEXT:    call void @unknown(ptr nocapture nonnull [[M1]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[M1]], inttoptr (i64 2048 to ptr)
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %m = alloca i8, i32 4
   call void @unknown(ptr nocapture %m)
@@ -261,14 +257,12 @@ define i1 @consistent_nocapture_offset() {
 }
 
 @gp = global ptr null, align 8
-; TODO: Missing optimization
+
 define i1 @consistent_nocapture_through_global() {
 ; CHECK-LABEL: @consistent_nocapture_through_global(
 ; CHECK-NEXT:    [[M1:%.*]] = alloca [4 x i8], align 1
 ; CHECK-NEXT:    call void @unknown(ptr nocapture nonnull [[M1]])
-; CHECK-NEXT:    [[LGP:%.*]] = load ptr, ptr @gp, align 8, !nonnull !0
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq ptr [[M1]], [[LGP]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %m = alloca i8, i32 4
   call void @unknown(ptr nocapture %m)