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;
@@ -41,6 +43,20 @@
   void run();
 
 private:
+  struct Phdr {
+    Phdr(unsigned Type = 0, unsigned Flags = 0) : Type(Type), Flags(Flags) {}
+    unsigned Type;
+    unsigned Flags;
+    std::list<OutputSectionBase<ELFT> *> OutSections;
+  };
+  struct PhdrsContext {
+    std::vector<std::unique_ptr<Phdr>> PhdrList;
+    size_t FirstLoadId = (size_t)-1;
+    size_t TlsPhdrId = (size_t)-1;
+    size_t RelroPhdrId = (size_t)-1;
+  };
+  using PhdrMap = PhdrsContext;
+
   void copyLocalSymbols();
   void addReservedSymbols();
   void createSections();
@@ -53,7 +69,9 @@
   void scanRelocs(InputSection<ELFT> &C);
   void scanRelocs(InputSectionBase<ELFT> &S, const Elf_Shdr &RelSec);
   void updateRelro(Elf_Phdr *Cur, Elf_Phdr *GnuRelroPhdr, uintX_t VA);
+  PhdrMap scanHeaders();
   void assignAddresses();
+  void assignTlsAddresses(Elf_Phdr *PH, const Phdr &Map);
   void buildSectionMap();
   void fixAbsoluteSymbols();
   void openFile(StringRef OutputPath);
@@ -67,7 +85,6 @@
   bool isOutputDynamic() const {
     return !Symtab.getSharedFiles().empty() || Config->Shared;
   }
-  int getPhdrsNum() const;
 
   OutputSection<ELFT> *getBss();
   void addCommonSymbols(std::vector<DefinedCommon *> &Syms);
