From c683e281cdc99b04497d5102c28be0d2e991c9f6 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Mon, 3 Oct 2022 16:01:01 -0700
Subject: [PATCH] [BOLT] Properly set _end symbol

To properly set the "_end" symbol, we need to track the last allocatable
address. Simply emitting "_end" at the end of some section is not
sufficient since the order of section allocation is unknown during the
emission step.

Reviewed By: rafauler

Differential Revision: https://reviews.llvm.org/D135121
---
 bolt/lib/Core/BinaryEmitter.cpp      |  2 --
 bolt/lib/Rewrite/RewriteInstance.cpp | 26 ++++++++++++++------------
 bolt/test/X86/end-symbol.test        |  9 +++++++++
 3 files changed, 23 insertions(+), 14 deletions(-)
 create mode 100644 bolt/test/X86/end-symbol.test

diff --git a/bolt/lib/Core/BinaryEmitter.cpp b/bolt/lib/Core/BinaryEmitter.cpp
index 96d6d8e..c4f887d 100644
--- a/bolt/lib/Core/BinaryEmitter.cpp
+++ b/bolt/lib/Core/BinaryEmitter.cpp
@@ -211,8 +211,6 @@ void BinaryEmitter::emitAll(StringRef OrgSecPrefix) {
   }
 
   emitDataSections(OrgSecPrefix);
-
-  Streamer.emitLabel(BC.Ctx->getOrCreateSymbol("_end"));
 }
 
 void BinaryEmitter::emitFunctions() {
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 0208c0c..16e2db8 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -4730,27 +4730,29 @@ void RewriteInstance::updateELFSymbolTable(
     Expected<StringRef> SymbolName = Symbol.getName(StringSection);
     assert(SymbolName && "cannot get symbol name");
 
-    auto updateSymbolValue = [&](const StringRef Name, unsigned &IsUpdated) {
-      NewSymbol.st_value = getNewValueForSymbol(Name);
+    auto updateSymbolValue = [&](const StringRef Name,
+                                 Optional<uint64_t> Value = NoneType()) {
+      NewSymbol.st_value = Value ? *Value : getNewValueForSymbol(Name);
       NewSymbol.st_shndx = ELF::SHN_ABS;
       outs() << "BOLT-INFO: setting " << Name << " to 0x"
              << Twine::utohexstr(NewSymbol.st_value) << '\n';
-      ++IsUpdated;
     };
 
     if (opts::HotText &&
-        (*SymbolName == "__hot_start" || *SymbolName == "__hot_end"))
-      updateSymbolValue(*SymbolName, NumHotTextSymsUpdated);
-
-    if (opts::HotData &&
-        (*SymbolName == "__hot_data_start" || *SymbolName == "__hot_data_end"))
-      updateSymbolValue(*SymbolName, NumHotDataSymsUpdated);
+        (*SymbolName == "__hot_start" || *SymbolName == "__hot_end")) {
+      updateSymbolValue(*SymbolName);
+      ++NumHotTextSymsUpdated;
+    }
 
-    if (*SymbolName == "_end") {
-      unsigned Ignored;
-      updateSymbolValue(*SymbolName, Ignored);
+    if (opts::HotData && (*SymbolName == "__hot_data_start" ||
+                          *SymbolName == "__hot_data_end")) {
+      updateSymbolValue(*SymbolName);
+      ++NumHotDataSymsUpdated;
     }
 
+    if (*SymbolName == "_end")
+      updateSymbolValue(*SymbolName, NextAvailableAddress);
+
     if (IsDynSym)
       Write((&Symbol - cantFail(Obj.symbols(&SymTabSection)).begin()) *
                 sizeof(ELFSymTy),
diff --git a/bolt/test/X86/end-symbol.test b/bolt/test/X86/end-symbol.test
new file mode 100644
index 0000000..98018c3
--- /dev/null
+++ b/bolt/test/X86/end-symbol.test
@@ -0,0 +1,9 @@
+# RUN: yaml2obj %p/Inputs/plt-sec.yaml &> %t.exe
+# RUN: llvm-bolt %t.exe -o %t.out
+# RUN: (llvm-readelf --program-headers %t.out | tac ; llvm-nm %t.out) \
+# RUN:   | FileCheck %s
+
+## Check that llvm-bolt correctly updates _end symbol to match the end of the
+## last loadable segment.
+CHECK: LOAD 0x[[#%x,OFFSET:]] 0x[[#%x,VMA:]] 0x[[#%x,PA:]] 0x[[#%x,FILESIZE:]] 0x[[#%x,MEMSIZE:]]
+CHECK: [[#%.16x, VMA + MEMSIZE]] {{.}} _end
-- 
2.7.4