3 files changed, 68 insertions, 3 deletions

diff --git a/lld/ELF/LinkerScript.cpp b/lld/ELF/LinkerScript.cpp
index 2240f7a0887..3f9b24ae8a9 100644
--- a/lld/ELF/LinkerScript.cpp
+++ b/lld/ELF/LinkerScript.cpp
@@ -112,8 +112,7 @@ static void expandMemoryRegion(MemoryRegion *MemRegion, uint64_t Size,
           "': overflowed by " + Twine(NewSize - MemRegion->Length) + " bytes");
 }
 
-void LinkerScript::expandOutputSection(uint64_t Size) {
-  Ctx->OutSec->Size += Size;
+void LinkerScript::expandMemoryRegions(uint64_t Size) {
   if (Ctx->MemRegion)
     expandMemoryRegion(Ctx->MemRegion, Size, Ctx->MemRegion->Name,
                        Ctx->OutSec->Name);
@@ -122,6 +121,11 @@ void LinkerScript::expandOutputSection(uint64_t Size) {
                        Ctx->OutSec->Name);
 }
 
+void LinkerScript::expandOutputSection(uint64_t Size) {
+  Ctx->OutSec->Size += Size;
+  expandMemoryRegions(Size);
+}
+
 void LinkerScript::setDot(Expr E, const Twine &Loc, bool InSec) {
   uint64_t Val = E().getValue();
   if (Val < Dot && InSec)
@@ -707,9 +711,11 @@ void LinkerScript::output(InputSection *S) {
 void LinkerScript::switchTo(OutputSection *Sec) {
   if (Ctx->OutSec == Sec)
     return;
-
   Ctx->OutSec = Sec;
+
+  uint64_t Before = advance(0, 1);
   Ctx->OutSec->Addr = advance(0, Ctx->OutSec->Alignment);
+  expandMemoryRegions(Ctx->OutSec->Addr - Before);
 }
 
 // This function searches for a memory region to place the given output
diff --git a/lld/ELF/LinkerScript.h b/lld/ELF/LinkerScript.h
index 4c223030a28..f808849820d 100644
--- a/lld/ELF/LinkerScript.h
+++ b/lld/ELF/LinkerScript.h
@@ -232,6 +232,7 @@ class LinkerScript final {
   void assignSymbol(SymbolAssignment *Cmd, bool InSec);
   void setDot(Expr E, const Twine &Loc, bool InSec);
   void expandOutputSection(uint64_t Size);
+  void expandMemoryRegions(uint64_t Size);
 
   std::vector<InputSection *>
   computeInputSections(const InputSectionDescription *);
diff --git a/lld/test/ELF/linkerscript/memory-region-alignment.test b/lld/test/ELF/linkerscript/memory-region-alignment.test
new file mode 100644
index 00000000000..ba9d4e3bab3
--- /dev/null
+++ b/lld/test/ELF/linkerscript/memory-region-alignment.test
@@ -0,0 +1,58 @@
+# REQUIRES: x86
+# RUN: echo '.section .foo,"a"; .quad 0; .section .zed,"M",@progbits,1; .byte 0' > %t.s
+# RUN: llvm-mc -filetype=obj -triple=x86_64-unknown-linux %t.s -o %t.o
+
+MEMORY {
+  ram (rwx): org = 0x1, len = 96K
+}
+
+SECTIONS {
+  .foo : ALIGN(8) {
+    *(.foo)
+  } > ram
+  
+  .zed : {
+    *(.zed)
+  } > ram
+}
+
+# RUN: ld.lld %t.o -o %t --script %s
+# RUN: llvm-readobj -sections %t | FileCheck %s
+
+# CHECK:        Name: .foo
+# CHECK-NEXT:   Type: SHT_PROGBITS
+# CHECK-NEXT:   Flags [
+# CHECK-NEXT:     SHF_ALLOC
+# CHECK-NEXT:   ]
+# CHECK-NEXT:   Address: 0x8
+# CHECK-NEXT:   Offset: 0x1008
+# CHECK-NEXT:   Size: 8
+
+# CHECK:        Name: .text
+# CHECK-NEXT:   Type: SHT_PROGBITS
+# CHECK-NEXT:   Flags [
+# CHECK-NEXT:     SHF_ALLOC
+# CHECK-NEXT:     SHF_EXECINSTR
+# CHECK-NEXT:   ]
+# CHECK-NEXT:   Address: 0x10
+# CHECK-NEXT:   Offset: 0x1010
+# CHECK-NEXT:   Size: 0
+
+# CHECK:        Name: .zed
+# CHECK-NEXT:   Type: SHT_PROGBITS
+# CHECK-NEXT:   Flags [
+# CHECK-NEXT:     SHF_MERGE
+# CHECK-NEXT:   ]
+# CHECK-NEXT:   Address: 0x10
+# CHECK-NEXT:   Offset: 0x1010
+# CHECK-NEXT:   Size: 1
+
+# CHECK:        Name: .comment
+# CHECK-NEXT:   Type: SHT_PROGBITS
+# CHECK-NEXT:   Flags [
+# CHECK-NEXT:     SHF_MERGE
+# CHECK-NEXT:     SHF_STRINGS
+# CHECK-NEXT:   ]
+# CHECK-NEXT:   Address: 0x0
+# CHECK-NEXT:   Offset: 0x1011
+# CHECK-NEXT:   Size: 8