@@ -1065,136 +1082,209 @@
 void Writer<ELFT>::updateRelro(Elf_Phdr *Cur, Elf_Phdr *GnuRelroPhdr,
                                uintX_t VA) {
   if (!GnuRelroPhdr->p_type)
-    setPhdr(GnuRelroPhdr, PT_GNU_RELRO, PF_R, Cur->p_offset, Cur->p_vaddr,
+    setPhdr(GnuRelroPhdr, PT_GNU_RELRO, 0, Cur->p_offset, Cur->p_vaddr,
             VA - Cur->p_vaddr, 1 /*p_align*/);
   GnuRelroPhdr->p_filesz = VA - Cur->p_vaddr;
   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();
+// This method generates some kind of map where each phdr
+// is associated with zero or more output sections.
+// Output structure contains final list of phdrs used.
+// assignAddresses() will use this structure to assign addressed to sections.
+template <class ELFT>
+typename Writer<ELFT>::PhdrMap Writer<ELFT>::scanHeaders() {
+  PhdrMap Map;
+  auto AddHdr = [this, &Map](unsigned Type, unsigned Flags = 0) {
+    std::unique_ptr<Phdr> U = std::make_unique<Phdr>(Type, Flags);
+    return Map.PhdrList.emplace(Map.PhdrList.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;
+  AddHdr(PT_PHDR, PF_R);
 
   // PT_INTERP must be the second entry if exists.
-  int PhdrIdx = 0;
-  Elf_Phdr *Interp = nullptr;
   if (needsInterpSection())
-    Interp = &Phdrs[++PhdrIdx];
+    AddHdr(PT_INTERP, toPhdrFlags(Out<ELFT>::Interp->getFlags()));
 
   // Add the first PT_LOAD segment for regular output sections.
-  setPhdr(&Phdrs[++PhdrIdx], PT_LOAD, PF_R, 0, Target->getVAStart(), FileOff,
-          Target->getPageSize());
+  uintX_t Flags = PF_R;
+  Map.FirstLoadId = Map.PhdrList.size();
+  Phdr *Load = AddHdr(PT_LOAD, Flags);
 
-  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.
+  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 = AddHdr(LoadType, 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) {
+    Map.TlsPhdrId = Map.PhdrList.size();
+    Map.PhdrList.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);
-  }
+  if (isOutputDynamic())
+    AddHdr(PT_DYNAMIC, toPhdrFlags(Out<ELFT>::Dynamic->getFlags()));
 
+  // PT_GNU_RELRO includes all sections that should be marked as
+  // read-only by dynamic linker after proccessing relocations.
   if (HasRelro) {
-    Elf_Phdr *PH = &Phdrs[++PhdrIdx];
-    *PH = GnuRelroPhdr;
+    Map.RelroPhdrId = Map.PhdrList.size();
+    AddHdr(PT_GNU_RELRO, PF_R);
   }
 
-  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.
+  if (Out<ELFT>::EhFrameHdr->Live)
+    AddHdr(PT_GNU_EH_FRAME, toPhdrFlags(Out<ELFT>::EhFrameHdr->getFlags()));
 
   // 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;
+  if (!Config->ZExecStack)
+    AddHdr(PT_GNU_STACK, PF_R | PF_W);
+
+  return Map;
+}
+
+// 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.PhdrList.size());
+
+  // Main pass, here final addresses are assigned to all sections
+  // except TLS. Flags and header types are not assigned during pass.
+  // Them are assigned once for all headers at second pass.
+  bool RelroAligned = false;
+  uintX_t VA = Target->getVAStart() + sizeof(Elf_Ehdr);
+  uintX_t FileOff = sizeof(Elf_Ehdr);
+  for (size_t I = 0, E = Map.PhdrList.size(); I != E; ++I) {
+    const auto &PHdr = Map.PhdrList[I];
+    Elf_Phdr *PH = &Phdrs[I];
+    switch (PHdr->Type) {
+    case PT_LOAD:
+    case PT_AMDGPU_HSA_LOAD_CODE_AGENT:
+    case PT_AMDGPU_HSA_LOAD_GLOBAL_PROGRAM: {
+      // New PT_LOAD is usually created when access flags were changed.
+      // If it is the very first load then it is special case,
+      // it contains elf header and program headers.
+      // Otherwise we need to align VA and FileOff to new memory page
+      // and create a new load.
+      if (Map.FirstLoadId == I) {
+        setPhdr(PH, 0, 0, 0, Target->getVAStart(), FileOff,
+                Target->getPageSize());
+      } else {
+        VA = alignTo(VA, Target->getPageSize());
+        FileOff = alignTo(FileOff, Target->getPageSize());
+        setPhdr(PH, 0, 0, FileOff, VA, 0, Target->getPageSize());
+      }
+
+      for (OutputSectionBase<ELFT> *Sec : PHdr->OutSections) {
+        bool InRelRo = HasRelro && (PHdr->Flags & PF_W) && isRelroSection(Sec);
+        bool FirstNonRelRo = HasRelro && Phdrs[Map.RelroPhdrId].p_type &&
+                             needsPhdr<ELFT>(Sec) && !RelroAligned && !InRelRo;
+        // If this is a first non relro section after relro sequence then
+        // we want to align to page boundaries. We do this because we dont want
+        // any other sections except relro were marked as read only by dynamic
+        // loader. Them should be placed on a different memory page for that.
+        if (FirstNonRelRo) {
+          VA = alignTo(VA, Target->getPageSize());
+          FileOff = alignTo(FileOff, Target->getPageSize());
+          RelroAligned = true;
+        }
+
+        if (needsPhdr<ELFT>(Sec)) {
+          bool IsTls = Sec->getFlags() & SHF_TLS;
+          // If we found TLS section we want to fill phdr parameters early
+          // because later we will not have its VA and FileOff.
+          if (IsTls && !Phdrs[Map.TlsPhdrId].p_vaddr)
+            setPhdr(&Phdrs[Map.TlsPhdrId], 0, 0, FileOff, VA, 0, 0);
+
+          // 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, &Phdrs[Map.RelroPhdrId], 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;
+        }
+      }
+      break;
+    }
+    case PT_PHDR: {
+      size_t PhdrSize = sizeof(Elf_Phdr) * Phdrs.size();
+      setPhdr(PH, 0, 0, FileOff, VA, PhdrSize, /*Align=*/8);
+      FileOff += PhdrSize;
+      VA += PhdrSize;
+      break;
+    }
+    case PT_DYNAMIC:
+    case PT_GNU_EH_FRAME:
+    case PT_INTERP:
+    case PT_GNU_RELRO:
+    case PT_GNU_STACK:
+      // Do nothing here.
+      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);
+  // Final pass. Fix up PT_INTERP as we now know the address of .interp
+  // section, also assign TLS addresses, set types and flags and finish with
+  // other ones.
+  for (size_t I = 0, E = Map.PhdrList.size(); I != E; ++I) {
+    const auto &PHdr = Map.PhdrList[I];
+    Elf_Phdr *PH = &Phdrs[I];
+    switch (PHdr->Type) {
+    case PT_DYNAMIC:
+      copyPhdr(PH, Out<ELFT>::Dynamic);
+      break;
+    case PT_GNU_EH_FRAME:
+      copyPhdr(PH, Out<ELFT>::EhFrameHdr);
+      break;
+    case PT_INTERP:
+      copyPhdr(PH, Out<ELFT>::Interp);
+      break;
+    case PT_TLS:
+      assignTlsAddresses(PH, *PHdr.get());
+      break;
+    }
+    PH->p_type = PHdr->Type;
+    PH->p_flags = PHdr->Flags;
   }
 
   // Add space for section headers.
@@ -1206,35 +1296,32 @@
   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;
+// We want to assign proper addresses for all sections TLS includes.
+template <class ELFT>
+void Writer<ELFT>::assignTlsAddresses(Elf_Phdr *TlsPhdr, const Phdr &PHdr) {
+  uintX_t ThreadBssOffset = 0;
+  uintX_t VA = TlsPhdr->p_vaddr;
+  for (OutputSectionBase<ELFT> *Sec : PHdr.OutSections) {
+    if (Sec->getType() != SHT_NOBITS)
+      VA = alignTo(VA, Sec->getAlign());
+
+    uintX_t TVA = alignTo(VA + ThreadBssOffset, Sec->getAlign());
+    Sec->setVA(TVA);
+
+    if (Sec->getType() == SHT_NOBITS) {
+      ThreadBssOffset = TVA - VA + Sec->getSize();
+    } else {
+      VA += Sec->getSize();
+      TlsPhdr->p_filesz += Sec->getSize();
     }
+
+    TlsPhdr->p_memsz += Sec->getSize();
+    TlsPhdr->p_align = std::max<uintX_t>(TlsPhdr->p_align, Sec->getAlign());
   }
-  if (Tls)
-    ++I;
-  if (HasRelro)
-    ++I;
-  if (Out<ELFT>::EhFrameHdr->Live)
-    ++I;
-  return I;
+  // 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);
+  Out<ELFT>::TlsPhdr = TlsPhdr;
 }
 
 static uint32_t getELFFlags() {