[SafeStack] Don't create SCEV min between pointer and integer (PR54784)

[nikic]

Summary

Rather than rewriting the alloca pointer to zero, use removePointerBase() to drop the base pointer. This will simply bail if the base pointer is not the alloca. We could try doing something more fancy here (like dropping the sources not based on the alloca on the premise that they aren't SafeStack-relevant or something), but I don't think that's worthwhile.

Fixes https://github.com/llvm/llvm-project/issues/54784.

[efriedma]

LGTM

In some cases, we could prove that the access is either in-range, or not based on the alloca. (For example, given an alloca "a" and an access "min(a,b)+1", the pointer is either "a+1", or "b+1". "a+1" is in range, and "b" is unrelated to the alloca, so the access is safe.) But it seems unlikely to be important.

llvm/lib/CodeGen/SafeStack.cpp

SafeStackUsePointerAddress("safestack-use-pointer-address",
                                   cl::init(false), cl::Hidden);
 
 static cl::opt<bool> ClColoring("safe-stack-coloring",
                                 cl::desc("enable safe stack coloring"),
                                 cl::Hidden, cl::init(true));
 
 namespace {
 
-/// Rewrite an SCEV expression for a memory access address to an expression that
-/// represents offset from the given alloca.
-///
-/// The implementation simply replaces all mentions of the alloca with zero.
-class AllocaOffsetRewriter : public SCEVRewriteVisitor<AllocaOffsetRewriter> {
-  const Value *AllocaPtr;
-
-public:
-  AllocaOffsetRewriter(ScalarEvolution &SE, const Value *AllocaPtr)
-      : SCEVRewriteVisitor(SE), AllocaPtr(AllocaPtr) {}
-
-  const SCEV *visitUnknown(const SCEVUnknown *Expr) {
-    if (Expr->getValue() == AllocaPtr)
-      return SE.getZero(Expr->getType());
-    return Expr;
-  }
-};
-
 /// The SafeStack pass splits the stack of each function into the safe
 /// stack, which is only accessed through memory safe dereferences (as
 /// determined statically), and the unsafe stack, which contains all
 /// local variables that are accessed in ways that we can't prove to
 /// be safe.
 class SafeStack {
   Function &F;
   const TargetLoweringBase &TL;
if (!C)
       return 0;
     Size *= C->getZExtValue();
   }
   return Size;
 }
 
 bool SafeStack::IsAccessSafe(Value *Addr, uint64_t AccessSize,
                              const Value *AllocaPtr, uint64_t AllocaSize) {
-  AllocaOffsetRewriter Rewriter(SE, AllocaPtr);
-  const SCEV *Expr = Rewriter.visit(SE.getSCEV(Addr));
+  const SCEV *AddrExpr = SE.getSCEV(Addr);
+  const auto *Base = dyn_cast<SCEVUnknown>(SE.getPointerBase(AddrExpr));
+  if (!Base || Base->getValue() != AllocaPtr) {
+    LLVM_DEBUG(
+        dbgs() << "[SafeStack] "
+               << (isa<AllocaInst>(AllocaPtr) ? "Alloca " : "ByValArgument ")
+               << *AllocaPtr << "\n"
+               << "SCEV " << *AddrExpr << " not directly based on alloca\n");
+    return false;
+  }
 
+  const SCEV *Expr = SE.removePointerBase(AddrExpr);
   uint64_t BitWidth = SE.getTypeSizeInBits(Expr->getType());
   ConstantRange AccessStartRange = SE.getUnsignedRange(Expr);
   ConstantRange SizeRange =
       ConstantRange(APInt(BitWidth, 0), APInt(BitWidth, AccessSize));
   ConstantRange AccessRange = AccessStartRange.add(SizeRange);
   ConstantRange AllocaRange =
       ConstantRange(APInt(BitWidth, 0), APInt(BitWidth, AllocaSize));
   bool Safe = AllocaRange.contains(AccessRange);

llvm/test/Transforms/SafeStack/pr54784.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -safe-stack < %s | FileCheck %s
+
+target triple = "x86_64-unknown-unknown"
+
+define void @min_of_pointers(i32* %p) safestack {
+; CHECK-LABEL: @min_of_pointers(
+; CHECK-NEXT:    [[UNSAFE_STACK_PTR:%.*]] = load i8*, i8** @__safestack_unsafe_stack_ptr, align 8
+; CHECK-NEXT:    [[UNSAFE_STACK_STATIC_TOP:%.*]] = getelementptr i8, i8* [[UNSAFE_STACK_PTR]], i32 -16
+; CHECK-NEXT:    store i8* [[UNSAFE_STACK_STATIC_TOP]], i8** @__safestack_unsafe_stack_ptr, align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i8, i8* [[UNSAFE_STACK_PTR]], i32 -4
+; CHECK-NEXT:    [[A_UNSAFE1:%.*]] = bitcast i8* [[TMP1]] to i32*
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32* [[A_UNSAFE1]], [[P:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr i8, i8* [[UNSAFE_STACK_PTR]], i32 -4
+; CHECK-NEXT:    [[A_UNSAFE:%.*]] = bitcast i8* [[TMP2]] to i32*
+; CHECK-NEXT:    [[S:%.*]] = select i1 [[CMP]], i32* [[A_UNSAFE]], i32* [[P]]
+; CHECK-NEXT:    store i32 0, i32* [[S]], align 4
+; CHECK-NEXT:    store i8* [[UNSAFE_STACK_PTR]], i8** @__safestack_unsafe_stack_ptr, align 8
+; CHECK-NEXT:    ret void
+;
+  %a = alloca i32
+  %cmp = icmp ult i32* %a, %p
+  %s = select i1 %cmp, i32* %a, i32* %p
+  store i32 0, i32* %s
+  ret void
+}