From 9120098319d5658788fe68421bc5f724ba4250c6 Mon Sep 17 00:00:00 2001
From: Roman Lebedev <lebedev.ri@gmail.com>
Date: Sun, 20 Jun 2021 12:18:15 +0300
Subject: [PATCH] [SimplifyCFG] HoistThenElseCodeToIf(): don't hoist if either
 block has it's address taken

This problem is exposed by D104598, after it tail-merges `ret` in
`@test_inline_constraint_S_label`, the verifier would start complaining
`invalid operand for inline asm constraint 'S'`.

Essentially, taking address of a block is mismodelled in IR.
It should probably be an explicit instruction, a first one in block,
that isn't identical to any other instruction of the same type,
so that it can't be hoisted.
---
 lib/Transforms/Utils/SimplifyCFG.cpp          |  6 ++++++
 .../CodeGen/AArch64/inlineasm-S-constraint.ll | 13 ++++++++++++
 .../hoist-from-addresstaken-block.ll          | 21 +++++++++++++++++++
 3 files changed, 40 insertions(+)
 create mode 100644 test/Transforms/SimplifyCFG/hoist-from-addresstaken-block.ll

diff --git a/lib/Transforms/Utils/SimplifyCFG.cpp b/lib/Transforms/Utils/SimplifyCFG.cpp
index 642320e3289..4f3559d9779 100644
--- a/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -1404,6 +1404,12 @@ bool SimplifyCFGOpt::HoistThenElseCodeToIf(BranchInst *BI,
   BasicBlock *BB1 = BI->getSuccessor(0); // The true destination.
   BasicBlock *BB2 = BI->getSuccessor(1); // The false destination
 
+  // If either of the blocks has it's address taken, then we can't do this fold,
+  // because the code we'd hoist would no longer run when we jump into the block
+  // by it's address.
+  if (BB1->hasAddressTaken() || BB2->hasAddressTaken())
+    return false;
+
   BasicBlock::iterator BB1_Itr = BB1->begin();
   BasicBlock::iterator BB2_Itr = BB2->begin();
 
diff --git a/test/CodeGen/AArch64/inlineasm-S-constraint.ll b/test/CodeGen/AArch64/inlineasm-S-constraint.ll
index 3fb2a3f32ce..0e1169500e6 100644
--- a/test/CodeGen/AArch64/inlineasm-S-constraint.ll
+++ b/test/CodeGen/AArch64/inlineasm-S-constraint.ll
@@ -18,3 +18,16 @@ loc:
 loc2:
   ret i32 42
 }
+define i32 @test_inline_constraint_S_label_tailmerged(i1 %in) {
+; CHECK-LABEL: test_inline_constraint_S_label_tailmerged:
+  call void asm sideeffect "adr x0, $0", "S"(i8* blockaddress(@test_inline_constraint_S_label_tailmerged, %loc))
+; CHECK: adr x0, .Ltmp{{[0-9]+}}
+br i1 %in, label %loc, label %loc2
+loc:
+  br label %common.ret
+loc2:
+  br label %common.ret
+common.ret:
+  %common.retval = phi i32 [ 0, %loc ], [ 42, %loc2 ]
+  ret i32 %common.retval
+}
diff --git a/test/Transforms/SimplifyCFG/hoist-from-addresstaken-block.ll b/test/Transforms/SimplifyCFG/hoist-from-addresstaken-block.ll
new file mode 100644
index 00000000000..d00593e94f9
--- /dev/null
+++ b/test/Transforms/SimplifyCFG/hoist-from-addresstaken-block.ll
@@ -0,0 +1,21 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -mem2reg -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s
+
+define i32 @test_inline_constraint_S_label_tailmerged(i1 %in) {
+; CHECK-LABEL: @test_inline_constraint_S_label_tailmerged(
+; CHECK-NEXT:    call void asm sideeffect "adr x0, $0", "S"(i8* blockaddress(@test_inline_constraint_S_label_tailmerged, [[COMMON_RET:%.*]]))
+; CHECK-NEXT:    [[COMMON_RETVAL:%.*]] = select i1 [[IN:%.*]], i32 0, i32 42
+; CHECK-NEXT:    br label [[COMMON_RET]]
+; CHECK:       common.ret:
+; CHECK-NEXT:    ret i32 [[COMMON_RETVAL]]
+;
+  call void asm sideeffect "adr x0, $0", "S"(i8* blockaddress(@test_inline_constraint_S_label_tailmerged, %loc))
+br i1 %in, label %loc, label %loc2
+loc:
+  br label %common.ret
+loc2:
+  br label %common.ret
+common.ret:
+  %common.retval = phi i32 [ 0, %loc ], [ 42, %loc2 ]
+  ret i32 %common.retval
+}