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

[Ron Ofir]

  Several fixes according to Rui's comments

Hi Bigcheese, rui314,

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

CHANGE SINCE LAST DIFF
  http://llvm-reviews.chandlerc.com/D1616?vs=4099&id=4305#toc

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
@@ -473,6 +476,10 @@
     _sectionHeader.SizeOfRawData = size();
     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
@@ -814,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());
     }
@@ -823,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);
 
@@ -912,6 +920,27 @@
     }
   }
 
+  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);
+  }
+
   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.2.patch
Type: text/x-patch
Size: 2585 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20130914/863f88ba/attachment.bin>