Index: ELF/Writer.cpp
===================================================================
--- ELF/Writer.cpp
+++ ELF/Writer.cpp
@@ -19,6 +19,8 @@
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Support/StringSaver.h"
 
+#include <list>
+
 using namespace llvm;
 using namespace llvm::ELF;
 using namespace llvm::object;
@@ -94,6 +96,16 @@
   uintX_t FileSize;
   uintX_t SectionHeaderOff;
 
+  struct Phdr {
+    Phdr(unsigned Type, unsigned Flags) : Type(Type), Flags(Flags) {}
+    unsigned Type;
+    unsigned Flags;
+    std::list<OutputSectionBase<ELFT> *> OutSections;
+  };
+
+  using PhdrMap = std::list<std::unique_ptr<Phdr>>;
+  PhdrMap scanHeaders();
+
   llvm::StringMap<llvm::StringRef> InputToOutputSection;
 
   // Flag to force GOT to be in output if we have relocations
@@ -1037,131 +1049,217 @@
   GnuRelroPhdr->p_memsz = VA - Cur->p_vaddr;
 }
 
-// Visits all sections to create PHDRs and to assign incremental,
-// non-overlapping addresses to output sections.
-template <class ELFT> void Writer<ELFT>::assignAddresses() {
-  uintX_t VA = Target->getVAStart() + sizeof(Elf_Ehdr);
-  uintX_t FileOff = sizeof(Elf_Ehdr);
-
-  // Calculate and reserve the space for the program header first so that
-  // the first section can start right after the program header.
-  Phdrs.resize(getPhdrsNum());
-  size_t PhdrSize = sizeof(Elf_Phdr) * Phdrs.size();
+template <class ELFT>
+typename Writer<ELFT>::PhdrMap Writer<ELFT>::scanHeaders() {
+  PhdrMap Phdrs;
+  auto AddIf = [this, &Phdrs](unsigned Type, bool C, unsigned Flags = 0) {
+    if (!C)
+      return (Phdr *)nullptr;
+    std::unique_ptr<Phdr> U = std::make_unique<Phdr>(Type, Flags);
+    return Phdrs.emplace(Phdrs.end(), std::move(U))->get();
+  };
 
   // The first phdr entry is PT_PHDR which describes the program header itself.
-  setPhdr(&Phdrs[0], PT_PHDR, PF_R, FileOff, VA, PhdrSize, /*Align=*/8);
-  FileOff += PhdrSize;
-  VA += PhdrSize;
+  AddIf(PT_PHDR, true, PF_R);
 
   // PT_INTERP must be the second entry if exists.
-  int PhdrIdx = 0;
-  Elf_Phdr *Interp = nullptr;
-  if (needsInterpSection())
-    Interp = &Phdrs[++PhdrIdx];
+  AddIf(PT_INTERP, needsInterpSection());
 
   // Add the first PT_LOAD segment for regular output sections.
-  setPhdr(&Phdrs[++PhdrIdx], PT_LOAD, PF_R, 0, Target->getVAStart(), FileOff,
-          Target->getPageSize());
-
-  Elf_Phdr GnuRelroPhdr = {};
-  Elf_Phdr TlsPhdr{};
-  bool RelroAligned = false;
-  uintX_t ThreadBssOffset = 0;
-  // Create phdrs as we assign VAs and file offsets to all output sections.
+  uintX_t Flags = PF_R;
+  Phdr *Load = AddIf(PT_LOAD, true, Flags);
+  std::unique_ptr<Phdr> TlsHdr;
   for (OutputSectionBase<ELFT> *Sec : OutputSections) {
-    Elf_Phdr *PH = &Phdrs[PhdrIdx];
     if (needsPhdr<ELFT>(Sec)) {
-      uintX_t Flags = toPhdrFlags(Sec->getFlags());
-      bool InRelRo = Config->ZRelro && (Flags & PF_W) && isRelroSection(Sec);
-      bool FirstNonRelRo = GnuRelroPhdr.p_type && !InRelRo && !RelroAligned;
-      if (FirstNonRelRo || PH->p_flags != Flags) {
-        VA = alignTo(VA, Target->getPageSize());
-        FileOff = alignTo(FileOff, Target->getPageSize());
-        if (FirstNonRelRo)
-          RelroAligned = true;
-      }
-
-      if (PH->p_flags != Flags) {
-        // Flags changed. Create a new PT_LOAD.
-        PH = &Phdrs[++PhdrIdx];
-        uint32_t PTType = (Config->EMachine != EM_AMDGPU) ? (uint32_t)PT_LOAD
-                                                          : getAmdgpuPhdr(Sec);
-        setPhdr(PH, PTType, Flags, FileOff, VA, 0, Target->getPageSize());
+      // If flags changed then we want new load segment.
+      uintX_t NewFlags = toPhdrFlags(Sec->getFlags());
+      if (Flags != NewFlags) {
+        uint32_t LoadType = (Config->EMachine == EM_AMDGPU) ? getAmdgpuPhdr(Sec)
+                                                            : (uint32_t)PT_LOAD;
+        Load = AddIf(LoadType, true, NewFlags);
+        Flags = NewFlags;
       }
-
+      // If we meet TLS section then we create TLS header
+      // and put all TLS sections inside for futher use when
+      // assign addresses.
       if (Sec->getFlags() & SHF_TLS) {
-        if (!TlsPhdr.p_vaddr)
-          setPhdr(&TlsPhdr, PT_TLS, PF_R, FileOff, VA, 0, Sec->getAlign());
-        if (Sec->getType() != SHT_NOBITS)
-          VA = alignTo(VA, Sec->getAlign());
-        uintX_t TVA = alignTo(VA + ThreadBssOffset, Sec->getAlign());
-        Sec->setVA(TVA);
-        TlsPhdr.p_memsz += Sec->getSize();
-        if (Sec->getType() == SHT_NOBITS) {
-          ThreadBssOffset = TVA - VA + Sec->getSize();
-        } else {
-          TlsPhdr.p_filesz += Sec->getSize();
-          VA += Sec->getSize();
-        }
-        TlsPhdr.p_align = std::max<uintX_t>(TlsPhdr.p_align, Sec->getAlign());
-      } else {
-        VA = alignTo(VA, Sec->getAlign());
-        Sec->setVA(VA);
-        VA += Sec->getSize();
-        if (InRelRo)
-          updateRelro(PH, &GnuRelroPhdr, VA);
+        if (!TlsHdr)
+          TlsHdr = std::make_unique<Phdr>(PT_TLS, PF_R);
+        TlsHdr->OutSections.push_back(Sec);
       }
     }
-
-    FileOff = alignTo(FileOff, Sec->getAlign());
-    Sec->setFileOffset(FileOff);
-    if (Sec->getType() != SHT_NOBITS)
-      FileOff += Sec->getSize();
-    if (needsPhdr<ELFT>(Sec)) {
-      PH->p_filesz = FileOff - PH->p_offset;
-      PH->p_memsz = VA - PH->p_vaddr;
-    }
+    Load->OutSections.push_back(Sec);
   }
 
-  if (TlsPhdr.p_vaddr) {
-    // The TLS pointer goes after PT_TLS. At least glibc will align it,
-    // so round up the size to make sure the offsets are correct.
-    TlsPhdr.p_memsz = alignTo(TlsPhdr.p_memsz, TlsPhdr.p_align);
-    Phdrs[++PhdrIdx] = TlsPhdr;
-    Out<ELFT>::TlsPhdr = &Phdrs[PhdrIdx];
-  }
+  // TLS header.
+  if (TlsHdr)
+    Phdrs.push_back(std::move(TlsHdr));
 
   // Add an entry for .dynamic.
-  if (isOutputDynamic()) {
-    Elf_Phdr *PH = &Phdrs[++PhdrIdx];
-    PH->p_type = PT_DYNAMIC;
-    copyPhdr(PH, Out<ELFT>::Dynamic);
-  }
+  AddIf(PT_DYNAMIC, isOutputDynamic());
 
-  if (HasRelro) {
-    Elf_Phdr *PH = &Phdrs[++PhdrIdx];
-    *PH = GnuRelroPhdr;
-  }
+  // PT_GNU_RELRO includes all sections that should be marked as
+  // read-only by dynamic linker after proccessing relocations.
+  AddIf(PT_GNU_RELRO, HasRelro);
 
-  if (Out<ELFT>::EhFrameHdr->Live) {
-    Elf_Phdr *PH = &Phdrs[++PhdrIdx];
-    PH->p_type = PT_GNU_EH_FRAME;
-    copyPhdr(PH, Out<ELFT>::EhFrameHdr);
-  }
+  // PT_GNU_EH_FRAME is a special section pointing on .eh_frame_hdr.
+  AddIf(PT_GNU_EH_FRAME, Out<ELFT>::EhFrameHdr->Live);
 
   // PT_GNU_STACK is a special section to tell the loader to make the
   // pages for the stack non-executable.
-  if (!Config->ZExecStack) {
-    Elf_Phdr *PH = &Phdrs[++PhdrIdx];
-    PH->p_type = PT_GNU_STACK;
-    PH->p_flags = PF_R | PF_W;
+  AddIf(PT_GNU_STACK, !Config->ZExecStack, PF_R | PF_W);
+
+  return Phdrs;
+}
+
+// Visits all headers in PhdrMap and assigns the adresses to
+// the output sections. Also creates common and special headers.
+template <class ELFT> void Writer<ELFT>::assignAddresses() {
+  // We know the final amount of program headers here.
+  PhdrMap Map = scanHeaders();
+  Phdrs.resize(Map.size());
+
+  uintX_t VA = Target->getVAStart() + sizeof(Elf_Ehdr);
+  uintX_t FileOff = sizeof(Elf_Ehdr);
+  uintX_t TlsVA = 0;
+  uintX_t TlsFileOff = 0;
+  Elf_Phdr *InterpHdr = nullptr;
+  Elf_Phdr GnuRelroPhdr = {};
+  bool RelroAligned = false;
+  bool FirstLoad = true;
+  uintX_t I = 0;
+  for (auto &PHdr : Map) {
+    Elf_Phdr *PH = &Phdrs[I++];
+    PH->p_type = PHdr->Type;
+    switch (PHdr->Type) {
+    case PT_LOAD:
+    case PT_AMDGPU_HSA_LOAD_CODE_AGENT:
+    case PT_AMDGPU_HSA_LOAD_GLOBAL_PROGRAM: {
+      // Add the first load segment for regular output sections.
+      if (FirstLoad)
+        setPhdr(PH, PHdr->Type, PHdr->Flags, 0, Target->getVAStart(), FileOff,
+                Target->getPageSize());
+
+      bool CreateLoad = !FirstLoad;
+      bool isW = PHdr->Flags & PF_W;
+      for (OutputSectionBase<ELFT> *Sec : PHdr->OutSections) {
+        bool InRelRo = Config->ZRelro && isW && isRelroSection(Sec);
+        bool FirstNonRelRo = needsPhdr<ELFT>(Sec) && !RelroAligned &&
+                             GnuRelroPhdr.p_type && !InRelRo;
+        // There are two situations when we want to align to page boundaries.
+        // 1) If this is a first non relro section after relro sequence.
+        //    We do this because we dont want any other sections except relro
+        //    were marked as read only. And that is why them should be placed on
+        //    a different memory page.
+        // 2) CreateLoad flag is set for all PT_LOADs except first one.
+        //    New PT_LOAD is usually created when access flags were changed.
+        //    So for each next PT_LOAD we create the new load header.
+        //    For first one it is created separately above.
+        if (FirstNonRelRo || CreateLoad) {
+          VA = alignTo(VA, Target->getPageSize());
+          FileOff = alignTo(FileOff, Target->getPageSize());
+          if (CreateLoad) {
+            setPhdr(PH, PHdr->Type, PHdr->Flags, FileOff, VA, 0,
+                    Target->getPageSize());
+            CreateLoad = false;
+          }
+          if (FirstNonRelRo)
+            RelroAligned = true;
+        }
+
+        if (needsPhdr<ELFT>(Sec)) {
+          // If we found TLS section we want to remember
+          // its VA and FileOff because we know that TLS block
+          // is continuos and we need that values for creating
+          // PT_TLS header below.
+          bool IsTls = Sec->getFlags() & SHF_TLS;
+          if (IsTls && !TlsVA) {
+            TlsVA = VA;
+            TlsFileOff = FileOff;
+          }
+          // Don't allocate VA space for TLS NOBITS sections. The PT_TLS PHDR is
+          // responsible for allocating space for them, not the PT_LOAD that
+          // contains the TLS initialization image.
+          if (!IsTls || Sec->getType() != SHT_NOBITS) {
+            VA = alignTo(VA, Sec->getAlign());
+            Sec->setVA(VA);
+            VA += Sec->getSize();
+          }
+          // If we are in relro then we need to update
+          // the GnuRelroPhdr header.
+          if (InRelRo)
+            updateRelro(PH, &GnuRelroPhdr, VA);
+        }
+
+        FileOff = alignTo(FileOff, Sec->getAlign());
+        Sec->setFileOffset(FileOff);
+        if (Sec->getType() != SHT_NOBITS)
+          FileOff += Sec->getSize();
+        if (needsPhdr<ELFT>(Sec)) {
+          PH->p_filesz = FileOff - PH->p_offset;
+          PH->p_memsz = VA - PH->p_vaddr;
+        }
+      }
+      FirstLoad = false;
+      break;
+    }
+    case PT_TLS: {
+      // We want to create PT_TLS header and assign proper
+      // addresses for all sections it includes.
+      uintX_t ThreadBssOffset = 0;
+      setPhdr(PH, PT_TLS, PHdr->Flags, TlsFileOff, TlsVA, 0, 0);
+      for (OutputSectionBase<ELFT> *Sec : PHdr->OutSections) {
+        if (Sec->getType() != SHT_NOBITS)
+          TlsVA = alignTo(TlsVA, Sec->getAlign());
+        uintX_t TVA = alignTo(TlsVA + ThreadBssOffset, Sec->getAlign());
+        Sec->setVA(TVA);
+
+        if (Sec->getType() == SHT_NOBITS) {
+          ThreadBssOffset = TVA - TlsVA + Sec->getSize();
+        } else {
+          TlsVA += Sec->getSize();
+          PH->p_filesz += Sec->getSize();
+        }
+
+        PH->p_memsz += Sec->getSize();
+        PH->p_align = std::max<uintX_t>(PH->p_align, Sec->getAlign());
+      }
+      // The TLS pointer goes after PT_TLS. At least glibc will align it,
+      // so round up the size to make sure the offsets are correct.
+      PH->p_memsz = alignTo(PH->p_memsz, PH->p_align);
+      Out<ELFT>::TlsPhdr = PH;
+      break;
+    }
+    case PT_PHDR: {
+      size_t PhdrSize = sizeof(Elf_Phdr) * Phdrs.size();
+      setPhdr(PH, PT_PHDR, PF_R, FileOff, VA, PhdrSize, /*Align=*/8);
+      FileOff += PhdrSize;
+      VA += PhdrSize;
+      break;
+    }
+    case PT_INTERP:
+      InterpHdr = PH;
+      break;
+    case PT_DYNAMIC:
+      copyPhdr(PH, Out<ELFT>::Dynamic);
+      break;
+    case PT_GNU_EH_FRAME:
+      copyPhdr(PH, Out<ELFT>::EhFrameHdr);
+      break;
+    case PT_GNU_RELRO:
+      *PH = GnuRelroPhdr;
+      break;
+    case PT_GNU_STACK:
+      PH->p_flags = PHdr->Flags;
+      break;
+    default:
+      assert("Unknown header type");
+    }
   }
 
   // Fix up PT_INTERP as we now know the address of .interp section.
-  if (Interp) {
-    Interp->p_type = PT_INTERP;
-    copyPhdr(Interp, Out<ELFT>::Interp);
-  }
+  if (InterpHdr)
+    copyPhdr(InterpHdr, Out<ELFT>::Interp);
 
   // Add space for section headers.
   SectionHeaderOff = alignTo(FileOff, ELFT::Is64Bits ? 8 : 4);
@@ -1172,37 +1270,6 @@
   ElfSym<ELFT>::End.st_value = VA;
 }
 
-// Returns the number of PHDR entries.
-template <class ELFT> int Writer<ELFT>::getPhdrsNum() const {
-  bool Tls = false;
-  int I = 2; // 2 for PT_PHDR and first PT_LOAD
-  if (needsInterpSection())
-    ++I;
-  if (isOutputDynamic())
-    ++I;
-  if (!Config->ZExecStack)
-    ++I;
-  uintX_t Last = PF_R;
-  for (OutputSectionBase<ELFT> *Sec : OutputSections) {
-    if (!needsPhdr<ELFT>(Sec))
-      continue;
-    if (Sec->getFlags() & SHF_TLS)
-      Tls = true;
-    uintX_t Flags = toPhdrFlags(Sec->getFlags());
-    if (Last != Flags) {
-      Last = Flags;
-      ++I;
-    }
-  }
-  if (Tls)
-    ++I;
-  if (HasRelro)
-    ++I;
-  if (Out<ELFT>::EhFrameHdr->Live)
-    ++I;
-  return I;
-}
-
 static uint32_t getELFFlags() {
   if (Config->EMachine != EM_MIPS)
     return 0;