[llvm-objcopy] Change -O binary to respect section removal and behave like GNU objcopy

[jakehehrlich]

Summary

The original -O binary implementation just copied segment data from the object and dumped it into a file. This doesn't take into account any operations performed on objects such as section removal. GNU objcopy has some specific behavior that we'd also like to respect. For instance using -O binary and -j <some_section> will dump <some_section> to a file. This change implements GNU objcopy style -O binary to as close of an approximation as I can determine.

[jhenderson]

Missing test for removing the first section in the first segment?

test/tools/llvm-objcopy/binary-remove-end.test

45

To make things slightly less confusing, could you make od use hex for the address, please (-Ax). Ditto for the other tests.

[jhenderson]

tools/llvm-objcopy/Object.cpp

705

I wonder if it might be useful to have an is_sorted assert here? What do you think?

Alternatively, why not just sort them here as were doing before?

706

I think we can be a little more specific, and say "one past the end of the last segment" in this comment. Unless I've missed something?

Nit: here and elsewhere layed -> laid.

737

Given that we now have a function called "LayoutSegments", I don't think you need to say "in some consistent fashion" here, since being laid out implies to me that the corresponding function has been called.

785

I'm not sure you need to have this and the next comment - the function names are quite clear as to what they do.

789

we need align -> we need to align.

852–863

I'm not a massive fan of this duplicate comment, but I also see that it's not necessarily trivial to remove the duplication. About the best idea I had was to pass a predicate into OrderSegments, which has the responsibility of returning the copied and sorted vector, filtered according to the predicate. If this sounds sensible to you, that would be a good way of getting rid of some of this duplication.

Looking further down, this sounds quite similar to what you are proposing with the LayoutSections in your new TODO. If you want to defer this to later, please add a TODO here as well.

868

Why do we need to std::unique this list?

871

This first sentence doesn't quite make sense - segment's don't have a sh_offset.

892

cosntructing -> constructing

903

