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

[Rui Ueyama]

Author: ruiu
Date: Mon Sep 16 12:39:26 2013
New Revision: 190799

URL: http://llvm.org/viewvc/llvm-project?rev=190799&view=rev
Log:
[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.

Modified:
    lld/trunk/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
    lld/trunk/test/pecoff/hello.test

Modified: lld/trunk/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
URL: http://llvm.org/viewvc/llvm-project/lld/trunk/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp?rev=190799&r1=190798&r2=190799&view=diff
==============================================================================
--- lld/trunk/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp (original)
+++ lld/trunk/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp Mon Sep 16 12:39:26 2013
@@ -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);
 

Modified: lld/trunk/test/pecoff/hello.test
URL: http://llvm.org/viewvc/llvm-project/lld/trunk/test/pecoff/hello.test?rev=190799&r1=190798&r2=190799&view=diff
==============================================================================
--- lld/trunk/test/pecoff/hello.test (original)
+++ lld/trunk/test/pecoff/hello.test Mon Sep 16 12:39:26 2013
@@ -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 \