[PECOFF] Take into account all sections when setting size fields in the PE header

This patch changes lld to go through all sections while calculating the size
for SizeOfCode, SizeOfInitializedData and SizeOfUninitializedData fields in the
PE header, instead of using only a small set of hard-coded sections.

This only really changes SizeOfInitializedData which didn't include .reloc
section before this patch.

Patch by Ron Ofir.

llvm-svn: 190799

2 files changed, 29 insertions, 4 deletions

diff --git a/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp b/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
index 8d058911346..892fcfb15af 100644
--- a/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
@@ -474,6 +474,10 @@ public:
     return _sectionHeader;
   }
 
+  ulittle32_t getSectionCharacteristics() {
+    return _sectionHeader.Characteristics;
+  }
+
   void appendAtom(const DefinedAtom *atom) {
     // Atom may have to be at a proper alignment boundary. If so, move the
     // pointer to make a room after the last atom before adding new one.
@@ -815,7 +819,7 @@ public:
 
     // Now that we know the size and file offset of sections. Set the file
     // header accordingly.
-    peHeader->setSizeOfCode(text->size());
+    peHeader->setSizeOfCode(calcSizeOfCode());
     if (text->size()) {
       peHeader->setBaseOfCode(text->getVirtualAddress());
     }
@@ -824,14 +828,35 @@ public:
     } else if (data->size()) {
       peHeader->setBaseOfData(data->getVirtualAddress());
     }
-    peHeader->setSizeOfInitializedData(rdata->size() + data->size());
-    peHeader->setSizeOfUninitializedData(bss->size());
+    peHeader->setSizeOfInitializedData(calcSizeOfInitializedData());
+    peHeader->setSizeOfUninitializedData(calcSizeOfUninitializedData());
     peHeader->setNumberOfSections(_numSections);
     peHeader->setSizeOfImage(_imageSizeInMemory);
 
     setAddressOfEntryPoint(text, peHeader);
   }
 
+  uint64_t calcSectionSize(llvm::COFF::SectionCharacteristics sectionType) {
+    uint64_t ret = 0;
+    for (auto &cp : _chunks)
+      if (SectionChunk *chunk = dyn_cast<SectionChunk>(&*cp))
+        if (chunk->getSectionCharacteristics() & sectionType)
+          ret += chunk->size();
+    return ret;
+  }
+
+  uint64_t calcSizeOfInitializedData() {
+    return calcSectionSize(llvm::COFF::IMAGE_SCN_CNT_INITIALIZED_DATA);
+  }
+
+  uint64_t calcSizeOfUninitializedData() {
+    return calcSectionSize(llvm::COFF::IMAGE_SCN_CNT_UNINITIALIZED_DATA);
+  }
+
+  uint64_t calcSizeOfCode() {
+    return calcSectionSize(llvm::COFF::IMAGE_SCN_CNT_CODE);
+  }
+
   virtual error_code writeFile(const File &linkedFile, StringRef path) {
     this->build(linkedFile);
 
diff --git a/lld/test/pecoff/hello.test b/lld/test/pecoff/hello.test
index 575ff709818..3830429fcc6 100644
--- a/lld/test/pecoff/hello.test
+++ b/lld/test/pecoff/hello.test
@@ -5,7 +5,7 @@
 # RUN:   && llvm-readobj -file-headers %t1 | FileCheck -check-prefix=FILE %s
 
 FILE: ImageOptionalHeader {
-FILE:   SizeOfInitializedData: 512
+FILE:   SizeOfInitializedData: 1024
 FILE: }
 
 # RUN: lld -flavor link /out:%t1 /subsystem:console /force /opt:noref \