[llvm] [CaptureTracking] Supports analysis for dervied pointers (PR #132744)

[via llvm-commits]

https://github.com/Camsyn created https://github.com/llvm/llvm-project/pull/132744

Fixes issue #132739.

This fixes a missed capture detection scenario where only the base pointer is captured, as follows:
```cpp
int arr[2] = {0, 0};
foo(arr); // arr is captured
arr[1] = 1; // arr + 1 is mis-analyzed as non-capture
```

Previously, PointerMayBeCaptured only performed forward def-use analysis on the given pointer, which worked for base pointers (Arguments/Allocas/Calls) but failed to detect captures of derived pointers when their base pointers were captured.
```cpp
// Forward Tracing
return analyzeCaptures(Ptr);
```

The fix: For the input pointer, first trace back to its base pointer using `getUnderlyingObject`, then perform capture analysis on the base. This ensures proper handling of derived pointers with bases captured.
```cpp
// Backward Tracing
const Value *Obj = getUnderlyingObject(Ptr);
// Forward Tracing
return analyzeCaptures(Obj);
```

Performance considerations:
- Most existing callers already pass base pointers (the common case), so the added backtracing has negligible overhead
- The few callers (ThreadSanitizer.cpp/SanitizerBinaryMetadata.cpp) passing arbitrary pointers are sanitizer-related components where compilation time is less critical compared to runtime overhead

This addresses false negatives in capture detection while maintaining reasonable compilation efficiency.

>From 87c0937b9cc29c312a379a72c07aea84a8dcbdf0 Mon Sep 17 00:00:00 2001
From: Camsyn <camsyn at foxmail.com>
Date: Mon, 24 Mar 2025 15:07:04 +0800
Subject: [PATCH] [CaptureTracking] Supports analysis for dervied pointers

Fixes issue #132739.

This fixes a missed capture detection scenario where only the base
pointer is captured.

Previously, PointerMayBeCaptured only performed forward def-use
analysis on the given pointer, which worked for base pointers
(Arguments/Allocas/Calls) but failed to detect captures of derived
pointers when their base pointers were captured.

The fix: For the input pointer, first trace back to its base pointer
using `getUnderlyingObject`, then perform capture analysis on the base.
This ensures proper handling of derived pointers with bases captured.

Performance considerations:
- Most existing callers already pass base pointers (the common case), so
the added backtracing has negligible overhead
- The few callers (ThreadSanitizer.cpp/SanitizerBinaryMetadata.cpp)
passing arbitrary pointers are sanitizer-related components where
compilation time is less critical compared to runtime overhead

This addresses false negatives in capture detection while maintaining
 reasonable compilation efficiency.
---
 llvm/lib/Analysis/CaptureTracking.cpp         |  4 +++
 .../Analysis/CaptureTrackingTest.cpp          | 28 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/llvm/lib/Analysis/CaptureTracking.cpp b/llvm/lib/Analysis/CaptureTracking.cpp
index aa53a27537567..0d598cd245464 100644
--- a/llvm/lib/Analysis/CaptureTracking.cpp
+++ b/llvm/lib/Analysis/CaptureTracking.cpp
@@ -424,6 +424,10 @@ void llvm::PointerMayBeCaptured(const Value *V, CaptureTracker *Tracker,
   if (MaxUsesToExplore == 0)
     MaxUsesToExplore = DefaultMaxUsesToExplore;
 
+  // When analyzing a derived pointer, we need to analyze its underlying
+  // object to determine whether it is captured.
+  // E.g., `ptr + 1` is captured if `ptr` is captured.
+  V = getUnderlyingObjectAggressive(V);
   SmallVector<const Use *, 20> Worklist;
   Worklist.reserve(getDefaultMaxUsesToExploreForCaptureTracking());
   SmallSet<const Use *, 20> Visited;
diff --git a/llvm/unittests/Analysis/CaptureTrackingTest.cpp b/llvm/unittests/Analysis/CaptureTrackingTest.cpp
index ea3f21efc014c..f35fc0a151470 100644
--- a/llvm/unittests/Analysis/CaptureTrackingTest.cpp
+++ b/llvm/unittests/Analysis/CaptureTrackingTest.cpp
@@ -132,3 +132,31 @@ TEST(CaptureTracking, MultipleUsesInSameInstruction) {
   EXPECT_EQ(ICmp, CT.Captures[6]->getUser());
   EXPECT_EQ(1u, CT.Captures[6]->getOperandNo());
 }
+
+TEST(CaptureTracking, DerivedPointerIfBasePointerCaptured) {
+  StringRef Assembly = R"(
+    declare void @bar(ptr)
+
+    define void @test() {
+      %stkobj = alloca [2 x i32]
+      %derived = getelementptr inbounds [2 x i32], ptr %stkobj, i64 0, i64 1
+      store i32 1, ptr %derived
+      call void @bar(ptr %stkobj)
+      ret void
+    }
+  )";
+
+  LLVMContext Context;
+  SMDiagnostic Error;
+  auto M = parseAssemblyString(Assembly, Error, Context);
+  ASSERT_TRUE(M) << "Bad assembly?";
+
+  Function *F = M->getFunction("test");
+  BasicBlock *BB = &F->getEntryBlock();
+  Instruction *StackObj = &*BB->begin();
+  Instruction *DerviedPtr = StackObj->getNextNode();
+  
+  // The base object and its derived pointer are both captured.
+  EXPECT_TRUE(PointerMayBeCaptured(StackObj, true));
+  EXPECT_TRUE(PointerMayBeCaptured(DerviedPtr, true));
+}