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

[Weihang Fan via llvm-commits]

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

>From 421553bd68daa268b245ccef9d38b490118b9ba6 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 | 47 ++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index ba598d8415b18..a19b334391b19 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..64aee4eb0f680 100644
--- a/llvm/test/Transforms/JumpThreading/select.ll
+++ b/llvm/test/Transforms/JumpThreading/select.ll
@@ -665,6 +665,53 @@ if.end:
   ret i32 %v1
 }
 
+; Do not unfold select when parent block address is taken and used
+define i64 @ba_unfold(i1 %cond, i32 %arg) {
+; CHECK-LABEL: @ba_unfold(
+; CHECK: B4:
+; CHECK-NEXT:   [[P:%.*]] = phi i1
+; CHECK-NEXT:   [[ADDR:%.*]] = select i1 [[P]], ptr blockaddress(@ba_unfold, %[[B5:.*]]), ptr blockaddress(@ba_unfold, %[[B6:.*]])
+; CHECK:        [[B5]]:
+; CHECK-NEXT:   ret i64 0
+; CHECK:        [[B6]]:
+; CHECK-NEXT:   ret i64 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 ]
+  %addr = select i1 %p, ptr blockaddress(@ba_unfold, %B5), ptr blockaddress(@ba_unfold, %B6)
+  indirectbr ptr %addr, [label %B5, label %B6]
+
+B5:
+  ret i64 0
+
+B6:
+  ret i64 ptrtoint (ptr blockaddress(@ba_unfold, %B4) to i64)
+}
+
 ; branch_weights overflowing uint32_t
 !0 = !{!"branch_weights", i64 1073741824, i64 3221225472}
 !1 = !{!"function_entry_count", i64 1984}