[llvm] [BOLT][AArch64] Check Last Element Instead of Returning `nullptr` in `lookupStubFromGroup` (PR #114015)

[via llvm-commits]

https://github.com/liusy58 updated https://github.com/llvm/llvm-project/pull/114015

>From 2acb5077d9b7b236cb8ad864e3946121deda8a83 Mon Sep 17 00:00:00 2001
From: liusy58 <liusy58 at linux.alibaba.com>
Date: Tue, 29 Oct 2024 16:09:31 +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/one-stub.s | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)
 create mode 100644 bolt/test/AArch64/one-stub.s

diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 0b2d00300f46b9..77acd2decc9bd9 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -130,9 +130,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/one-stub.s b/bolt/test/AArch64/one-stub.s
new file mode 100644
index 00000000000000..8d118a1bdd327c
--- /dev/null
+++ b/bolt/test/AArch64/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 -nostdlib -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