[llvm] [BOLT] Fix bug in lookupStubFromGroup (PR #119710)

[via llvm-commits]

https://github.com/liusy58 created https://github.com/llvm/llvm-project/pull/119710

The current implementation of `lookupStubFromGroup` is incorrect. The function is intended to find and return the closest stub using `lower_bound`, which identifies the first element in a sorted range that is not less than a specified value. However, if such an element is not found within `Candidates` and the list is not empty, the function returns `nullptr`. Instead, it should check whether the last element satisfies the condition.

>From e964fa91dcd8f61b12f8dbc1c88a7c9914a0d958 Mon Sep 17 00:00:00 2001
From: liusy58 <liusy58 at linux.alibaba.com>
Date: Thu, 12 Dec 2024 23:13:50 +0800
Subject: [PATCH] [BOLT] Fix bug in lookupStubFromGroup

The current implementation of `lookupStubFromGroup` is incorrect. The function
is intended to find and return the closest stub using `lower_bound`, which
identifies the first element in a sorted range that is not less than a specified
value. However, if such an element is not found within `Candidates` and the list
is not empty, the function returns `nullptr`. Instead, it should check whether
the last element satisfies the condition.
---
 bolt/lib/Passes/LongJmp.cpp           |  6 ++---
 bolt/test/AArch64/long-jmp-one-stub.s | 32 +++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)
 create mode 100644 bolt/test/AArch64/long-jmp-one-stub.s

diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index c1b8c03324e0e2..e6bd417705e6ff 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -138,9 +138,9 @@ BinaryBasicBlock *LongJmpPass::lookupStubFromGroup(
           const std::pair<uint64_t, BinaryBasicBlock *> &RHS) {
         return LHS.first < RHS.first;
       });
-  if (Cand == Candidates.end())
-    return nullptr;
-  if (Cand != Candidates.begin()) {
+  if (Cand == Candidates.end()) {
+    Cand = std::prev(Cand);
+  } else if (Cand != Candidates.begin()) {
     const StubTy *LeftCand = std::prev(Cand);
     if (Cand->first - DotAddress > DotAddress - LeftCand->first)
       Cand = LeftCand;
diff --git a/bolt/test/AArch64/long-jmp-one-stub.s b/bolt/test/AArch64/long-jmp-one-stub.s
new file mode 100644
index 00000000000000..1c45d95237cbe3
--- /dev/null
+++ b/bolt/test/AArch64/long-jmp-one-stub.s
@@ -0,0 +1,32 @@
+# This test verifies that no unnecessary stubs are inserted when each DotAddress increases during a lookup.
+
+# REQUIRES: system-linux, asserts
+
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown %s -o %t.o
+# RUN: %clang %cflags -O0 %t.o -o %t.exe -Wl,-q
+# RUN: link_fdata %s %t.o %t.fdata
+# RUN: llvm-bolt %t.exe -o %t.bolt  \
+# RUN:   --data %t.fdata  | FileCheck %s
+
+# CHECK: BOLT-INFO: Inserted 1 stubs in the hot area and 0 stubs in the cold area. 
+
+  .section .text
+  .global _start
+  .global far_away_func
+
+  .align 4
+  .global _start
+  .type _start, %function
+_start:
+# FDATA: 0 [unknown] 0 1 _start 0 0 100
+    bl far_away_func
+    bl far_away_func
+    ret  
+  .space 0x8000000
+  .global far_away_func
+  .type far_away_func, %function
+far_away_func:
+    add x0, x0, #1
+    ret
+
+.reloc 0, R_AARCH64_NONE