[lld] [lld-macho] Fix branch extension thunk estimation logic (PR #120529)

[via llvm-commits]

https://github.com/alx32 updated https://github.com/llvm/llvm-project/pull/120529

>From fe6369ae90d7a69dc013781678b54b6ac39ea9d0 Mon Sep 17 00:00:00 2001
From: Alex B <alexborcan at meta.com>
Date: Wed, 18 Dec 2024 23:37:48 -0800
Subject: [PATCH 1/2] [lld-macho] Fix branch extension thunk estimation logic

---
 lld/MachO/ConcatOutputSection.cpp | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/lld/MachO/ConcatOutputSection.cpp b/lld/MachO/ConcatOutputSection.cpp
index 6a9301f84a03ef..9610a80e309bac 100644
--- a/lld/MachO/ConcatOutputSection.cpp
+++ b/lld/MachO/ConcatOutputSection.cpp
@@ -184,6 +184,15 @@ uint64_t TextOutputSection::estimateStubsInRangeVA(size_t callIdx) const {
     InputSection *isec = inputs[i];
     isecEnd = alignToPowerOf2(isecEnd, isec->align) + isec->getSize();
   }
+
+  // Tally up any thunks that have already been placed that have address higher
+  // than the equivalent callIdx. We first find the index of the first thunk
+  // that is beyond the current inputs[callIdx].
+  auto itPostcallIdxThunks = std::partition_point(
+      thunks.begin(), thunks.end(),
+      [isecVA](const ConcatInputSection *t) { return t->getVA() <= isecVA; });
+  uint32_t existingForwardThunks = thunks.end() - itPostcallIdxThunks;
+
   // Estimate the address after which call sites can safely call stubs
   // directly rather than through intermediary thunks.
   uint64_t forwardBranchRange = target->forwardBranchRange;
@@ -191,8 +200,21 @@ uint64_t TextOutputSection::estimateStubsInRangeVA(size_t callIdx) const {
          "should not run thunk insertion if all code fits in jump range");
   assert(isecEnd - isecVA <= forwardBranchRange &&
          "should only finalize sections in jump range");
-  uint64_t stubsInRangeVA = isecEnd + maxPotentialThunks * target->thunkSize +
-                            in.stubs->getSize() - forwardBranchRange;
+
+  // Estimate the maximum size of the code, right before the stubs section
+  uint32_t maxTextSize = 0;
+  // Add the size of all the inputs, including the unprocessed ones.
+  maxTextSize += isecEnd;
+
+  // Add the size of the thunks that may be created in the future
+  maxTextSize += maxPotentialThunks * target->thunkSize;
+
+  // Add the size of the thunks that have already been created that are ahead
+  maxTextSize += existingForwardThunks * target->thunkSize;
+
+  uint64_t stubsInRangeVA =
+      maxTextSize + in.stubs->getSize() - forwardBranchRange;
+
   log("thunks = " + std::to_string(thunkMap.size()) +
       ", potential = " + std::to_string(maxPotentialThunks) +
       ", stubs = " + std::to_string(in.stubs->getSize()) + ", isecVA = " +

>From bc5e3426d0f71a2fc364258a4009b3848f856d95 Mon Sep 17 00:00:00 2001
From: Alex B <alexborcan at meta.com>
Date: Thu, 19 Dec 2024 09:53:27 -0800
Subject: [PATCH 2/2] uint32 => uint64 + extra comments

---
 lld/MachO/ConcatOutputSection.cpp | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/lld/MachO/ConcatOutputSection.cpp b/lld/MachO/ConcatOutputSection.cpp
index 9610a80e309bac..cf3ab72c351e4c 100644
--- a/lld/MachO/ConcatOutputSection.cpp
+++ b/lld/MachO/ConcatOutputSection.cpp
@@ -191,7 +191,7 @@ uint64_t TextOutputSection::estimateStubsInRangeVA(size_t callIdx) const {
   auto itPostcallIdxThunks = std::partition_point(
       thunks.begin(), thunks.end(),
       [isecVA](const ConcatInputSection *t) { return t->getVA() <= isecVA; });
-  uint32_t existingForwardThunks = thunks.end() - itPostcallIdxThunks;
+  uint64_t existingForwardThunks = thunks.end() - itPostcallIdxThunks;
 
   // Estimate the address after which call sites can safely call stubs
   // directly rather than through intermediary thunks.
@@ -201,26 +201,30 @@ uint64_t TextOutputSection::estimateStubsInRangeVA(size_t callIdx) const {
   assert(isecEnd - isecVA <= forwardBranchRange &&
          "should only finalize sections in jump range");
 
-  // Estimate the maximum size of the code, right before the stubs section
-  uint32_t maxTextSize = 0;
+  // Estimate the maximum size of the code, right before the stubs section.
+  uint64_t maxTextSize = 0;
   // Add the size of all the inputs, including the unprocessed ones.
   maxTextSize += isecEnd;
 
-  // Add the size of the thunks that may be created in the future
+  // Add the size of the thunks that may be created in the future. Since
+  // 'maxPotentialThunks' overcounts, this is an estimate of the upper limit.
   maxTextSize += maxPotentialThunks * target->thunkSize;
 
   // Add the size of the thunks that have already been created that are ahead
   maxTextSize += existingForwardThunks * target->thunkSize;
 
-  uint64_t stubsInRangeVA =
-      maxTextSize + in.stubs->getSize() - forwardBranchRange;
+  // Estimated maximum VA of last stub.
+  uint64_t maxVAOfLastStub = maxTextSize + in.stubs->getSize();
+
+  // Calculaate the first address that is gueranteed to not need a thunk to
+  // reach any stub.git
+  uint64_t stubsInRangeVA = maxVAOfLastStub - forwardBranchRange;
 
   log("thunks = " + std::to_string(thunkMap.size()) +
-      ", potential = " + std::to_string(maxPotentialThunks) +
-      ", stubs = " + std::to_string(in.stubs->getSize()) + ", isecVA = " +
-      utohexstr(isecVA) + ", threshold = " + utohexstr(stubsInRangeVA) +
-      ", isecEnd = " + utohexstr(isecEnd) +
-      ", tail = " + utohexstr(isecEnd - isecVA) +
+      ", potential = " + std::to_string(maxPotentialThunks) + ", stubs = " +
+      std::to_string(in.stubs->getSize()) + ", isecVA = " + utohexstr(isecVA) +
+      ", threshold = " + utohexstr(stubsInRangeVA) + ", isecEnd = " +
+      utohexstr(isecEnd) + ", tail = " + utohexstr(isecEnd - isecVA) +
       ", slop = " + utohexstr(forwardBranchRange - (isecEnd - isecVA)));
   return stubsInRangeVA;
 }