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

[Ron Ofir]

Hi Bigcheese, rui314,

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.

http://llvm-reviews.chandlerc.com/D1616

Files:
  lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
  test/pecoff/hello.test

Index: lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
===================================================================
--- lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
+++ lib/ReaderWriter/PECOFF/WriterPECOFF.cpp
@@ -468,7 +468,19 @@
     _sectionHeader.SizeOfRawData = size();
     return _sectionHeader;
   }
+  
+  virtual bool isCodeSection() {
+    return _sectionHeader.Characteristics & llvm::COFF::IMAGE_SCN_CNT_CODE;
+  }
+  
+  virtual bool isInitializedDataSection() {
+    return _sectionHeader.Characteristics & llvm::COFF::IMAGE_SCN_CNT_INITIALIZED_DATA;
+  }
 
+  virtual bool isUninitializedDataSection() {
+    return _sectionHeader.Characteristics & llvm::COFF::IMAGE_SCN_CNT_UNINITIALIZED_DATA;
+  }
+
   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.
@@ -809,7 +822,7 @@
 
     // 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());
     }
@@ -818,8 +831,8 @@
     } 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);
 
@@ -906,7 +919,34 @@
         peHeader->setAddressOfEntryPoint(entryPointAddress);
     }
   }
-
+  
+  uint64_t calcSizeOfInitializedData() {
+    uint64_t ret = 0;
+    for (auto &cp : _chunks)
+      if (SectionChunk *chunk = dyn_cast<SectionChunk>(&*cp))
+        if (chunk->isInitializedDataSection())
+          ret += chunk->size();
+    return ret;
+  }
+  
+  uint64_t calcSizeOfUninitializedData() {
+    uint64_t ret = 0;
+    for (auto &cp : _chunks)
+      if (SectionChunk *chunk = dyn_cast<SectionChunk>(&*cp))
+        if (chunk->isUninitializedDataSection())
+          ret += chunk->size();
+    return ret;
+  }
+  
+  uint64_t calcSizeOfCode() {
+    uint64_t ret = 0;
+    for (auto &cp : _chunks)
+      if (SectionChunk *chunk = dyn_cast<SectionChunk>(&*cp))
+        if (chunk->isCodeSection())
+          ret += chunk->size();
+    return ret;
+  }
+  
   std::vector<std::unique_ptr<Chunk>> _chunks;
   const PECOFFLinkingContext &_PECOFFLinkingContext;
   uint32_t _numSections;
Index: test/pecoff/hello.test
===================================================================
--- test/pecoff/hello.test
+++ 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 \
-------------- next part --------------
A non-text attachment was scrubbed...
Name: D1616.1.patch
Type: text/x-patch
Size: 3151 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20130906/0384ea6c/attachment.bin>