[CGP] Do not eliminate blocks which have their address taken

When eliminating a block, codegenprepare updates all blockaddress
expressions which reference the block. In case the blockaddress is
located in a different function, it leads to updates across function
boundaries, which is problematic for a function pass like
codegenprepare.

If blockaddress is in a function defined before the current one, the
update to the blockaddress is lost.

This change adds a check to avoid eliminations of any block which has
its address taken.

Fixes: #161164

Author: Frank Tetzel

llvm/lib/CodeGen/CodeGenPrepare.cpp

@@ -954,6 +954,12 @@ bool CodeGenPrepare::eliminateMostlyEmptyBlocks(Function &F) {
 bool CodeGenPrepare::isMergingEmptyBlockProfitable(BasicBlock *BB,
                                                    BasicBlock *DestBB,
                                                    bool isPreheader) {
+  // Do not eliminate blocks which have their address taken.
+  // This could lead to updates across functions which is problematic in a
+  // function pass like codegenprepare. The update to a blockaddress in a
+  // function defined before the function with the eliminated block is lost.
+  if(BB->hasAddressTaken()) return false;
+
   // Do not delete loop preheaders if doing so would create a critical edge.
   // Loop preheaders can be good locations to spill registers. If the
   // preheader is deleted and we create a critical edge, registers may be

llvm/test/Transforms/CodeGenPrepare/X86/blockaddress-reversed.ll

@@ -0,0 +1,53 @@
+; RUN: llc <%s | FileCheck %s
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK:      .Ltmp0: # Address of block that was removed by CodeGen
+; CHECK-NEXT: .Ltmp1: # Address of block that was removed by CodeGen
+
+define i64 @main(ptr %p, i1 %c1) {
+first:
+  br label %loop_head
+
+loop_head:                                        ; preds = %loop_end, %first
+  br i1 %c1, label %scope0, label %scope1
+
+scope1:                                           ; preds = %loop_head
+  br i1 %c1, label %scope1_exit, label %scope2
+
+scope2:                                           ; preds = %scope1
+  br i1 %c1, label %scope2_exit, label %inner
+
+inner:                                            ; preds = %scope2
+  br label %scope2_exit
+
+scope2_exit:                                      ; preds = %inner, %scope2
+  %phi2 = phi ptr [ %p, %inner ], [ null, %scope2 ]
+  br label %scope1_exit
+
+scope1_exit:                                      ; preds = %scope2_exit, %scope1
+  %phi1 = phi ptr [ %phi2, %scope2_exit ], [ null, %scope1 ]
+  %val1 = load i128, ptr %phi1, align 16
+  br label %loop_end
+
+scope0:                                           ; preds = %loop_head
+  %ptr0 = select i1 %c1, ptr null, ptr %p
+  %val0 = load i128, ptr %ptr0, align 16
+  br label %loop_end
+
+loop_end:                                         ; preds = %scope0, %scope1_exit
+  %storemerge = phi i128 [ %val1, %scope1_exit ], [ %val0, %scope0 ]
+  store i128 %storemerge, ptr %p, align 16
+  br label %loop_head
+}
+
+define void @foo0(ptr %jumpAddr) {
+; CHECK:        movq $.Ltmp0, (%rdi)
+  store ptr blockaddress(@main, %scope2_exit), ptr %jumpAddr, align 8
+  ret void
+}
+
+define void @foo1(ptr %jumpAddr) {
+; CHECK:        movq $.Ltmp1, (%rdi)
+  store ptr blockaddress(@main, %scope1_exit), ptr %jumpAddr, align 8
+  ret void
+}

llvm/test/Transforms/CodeGenPrepare/X86/blockaddress.ll

@@ -0,0 +1,53 @@
+; RUN: llc <%s | FileCheck %s
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @foo0(ptr %jumpAddr) {
+; CHECK:        movq $.Ltmp0, (%rdi)
+  store ptr blockaddress(@main, %scope2_exit), ptr %jumpAddr, align 8
+  ret void
+}
+
+define void @foo1(ptr %jumpAddr) {
+; CHECK:        movq $.Ltmp1, (%rdi)
+  store ptr blockaddress(@main, %scope1_exit), ptr %jumpAddr, align 8
+  ret void
+}
+
+; CHECK:      .Ltmp0: # Address of block that was removed by CodeGen
+; CHECK-NEXT: .Ltmp1: # Address of block that was removed by CodeGen
+
+define i64 @main(ptr %p, i1 %c1) {
+first:
+  br label %loop_head
+
+loop_head:                                        ; preds = %loop_end, %first
+  br i1 %c1, label %scope0, label %scope1
+
+scope1:                                           ; preds = %loop_head
+  br i1 %c1, label %scope1_exit, label %scope2
+
+scope2:                                           ; preds = %scope1
+  br i1 %c1, label %scope2_exit, label %inner
+
+inner:                                            ; preds = %scope2
+  br label %scope2_exit
+
+scope2_exit:                                      ; preds = %inner, %scope2
+  %phi2 = phi ptr [ %p, %inner ], [ null, %scope2 ]
+  br label %scope1_exit
+
+scope1_exit:                                      ; preds = %scope2_exit, %scope1
+  %phi1 = phi ptr [ %phi2, %scope2_exit ], [ null, %scope1 ]
+  %val1 = load i128, ptr %phi1, align 16
+  br label %loop_end
+
+scope0:                                           ; preds = %loop_head
+  %ptr0 = select i1 %c1, ptr null, ptr %p
+  %val0 = load i128, ptr %ptr0, align 16
+  br label %loop_end
+
+loop_end:                                         ; preds = %scope0, %scope1_exit
+  %storemerge = phi i128 [ %val1, %scope1_exit ], [ %val0, %scope0 ]
+  store i128 %storemerge, ptr %p, align 16
+  br label %loop_head
+}