[CodeExtractor] Prevent extraction of block involving blockaddress

BlockAddress are only valid within their function context, which does not
interact well with CodeExtractor. Detect this case and prevent it.

Differential Revision: https://reviews.llvm.org/D33839


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@306448 91177308-0d34-0410-b5e6-96231b3b80d8

lib/Transforms/Utils/CodeExtractor.cpp

@@ -59,6 +59,33 @@ bool CodeExtractor::isBlockValidForExtraction(const BasicBlock &BB) {
   // Landing pads must be in the function where they were inserted for cleanup.
   if (BB.isEHPad())
     return false;
+  // taking the address of a basic block moved to another function is illegal
+  if (BB.hasAddressTaken())
+    return false;
+
+  // don't hoist code that uses another basicblock address, as it's likely to
+  // lead to unexpected behavior, like cross-function jumps
+  SmallPtrSet<User const *, 16> Visited;
+  SmallVector<User const *, 16> ToVisit;
+
+  for (Instruction const &Inst : BB)
+    ToVisit.push_back(&Inst);
+
+  while (!ToVisit.empty()) {
+    User const *Curr = ToVisit.pop_back_val();
+    if (!Visited.insert(Curr).second)
+      continue;
+    if (isa<BlockAddress const>(Curr))
+      return false; // even a reference to self is likely to be not compatible
+
+    if (isa<Instruction>(Curr) && cast<Instruction>(Curr)->getParent() != &BB)
+      continue;
+
+    for (auto const &U : Curr->operands()) {
+      if (auto *UU = dyn_cast<User>(U))
+        ToVisit.push_back(UU);
+    }
+  }
 
   // Don't hoist code containing allocas, invokes, or vastarts.
   for (BasicBlock::const_iterator I = BB.begin(), E = BB.end(); I != E; ++I) {

test/Transforms/CodeExtractor/BlockAddressReference.ll

@@ -0,0 +1,36 @@
+; RUN: opt < %s -loop-extract -S | FileCheck %s
+
+@label = common local_unnamed_addr global i8* null
+
+; CHECK: define
+; no outlined function
+; CHECK-NOT: define
+define i32 @sterix(i32 %n) {
+entry:
+  %tobool = icmp ne i32 %n, 0
+  ; this blockaddress references a basic block that goes in the extracted loop
+  %cond = select i1 %tobool, i8* blockaddress(@sterix, %for.cond), i8* blockaddress(@sterix, %exit)
+  store i8* %cond, i8** @label
+  %cmp5 = icmp sgt i32 %n, 0
+  br i1 %cmp5, label %for.body, label %exit
+
+for.cond:
+  %mul = shl nsw i32 %s.06, 1
+  %exitcond = icmp eq i32 %inc, %n
+  br i1 %exitcond, label %exit.loopexit, label %for.body
+
+for.body:
+  %i.07 = phi i32 [ %inc, %for.cond ], [ 0, %entry ]
+  %s.06 = phi i32 [ %mul, %for.cond ], [ 1, %entry ]
+  %inc = add nuw nsw i32 %i.07, 1
+  br label %for.cond
+
+exit.loopexit:
+  %phitmp = icmp ne i32 %s.06, 2
+  %phitmp8 = zext i1 %phitmp to i32
+  br label %exit
+
+exit:
+  %s.1 = phi i32 [ 1, %entry ], [ %phitmp8, %exit.loopexit ]
+  ret i32 %s.1
+}

test/Transforms/CodeExtractor/BlockAddressSelfReference.ll

@@ -0,0 +1,43 @@
+; RUN: opt < %s -loop-extract -S | FileCheck %s
+
+@choum.addr = internal unnamed_addr constant [3 x i8*] [i8* blockaddress(@choum, %12), i8* blockaddress(@choum, %16), i8* blockaddress(@choum, %20)]
+
+; CHECK: define
+; no outlined function
+; CHECK-NOT: define
+
+define void @choum(i32, i32* nocapture, i32) {
+  %4 = icmp sgt i32 %0, 0
+  br i1 %4, label %5, label %26
+
+  %6 = sext i32 %2 to i64
+  %7 = getelementptr inbounds [3 x i8*], [3 x i8*]* @choum.addr, i64 0, i64 %6
+  %8 = load i8*, i8** %7
+  %9 = zext i32 %0 to i64
+  br label %10
+
+  %11 = phi i64 [ 0, %5 ], [ %24, %20 ]
+  indirectbr i8* %8, [label %12, label %16, label %20]
+
+  %13 = getelementptr inbounds i32, i32* %1, i64 %11
+  %14 = load i32, i32* %13
+  %15 = add nsw i32 %14, 1
+  store i32 %15, i32* %13
+  br label %16
+
+  %17 = getelementptr inbounds i32, i32* %1, i64 %11
+  %18 = load i32, i32* %17
+  %19 = shl nsw i32 %18, 1
+  store i32 %19, i32* %17
+  br label %20
+
+  %21 = getelementptr inbounds i32, i32* %1, i64 %11
+  %22 = load i32, i32* %21
+  %23 = add nsw i32 %22, -3
+  store i32 %23, i32* %21
+  %24 = add nuw nsw i64 %11, 1
+  %25 = icmp eq i64 %24, %9
+  br i1 %25, label %26, label %10
+
+  ret void
+}