I'd be explicit, and say that "total file size". I'm not sure I understand the second sentence at all (and therefore why we can't just use Offset).

[jakehehrlich]

tools/llvm-objcopy/Object.cpp

705

Adding the assert seems like a good idea. I separated sorting out because I need to know the OriginalOffset of the first segment in order to give LayoutSegments the starting offset.

852–863

Yeah This should be solvable by the same filter range idea that I propose below. I'll add a TODO for that since I think I'd like the filter range to reside in Iterator.h not inside this code.

868

Notice that we construct OrderedSegments by adding the ParentSegment of each section. There might be a single segment for multiple sections so we need to dedup these so that LayoutSegment doesn't do crazy things.

903

Say the last section of the last segment was removed. Then Offset will be past the last section but to be consistent with GNU objcopy we need TotalSize to be such that the output ends at the end of the section, not the end of the segment. Offset == TotalSize is true if there either isn't a gap at the end of the section or if an extra allocated section is added that isn't in a segment.

[jhenderson]

test/tools/llvm-objcopy/two-seg-remove-end.test

60

It's not particularly important, but this and several of the other tests look like they're not testing for the correct number of digits in the address.

tools/llvm-objcopy/Object.cpp

708

I don't think you need to repeat "This function" at the start of the second sentence. Simply "It", or combining the two sentences would suffice.

739

Please comment what exactly is returned here. By my understanding, it will be the end offset of the last segment, if there are no sections outside segments, or the end offset of the last section.

770

it's -> its

852–863

Sounds fair. Thinking about it, this loop is basically just a std::copy_if followed by a std::transform.

868

Okay, makes sense. Please add a comment to this effect.

903

Ok, I understand. I still think the comment needs rewriting though. How about something like:

"Now that every section has been laid out we just need to compute the total file size. This might not be the same as the offset returned by LayoutSections, because we want to truncate the last segment to the end of its last section, to match GNU objcopy's behaviour."

905

const auto &

[jakehehrlich]

1. Fixed some comments
2. Added reference in loop over AllocatedSections

[jakehehrlich]

tools/llvm-objcopy/Object.cpp

905

AllocatedSections is a vector of pointers so we shouldn't need to add that. Is there a reason in to add the reference that I'm not aware of? It's not as if it harms anything by being there, I'm just curious. I went ahead and added it in case there's some reason it should be done so that this doesn't block an LGTM.

[jhenderson]

Still quite a few comment changes, and one or two that were missed, so I'm going to ask you to update them again, before giving a LGTM. Sorry!

tools/llvm-objcopy/Object.cpp

708

Okay, my bad. That didn't come out quite as I meant/felt it would, and it now feels a little unnatural. As I believe that the LLVM standard is for full sentence, that first sentence needs to be either "This function finds a..." or "Find a...". Also, you've got "of list" when you need "list of".

Depending on which of the two suggestions you use for the first sentence, you should change the start of the second and third sentences accordingly. If the first approach is used, don't change the second sentence and merge the two: "It assumes ... OrderSegments, and returns ...". If using the second approach, you need to put "This function" instead of the first "It". Optionally, you can then fold the two sentences together as in the previous case.

739

You've used too many "returns" in this comment now. Should be "It returns either the ... ParentSegment, or an offset ..."

868

Ping here.

903

Ping here.

905

Good point, I missed this. Feel free to remove it again, if you prefer.

[jakehehrlich]

Fixed more comments

[jhenderson]

Three more minor points, but LGTM with the suggested fixes.

tools/llvm-objcopy/Object.cpp

708

OrderSegments and It returns -> OrderSegments and returns

736

layout of -> layout for

869

stranges -> strange

[jakehehrlich]

Whoops! Turns out I added a bug in this change where certain kinds of cycles in ParentSegment could form because there wasn't a check to make sure that compareSegments(Parent, Child) held if Child->ParentSegment was nullptr. This caused -O binary to fail in our kernel dump. I fixed this and added a test case. Additionally the way I modified the first segment wasn't exactly right.

@jhenderson could you please review the additional changes I made?

1. Added additional compareSegment check to make sure that bad ParentSegments are not set
2. Set alignment to zero because sometimes the first used segment won't even have an OriginalOffset and Align that let's us place it at offset 0x0 but we need to force that to happen regardless. This is also checked by the new test.

[jhenderson]

Oops, I should have spotted that as well.

Although I think the alignment change is tested as part of parent-loop-check.test, I think it should have its own explicit test, to make it easier to see what is being tested. Please could you add one.

test/tools/llvm-objcopy/parent-loop-check.test

1 ↗

(On Diff #122762)

Could you add a comment to the top of this test case to explain what is special about this case, please.

[jakehehrlich]

1. Added test for setting alignment of first segment to zero
2. Added description of the parent cycle test.

[jhenderson]

As mentioned inline, I'd like the new test renaming slightly, but otherwise LGTM.

test/tools/llvm-objcopy/binary-align-first.test

1 ↗

(On Diff #122896)

The test itself is fine, but I think it needs a different name, and a brief comment explaining the purpose of the test (i.e. that the first segment is always at offset zero). Possible new test name: binary-first-seg-offset-zero.test (but other names are aacceptable, as long as they convey some form of meaning).

test/tools/llvm-objcopy/binary-remove-all-but-one.test

+# RUN: yaml2obj %s -o %t
+# RUN: llvm-objcopy -R .text -R .text3 -O binary %t %t2
+# RUN: od -Ax -t x1 %t2 | FileCheck %s
+# RUN: wc -c %t2 | FileCheck %s -check-prefix=SIZE
+
+!ELF
+FileHeader:
+  Class:           ELFCLASS64
+  Data:            ELFDATA2LSB
+  Type:            ET_EXEC
+  Machine:         EM_X86_64
+Sections:
+  - Name:            .text
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x1000
+    AddressAlign:    0x1000
+    Content:         "c3c3c3c3"
+    Size:            0x1000
+  - Name:            .text2
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x2000
+    AddressAlign:    0x1000
+    Content:         "DEADBEEF"
+    Size:            0x1000
+  - Name:            .text3
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x3000
+    AddressAlign:    0x1000
+    Content:         "32323232"
+    Size:            0x1000
+ProgramHeaders:
+  - Type: PT_LOAD
+    Flags: [ PF_R ]
+    VAddr: 0x1000
+    PAddr: 0x1000
+    Sections:
+      - Section: .text
+      - Section: .text2
+      - Section: .text3
+
+# CHECK: 000000 de ad be ef
+
+# SIZE: 4096

test/tools/llvm-objcopy/binary-remove-end.test

+# RUN: yaml2obj %s -o %t
+# RUN: llvm-objcopy -R .text3 -O binary %t %t2
+# RUN: od -Ax -v -t x1 %t2 | FileCheck %s
+# RUN: wc -c %t2 | FileCheck %s -check-prefix=SIZE
+
+!ELF
+FileHeader:
+  Class:           ELFCLASS64
+  Data:            ELFDATA2LSB
+  Type:            ET_EXEC
+  Machine:         EM_X86_64
+Sections:
+  - Name:            .text
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x1000
+    AddressAlign:    0x1000
+    Content:         "c3c3c3c3"
+    Size:            0x1000
+  - Name:            .text2
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x2000
+    AddressAlign:    0x1000
+    Content:         "DEADBEEF"
+    Size:            0x1000
+  - Name:            .text3
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x3000
+    AddressAlign:    0x1000
+    Content:         "32323232"
+    Size:            0x1000
+ProgramHeaders:
+  - Type: PT_LOAD
+    Flags: [ PF_R ]
+    VAddr: 0x1000
+    PAddr: 0x1000
+    Sections:
+      - Section: .text
+      - Section: .text2
+      - Section: .text3
+
+# CHECK: 000000 c3 c3 c3 c3
+# CHECK: 001000 de ad be ef
+
+# SIZE: 8192

test/tools/llvm-objcopy/binary-remove-middle.test

+# RUN: yaml2obj %s -o %t
+# RUN: llvm-objcopy -R .text2 -O binary %t %t2
+# RUN: od -Ax -v -t x1 %t2 | FileCheck %s
+# RUN: wc -c %t2 | FileCheck %s -check-prefix=SIZE
+
+!ELF
+FileHeader:
+  Class:           ELFCLASS64
+  Data:            ELFDATA2LSB
+  Type:            ET_EXEC
+  Machine:         EM_X86_64
+Sections:
+  - Name:            .text
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x1000
+    AddressAlign:    0x1000
+    Content:         "c3c3c3c3"
+    Size:            0x1000
+  - Name:            .text2
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x2000
+    AddressAlign:    0x1000
+    Content:         "DEADBEEF"
+    Size:            0x1000
+  - Name:            .text3
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x3000
+    AddressAlign:    0x1000
+    Content:         "32323232"
+    Size:            0x1000
+ProgramHeaders:
+  - Type: PT_LOAD
+    Flags: [ PF_R ]
+    VAddr: 0x1000
+    PAddr: 0x1000
+    Sections:
+      - Section: .text
+      - Section: .text2
+      - Section: .text3
+
+# CHECK: 000000 c3 c3 c3 c3
+# CHECK: 001000 00 00 00 00
+# CHECK: 002000 32 32 32 32
+
+# SIZE: 12288

test/tools/llvm-objcopy/two-seg-remove-end.test

+# RUN: yaml2obj %s -o %t
+# RUN: llvm-objcopy -R .text4 -O binary %t %t2
+# RUN: od -Ax -v -t x1 %t2 | FileCheck %s
+# RUN: wc -c %t2 | FileCheck %s -check-prefix=SIZE
+
+!ELF
+FileHeader:
+  Class:           ELFCLASS64
+  Data:            ELFDATA2LSB
+  Type:            ET_EXEC
+  Machine:         EM_X86_64
+Sections:
+  - Name:            .text
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x1000
+    AddressAlign:    0x1000
+    Content:         "c3c3c3c3"
+    Size:            0x1000
+  - Name:            .text2
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x2000
+    AddressAlign:    0x1000
+    Content:         "DEADBEEF"
+    Size:            0x1000
+  - Name:            .text3
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x3000
+    AddressAlign:    0x1000
+    Content:         "32323232"
+    Size:            0x1000
+  - Name:            .text4
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x4000
+    AddressAlign:    0x1000
+    Content:         "FFFFFFFF"
+    Size:            0x1000
+ProgramHeaders:
+  - Type: PT_LOAD
+    Flags: [ PF_R ]
+    VAddr: 0x1000
+    PAddr: 0x1000
+    Sections:
+      - Section: .text
+      - Section: .text2
+  - Type: PT_LOAD
+    Flags: [ PF_R ]
+    VAddr: 0x1000
+    PAddr: 0x1000
+    Sections:
+      - Section: .text3
+      - Section: .text4
+
+# CHECK: 000000 c3 c3 c3 c3
+# CHECK: 001000 de ad be ef
+# CHECK: 002000 32 32 32 32
+
+# SIZE: 12288

test/tools/llvm-objcopy/two-seg-remove-first.test

+# RUN: yaml2obj %s -o %t
+# RUN: llvm-objcopy -R .text -O binary %t %t2
+# RUN: od -Ax -v -t x1 %t2 | FileCheck %s
+# RUN: wc -c %t2 | FileCheck %s -check-prefix=SIZE
+
+!ELF
+FileHeader:
+  Class:           ELFCLASS64
+  Data:            ELFDATA2LSB
+  Type:            ET_EXEC
+  Machine:         EM_X86_64
+Sections:
+  - Name:            .text
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x1000
+    AddressAlign:    0x1000
+    Content:         "c3c3c3c3"
+    Size:            0x1000
+  - Name:            .text2
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x2000
+    AddressAlign:    0x1000
+    Content:         "DEADBEEF"
+    Size:            0x1000
+  - Name:            .text3
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x3000
+    AddressAlign:    0x1000
+    Content:         "32323232"
+    Size:            0x1000
+  - Name:            .text4
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x4000
+    AddressAlign:    0x1000
+    Content:         "FFFFFFFF"
+    Size:            0x1000
+ProgramHeaders:
+  - Type: PT_LOAD
+    Flags: [ PF_R ]
+    VAddr: 0x1000
+    PAddr: 0x1000
+    Sections:
+      - Section: .text
+      - Section: .text2
+  - Type: PT_LOAD
+    Flags: [ PF_R ]
+    VAddr: 0x1000
+    PAddr: 0x1000
+    Sections:
+      - Section: .text3
+      - Section: .text4
+
+# CHECK: 000000 de ad be ef
+# CHECK: 001000 32 32 32 32
+# CHECK: 002000 ff ff ff ff
+
+# SIZE: 12288

test/tools/llvm-objcopy/two-seg-remove-third-sec.test

+# RUN: yaml2obj %s -o %t
+# RUN: llvm-objcopy -R .text3 -O binary %t %t2
+# RUN: od -Ax -v -t x1 %t2 | FileCheck %s
+# RUN: wc -c %t2 | FileCheck %s -check-prefix=SIZE
+
+!ELF
+FileHeader:
+  Class:           ELFCLASS64
+  Data:            ELFDATA2LSB
+  Type:            ET_EXEC
+  Machine:         EM_X86_64
+Sections:
+  - Name:            .text
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x1000
+    AddressAlign:    0x1000
+    Content:         "c3c3c3c3"
+    Size:            0x1000
+  - Name:            .text2
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x2000
+    AddressAlign:    0x1000
+    Content:         "DEADBEEF"
+    Size:            0x1000
+  - Name:            .text3
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x3000
+    AddressAlign:    0x1000
+    Content:         "32323232"
+    Size:            0x1000
+  - Name:            .text4
+    Type:            SHT_PROGBITS
+    Flags:           [ SHF_ALLOC, SHF_EXECINSTR ]
+    Address:         0x4000
+    AddressAlign:    0x1000
+    Content:         "FFFFFFFF"
+    Size:            0x1000
+ProgramHeaders:
+  - Type: PT_LOAD
+    Flags: [ PF_R ]
+    VAddr: 0x1000
+    PAddr: 0x1000
+    Sections:
+      - Section: .text
+      - Section: .text2
+  - Type: PT_LOAD
+    Flags: [ PF_R ]
+    VAddr: 0x1000
+    PAddr: 0x1000
+    Sections:
+      - Section: .text3
+      - Section: .text4
+
+# CHECK: 000000 c3 c3 c3 c3
+# CHECK: 001000 de ad be ef
+# CHECK: 002000 00 00 00 00
+# CHECK: 003000 ff ff ff ff
+
+# SIZE: 16384

tools/llvm-objcopy/Object.cpp

 // range.
 static bool segmentOverlapsSegment(const Segment &Child,
                                    const Segment &Parent) {
 
   return Parent.OriginalOffset <= Child.OriginalOffset &&
          Parent.OriginalOffset + Parent.FileSize > Child.OriginalOffset;
 }
 
+static bool compareSegments(const Segment *A, const Segment *B) {
+  // Any segment without a parent segment should come before a segment
+  // that has a parent segment.
+  if (A->OriginalOffset < B->OriginalOffset)
+    return true;
+  if (A->OriginalOffset > B->OriginalOffset)
+    return false;
+  return A->Index < B->Index;
+}
+
 template <class ELFT>
 void Object<ELFT>::readProgramHeaders(const ELFFile<ELFT> &ElfFile) {
   uint32_t Index = 0;
   for (const auto &Phdr : unwrapOrError(ElfFile.program_headers())) {
     ArrayRef<uint8_t> Data{ElfFile.base() + Phdr.p_offset,
                            (size_t)Phdr.p_filesz};
     Segments.emplace_back(llvm::make_unique<Segment>(Data));
     Segment &Seg = *Segments.back();
void Object<ELFT>::readProgramHeaders(const ELFFile<ELFT> &ElfFile) {
   // segments.
   for (auto &Child : Segments) {
     for (auto &Parent : Segments) {
       // Every segment will overlap with itself but we don't want a segment to
       // be it's own parent so we avoid that situation.
       if (&Child != &Parent && segmentOverlapsSegment(*Child, *Parent)) {
         // We want a canonical "most parental" segment but this requires
         // inspecting the ParentSegment.
-        if (Child->ParentSegment != nullptr) {
-          if (Child->ParentSegment->OriginalOffset > Parent->OriginalOffset) {
-            Child->ParentSegment = Parent.get();
-          } else if (Child->ParentSegment->Index > Parent->Index) {
-            // They must have equal OriginalOffsets so we need to disambiguate.
-            // To decide which is the parent we'll choose the one with the
-            // higher index.
-            Child->ParentSegment = Parent.get();
-          }
-        } else {
+        if (Child->ParentSegment == nullptr ||
+            compareSegments(Parent.get(), Child->ParentSegment)) {
           Child->ParentSegment = Parent.get();
         }
       }
     }
   }
 }
 
 template <class ELFT>
auto Diff =
       static_cast<int64_t>(Addr % Align) - static_cast<int64_t>(Offset % Align);
   // We only want to add to Offset, however, so if Diff < 0 we can add Align and
   // (Offset + Diff) & -Align == Addr & -Align will still hold.
   if (Diff < 0)
     Diff += Align;
   return Offset + Diff;
 }
 
-template <class ELFT> void ELFObject<ELFT>::assignOffsets() {
-  // We need a temporary list of segments that has a special order to it
-  // so that we know that anytime ->ParentSegment is set that segment has
-  // already had it's offset properly set.
-  std::vector<Segment *> OrderedSegments;
-  for (auto &Segment : this->Segments)
-    OrderedSegments.push_back(Segment.get());
-  auto CompareSegments = [](const Segment *A, const Segment *B) {
-    // Any segment without a parent segment should come before a segment
-    // that has a parent segment.
-    if (A->OriginalOffset < B->OriginalOffset)
-      return true;
-    if (A->OriginalOffset > B->OriginalOffset)
-      return false;
-    return A->Index < B->Index;
-  };
-  std::stable_sort(std::begin(OrderedSegments), std::end(OrderedSegments),
-                   CompareSegments);
-  // The size of ELF + program headers will not change so it is ok to assume
-  // that the first offset of the first segment is a good place to start
-  // outputting sections. This covers both the standard case and the PT_PHDR
-  // case.
-  uint64_t Offset;
-  if (!OrderedSegments.empty()) {
-    Offset = OrderedSegments[0]->Offset;
-  } else {
-    Offset = sizeof(Elf_Ehdr);
-  }
+// Orders segments such that if x = y->ParentSegment then y comes before x.
+static void OrderSegments(std::vector<Segment *> &Segments) {
+  std::stable_sort(std::begin(Segments), std::end(Segments), compareSegments);
+
+// This function finds a consistent layout for a list of segments starting from
+// an Offset. It assumes that Segments have been sorted by OrderSegments and It
+// returns an Offset one past the end of the last segment.
+static uint64_t LayoutSegments(std::vector<Segment *> &Segments,
+                               uint64_t Offset) {
+  assert(std::is_sorted(std::begin(Segments), std::end(Segments),
+                        compareSegments));
   // The only way a segment should move is if a section was between two
   // segments and that section was removed. If that section isn't in a segment
   // then it's acceptable, but not ideal, to simply move it to after the
   // segments. So we can simply layout segments one after the other accounting
   // for alignment.
-  for (auto &Segment : OrderedSegments) {
+  for (auto &Segment : Segments) {
     // We assume that segments have been ordered by OriginalOffset and Index
     // such that a parent segment will always come before a child segment in
     // OrderedSegments. This means that the Offset of the ParentSegment should
     // already be set and we can set our offset relative to it.
     if (Segment->ParentSegment != nullptr) {
       auto Parent = Segment->ParentSegment;
       Segment->Offset =
           Parent->Offset + Segment->OriginalOffset - Parent->OriginalOffset;
     } else {
       Offset = alignToAddr(Offset, Segment->VAddr, Segment->Align);
       Segment->Offset = Offset;
     }
     Offset = std::max(Offset, Segment->Offset + Segment->FileSize);
   }
+  return Offset;
+}
+
+// This function finds a consistent layout of a list of sections. It assumes
+// that the ->ParentSegment of each section has already been laid out. The
+// supplied starting Offset is used for the starting offset of any section that
+// does not have a ParentSegment. It returns either the offset given if all
+// sections had a ParentSegment or an offset one past the last section if there
+// was a section that didn't have a ParentSegment.
+template <class SecPtr>
+static uint64_t LayoutSections(std::vector<SecPtr> &Sections, uint64_t Offset) {
   // Now the offset of every segment has been set we can assign the offsets
   // of each section. For sections that are covered by a segment we should use
   // the segment's original offset and the section's original offset to compute
   // the offset from the start of the segment. Using the offset from the start
   // of the segment we can assign a new offset to the section. For sections not
   // covered by segments we can just bump Offset to the next valid location.
   uint32_t Index = 1;
-  for (auto &Section : this->Sections) {
+  for (auto &Section : Sections) {
     Section->Index = Index++;
     if (Section->ParentSegment != nullptr) {
       auto Segment = Section->ParentSegment;
       Section->Offset =
           Segment->Offset + (Section->OriginalOffset - Segment->OriginalOffset);
     } else {
       Offset = alignTo(Offset, Section->Align == 0 ? 1 : Section->Align);
       Section->Offset = Offset;
       if (Section->Type != SHT_NOBITS)
         Offset += Section->Size;
     }
   }
+  return Offset;
+}
 
-  if (this->WriteSectionHeaders) {
-    Offset = alignTo(Offset, sizeof(typename ELFT::Addr));
+template <class ELFT> void ELFObject<ELFT>::assignOffsets() {
+  // We need a temporary list of segments that has a special order to it
+  // so that we know that anytime ->ParentSegment is set that segment has
+  // already had its offset properly set.
+  std::vector<Segment *> OrderedSegments;
+  for (auto &Segment : this->Segments)
+    OrderedSegments.push_back(Segment.get());
+  OrderSegments(OrderedSegments);
+  // The size of ELF + program headers will not change so it is ok to assume
+  // that the first offset of the first segment is a good place to start
+  // outputting sections. This covers both the standard case and the PT_PHDR
+  // case.
+  uint64_t Offset;
+  if (!OrderedSegments.empty()) {
+    Offset = OrderedSegments[0]->Offset;
+  } else {
+    Offset = sizeof(Elf_Ehdr);
   }
+  Offset = LayoutSegments(OrderedSegments, Offset);
+  Offset = LayoutSections(this->Sections, Offset);
+  // If we need to write the section header table out then we need to align the
+  // Offset so that SHOffset is valid.
+  if (this->WriteSectionHeaders)
+    Offset = alignTo(Offset, sizeof(typename ELFT::Addr));
   this->SHOffset = Offset;
 }
 
 template <class ELFT> size_t ELFObject<ELFT>::totalSize() const {
   // We already have the section header offset so we can calculate the total
   // size by just adding up the size of each section header.
   auto NullSectionSize = this->WriteSectionHeaders ? sizeof(Elf_Shdr) : 0;
   return this->SHOffset + this->Sections.size() * sizeof(Elf_Shdr) +

 }
 
 template <class ELFT> size_t BinaryObject<ELFT>::totalSize() const {
   return TotalSize;
 }
 
 template <class ELFT>
 void BinaryObject<ELFT>::write(FileOutputBuffer &Out) const {
-  for (auto &Segment : this->Segments) {
-    // GNU objcopy does not output segments that do not cover a section. Such
-    // segments can sometimes be produced by LLD due to how LLD handles PT_PHDR.
-    if (Segment->Type == PT_LOAD && Segment->firstSection() != nullptr) {
-      Segment->writeSegment(Out);
-    }
+  for (auto &Section : this->Sections) {
+    if ((Section->Flags & SHF_ALLOC) == 0)
+      continue;
+    Section->writeSection(Out);
   }
 }
 
 template <class ELFT> void BinaryObject<ELFT>::finalize() {
-  // Put all segments in offset order.
-  auto CompareSegments = [](const SegPtr &A, const SegPtr &B) {
-    return A->Offset < B->Offset;
-  };
-  std::sort(std::begin(this->Segments), std::end(this->Segments),
-            CompareSegments);
-
-  uint64_t Offset = 0;
-  for (auto &Segment : this->Segments) {
-    if (Segment->Type == llvm::ELF::PT_LOAD &&
-        Segment->firstSection() != nullptr) {
-      Offset = alignToAddr(Offset, Segment->VAddr, Segment->Align);
-      Segment->Offset = Offset;
-      Offset += Segment->FileSize;
-    }
+  // TODO: Create a filter range to construct OrderedSegments from so that this
+  // code can be deduped with assignOffsets above. This should also solve the
+  // todo below for LayoutSections.
+  // We need a temporary list of segments that has a special order to it
+  // so that we know that anytime ->ParentSegment is set that segment has
+  // already had it's offset properly set. We only want to consider the segments
+  // that will affect layout of allocated sections so we only add those.
+  std::vector<Segment *> OrderedSegments;
+  for (auto &Section : this->Sections) {
+    if ((Section->Flags & SHF_ALLOC) != 0 &&
+        Section->ParentSegment != nullptr) {
+      OrderedSegments.push_back(Section->ParentSegment);
+    }
+  }
+  OrderSegments(OrderedSegments);
+  // Because we add a ParentSegment for each section we might have duplicate
+  // segments in OrderedSegments. If there were duplicates then LayoutSegments
+  // would do very stranges things.
+  auto End =
+      std::unique(std::begin(OrderedSegments), std::end(OrderedSegments));
+  OrderedSegments.erase(End, std::end(OrderedSegments));
+
+  // Modify the first segment so that there is no gap at the start. This allows
+  // our layout algorithm to proceed as expected while not out writing out the
+  // gap at the start.
+  if (!OrderedSegments.empty()) {
+    auto Seg = OrderedSegments[0];
+    auto Sec = Seg->firstSection();
+    auto Diff = Sec->OriginalOffset - Seg->OriginalOffset;
+    Seg->OriginalOffset += Diff;
+    // The size needs to be shrunk as well
+    Seg->FileSize -= Diff;
+    Seg->MemSize -= Diff;
+    // The VAddr needs to be adjusted so that the alignment is correct as well
+    Seg->VAddr += Diff;
+    Seg->PAddr = Seg->VAddr;
+  }
+
+  uint64_t Offset = LayoutSegments(OrderedSegments, 0);
+
+  // TODO: generalize LayoutSections to take a range. Pass a special range
+  // constructed from an iterator that skips values for which a predicate does
+  // not hold. Then pass such a range to LayoutSections instead of constructing
+  // AllocatedSections here.
+  std::vector<SectionBase *> AllocatedSections;
+  for (auto &Section : this->Sections) {
+    if ((Section->Flags & SHF_ALLOC) == 0)
+      continue;
+    AllocatedSections.push_back(Section.get());
+  LayoutSections(AllocatedSections, Offset);
+
+  // Now that every section has been laid out we just need to compute the total
+  // file size. This might not be the same as the offset returned by
+  // LayoutSections, because we want to truncate the last segment to the end of
+  // its last section, to match GNU objcopy's behaviour.
+  TotalSize = 0;
+  for (const auto &Section : AllocatedSections) {
+    if (Section->Type != SHT_NOBITS)
+      TotalSize = std::max(TotalSize, Section->Offset + Section->Size);
   }
-  TotalSize = Offset;
 }
 
 namespace llvm {
 
 template class Object<ELF64LE>;
 template class Object<ELF64BE>;
 template class Object<ELF32LE>;
 template class Object<ELF32BE>;