From 0d4b6d5c9833f6c577865817649f5fa644ad5e73 Mon Sep 17 00:00:00 2001
From: Rafael Espindola <rafael.espindola@gmail.com>
Date: Thu, 22 Sep 2016 16:47:21 +0000
Subject: [PATCH] Avoid duplicated code.

This also fixes the linker script accounting for the ELF header in
some places but not in others.

llvm-svn: 282173
---
 lld/ELF/LinkerScript.cpp      | 5 ++---
 lld/ELF/OutputSections.h      | 6 ++++++
 lld/ELF/Writer.cpp            | 6 +-----
 lld/test/ELF/oformat-binary.s | 4 ++++
 4 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/lld/ELF/LinkerScript.cpp b/lld/ELF/LinkerScript.cpp
index 37b51ce19240..d4577c2f41d6 100644
--- a/lld/ELF/LinkerScript.cpp
+++ b/lld/ELF/LinkerScript.cpp
@@ -587,8 +587,7 @@ template <class ELFT> void LinkerScript<ELFT>::assignAddresses() {
       Sec->setVA(0);
   }
 
-  uintX_t HeaderSize =
-      Out<ELFT>::ElfHeader->getSize() + Out<ELFT>::ProgramHeaders->getSize();
+  uintX_t HeaderSize = getHeaderSize();
   if (HeaderSize > MinVA)
     fatal("Not enough space for ELF and program headers");
 
@@ -723,7 +722,7 @@ uint64_t LinkerScript<ELFT>::getOutputSectionAlign(StringRef Name) {
 }
 
 template <class ELFT> uint64_t LinkerScript<ELFT>::getHeaderSize() {
-  return Out<ELFT>::ElfHeader->getSize() + Out<ELFT>::ProgramHeaders->getSize();
+  return elf::getHeaderSize<ELFT>();
 }
 
 template <class ELFT> uint64_t LinkerScript<ELFT>::getSymbolValue(StringRef S) {
diff --git a/lld/ELF/OutputSections.h b/lld/ELF/OutputSections.h
index bf799e27e028..fedda160f702 100644
--- a/lld/ELF/OutputSections.h
+++ b/lld/ELF/OutputSections.h
@@ -812,6 +812,12 @@ private:
   llvm::SmallDenseMap<Key, OutputSectionBase<ELFT> *> Map;
 };
 
+template <class ELFT> uint64_t getHeaderSize() {
+  if (Config->OFormatBinary)
+    return 0;
+  return Out<ELFT>::ElfHeader->getSize() + Out<ELFT>::ProgramHeaders->getSize();
+}
+
 template <class ELFT> BuildIdSection<ELFT> *Out<ELFT>::BuildId;
 template <class ELFT> DynamicSection<ELFT> *Out<ELFT>::Dynamic;
 template <class ELFT> EhFrameHeader<ELFT> *Out<ELFT>::EhFrameHdr;
diff --git a/lld/ELF/Writer.cpp b/lld/ELF/Writer.cpp
index 6cdf0978e5ac..e42240fdd830 100644
--- a/lld/ELF/Writer.cpp
+++ b/lld/ELF/Writer.cpp
@@ -1148,11 +1148,7 @@ template <class ELFT> void Writer<ELFT>::fixHeaders() {
 
 // Assign VAs (addresses at run-time) to output sections.
 template <class ELFT> void Writer<ELFT>::assignAddresses() {
-  uintX_t VA = Config->ImageBase;
-  if (!Config->OFormatBinary)
-    VA +=
-        Out<ELFT>::ElfHeader->getSize() + Out<ELFT>::ProgramHeaders->getSize();
-
+  uintX_t VA = Config->ImageBase + getHeaderSize<ELFT>();
   uintX_t ThreadBssOffset = 0;
   for (OutputSectionBase<ELFT> *Sec : OutputSections) {
     uintX_t Alignment = Sec->getAlignment();
diff --git a/lld/test/ELF/oformat-binary.s b/lld/test/ELF/oformat-binary.s
index 3f6a2497d006..acd95c7cef36 100644
--- a/lld/test/ELF/oformat-binary.s
+++ b/lld/test/ELF/oformat-binary.s
@@ -11,6 +11,10 @@
 # RUN: ld.lld -o %t2.out --script %t.script %t --oformat binary
 # RUN: od -t x1 -v %t2.out | FileCheck %s
 
+# RUN: echo "SECTIONS { }" > %t.script
+# RUN: ld.lld -o %t2.out --script %t.script %t --oformat binary
+# RUN: od -t x1 -v %t2.out | FileCheck %s
+
 # RUN: not ld.lld -o %t3.out %t --oformat foo 2>&1 \
 # RUN:   | FileCheck %s --check-prefix ERR
 # ERR: unknown --oformat value: foo