[llvm] [JumpThreading] Do not thread if block has address taken and used (PR #135106)

[Weihang Fan via llvm-commits]

https://github.com/weihangf-apple created https://github.com/llvm/llvm-project/pull/135106

In an internal test case, JumpThreadingPass threaded across a basic block that has its address taken, resulting in the block being removed and accidentally invalidating a blockaddress. This caused SimplifyCFG to insert an unreachable instruction which ultimately caused a crash at runtime.

The issue in the case was that the select instruction was unfolded, causing two basic blocks to be created, which meant that the blockaddress of the original block is now ambiguous.

This PR guards against unfolding selects in basic blocks that have their address taken and used and adds a unit test.

>From 313104dd1b25a2ef6cdfd405799e606fe7df54c5 Mon Sep 17 00:00:00 2001
From: Weihang Fan <weihang_fan at apple.com>
Date: Fri, 4 Apr 2025 19:47:02 -0700
Subject: [PATCH] [JumpThreading] Do not thread if block has address taken and
 used

---
 llvm/lib/Transforms/Scalar/JumpThreading.cpp |  5 +++
 llvm/test/Transforms/JumpThreading/select.ll | 46 ++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index a518e02d762f6..0c189ae769e88 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -2933,6 +2933,11 @@ bool JumpThreadingPass::tryToUnfoldSelectInCurrBB(BasicBlock *BB) {
   if (LoopHeaders.count(BB))
     return false;
 
+  // If a block has its address taken, it would break the semantics of its block
+  // address. To be safe, don't thread the edge.
+  if (hasAddressTakenAndUsed(BB))
+    return false;
+
   for (BasicBlock::iterator BI = BB->begin();
        PHINode *PN = dyn_cast<PHINode>(BI); ++BI) {
     // Look for a Phi having at least one constant incoming value.
diff --git a/llvm/test/Transforms/JumpThreading/select.ll b/llvm/test/Transforms/JumpThreading/select.ll
index 4ec55a66bb8ac..1774ce9594337 100644
--- a/llvm/test/Transforms/JumpThreading/select.ll
+++ b/llvm/test/Transforms/JumpThreading/select.ll
@@ -665,6 +665,52 @@ if.end:
   ret i32 %v1
 }
 
+; Do not unfold select when parent block address is taken and used
+define i32 @ba_unfold(i1 %cond, i32 %arg) {
+; CHECK-LABEL: @ba_unfold(
+; CHECK: B4:
+; CHECK-NEXT:   [[P:%.*]] = phi i1 [ true, [[B1:%.*]] ], [ [[COND_NOT:%.*]], [[LOOPEXIT:%.*]] ]
+; CHECK-NEXT:   [[OFFSET:%.*]] = select i1 [[P]], i64 8, i64 16
+; CHECK-NEXT:   [[ADDR_I:%.*]] = add i64 [[OFFSET]], ptrtoint (ptr blockaddress(@ba_unfold, %B4) to i64)
+;
+entry:
+  br label %B1
+
+B1:
+  br i1 %cond, label %B2, label %B4
+
+B2:
+  br label %while.body
+
+B3:
+  %new_i = add nuw nsw i32 1, %i
+  %exitcond = icmp eq i32 %new_i, 5
+  br i1 %exitcond, label %loopexit, label %while.body
+
+while.body:
+  %i = phi i32 [ 0, %B2 ], [ %new_i, %B3 ]
+  %xor = xor i32 %i, 16
+  %cmp = icmp eq i32 %xor, %arg
+  br i1 %cmp, label %B3, label %loopexit
+
+loopexit:
+  %cond.not = xor i1 %cmp, true
+  br label %B4
+
+B4:
+  %p = phi i1 [true, %B1], [ %cond.not, %loopexit ]
+  %offset = select i1 %p, i64 8, i64 16
+  %addr_i = add i64 %offset, ptrtoint (ptr blockaddress(@ba_unfold, %B4) to i64)
+  %addr = inttoptr i64 %addr_i to ptr
+  indirectbr ptr %addr, [label %B5, label %B6]
+
+B5:
+  br label %B6
+
+B6:
+  ret i32 0
+}
+
 ; branch_weights overflowing uint32_t
 !0 = !{!"branch_weights", i64 1073741824, i64 3221225472}
 !1 = !{!"function_entry_count", i64 1984}