3 files changed, 32 insertions, 5 deletions
diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index fae655e119d1..02588c29462a 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -353,7 +353,7 @@ public:
   }
 
   virtual bool isUnconditionalBranch(const MCInst &Inst) const {
-    return Analysis->isUnconditionalBranch(Inst);
+    return Analysis->isUnconditionalBranch(Inst) && !isTailCall(Inst);
   }
 
   virtual bool isIndirectBranch(const MCInst &Inst) const {
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index ef6ae1e500ee..b7678cbff06b 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -78,10 +78,6 @@ public:
     return Analysis->isBranch(Inst) && !isTailCall(Inst);
   }
 
-  bool isUnconditionalBranch(const MCInst &Inst) const override {
-    return Analysis->isUnconditionalBranch(Inst) && !isTailCall(Inst);
-  }
-
   bool isNoop(const MCInst &Inst) const override {
     return X86::isNOP(Inst.getOpcode());
   }
diff --git a/bolt/test/AArch64/ext-double-jump.s b/bolt/test/AArch64/ext-double-jump.s
new file mode 100644
index 000000000000..a039cd424c9e
--- /dev/null
+++ b/bolt/test/AArch64/ext-double-jump.s
@@ -0,0 +1,31 @@
+# This test checks that remove double jumps pass works properly with
+# non-local branches.
+
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown %s -o %t.o
+# RUN: %clang %cflags -nostartfiles -nodefaultlibs %t.o -o %t.exe -Wl,-q
+# RUN: llvm-bolt %t.exe -o %t.bolt -peepholes=double-jumps
+
+  .text
+  .align 4
+  .global dummy1
+  .type dummy1, %function
+dummy1:
+  mov x2, x0
+  ret
+  .size dummy1, .-dummy1
+
+  .global dummy2
+  .type dummy2, %function
+dummy2:
+  mov x1, x0
+  ret
+  .size dummy2, .-dummy2
+
+  .global _start
+  .type _start, %function
+_start:
+  cbz  x10, 1f
+  b dummy1
+1:
+  b dummy2
+  .size _start, .-_start