[lld] Sort data chunks before code chunks on ARM64EC. (#70722)

3 files changed, 47 insertions, 3 deletions

diff --git a/lld/COFF/Chunks.h b/lld/COFF/Chunks.h
index b82bc94..156e7a8 100644
--- a/lld/COFF/Chunks.h
+++ b/lld/COFF/Chunks.h
@@ -116,7 +116,7 @@ public:
   bool isHotPatchable() const;
 
   MachineTypes getMachine() const;
-  chpe_range_type getArm64ECRangeType() const;
+  std::optional<chpe_range_type> getArm64ECRangeType() const;
 
 protected:
   Chunk(Kind k = OtherKind) : chunkKind(k), hasData(true), p2Align(0) {}
@@ -437,7 +437,11 @@ inline MachineTypes Chunk::getMachine() const {
   return static_cast<const NonSectionChunk *>(this)->getMachine();
 }
 
-inline chpe_range_type Chunk::getArm64ECRangeType() const {
+inline std::optional<chpe_range_type> Chunk::getArm64ECRangeType() const {
+  // Data sections don't need codemap entries.
+  if (!(getOutputCharacteristics() & llvm::COFF::IMAGE_SCN_MEM_EXECUTE))
+    return std::nullopt;
+
   switch (getMachine()) {
   case AMD64:
     return chpe_range_type::Amd64;
diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
index 43d8e7c..895a6e0 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -1389,7 +1389,9 @@ void Writer::sortECChunks() {
   for (OutputSection *sec : ctx.outputSections) {
     if (sec->isCodeSection())
       llvm::stable_sort(sec->chunks, [=](const Chunk *a, const Chunk *b) {
-        return a->getArm64ECRangeType() < b->getArm64ECRangeType();
+        std::optional<chpe_range_type> aType = a->getArm64ECRangeType(),
+                                       bType = b->getArm64ECRangeType();
+        return !aType || (bType && *aType < *bType);
       });
   }
 }
diff --git a/lld/test/COFF/arm64ec-codemap.test b/lld/test/COFF/arm64ec-codemap.test
index 424456a..fcc7db2 100644
--- a/lld/test/COFF/arm64ec-codemap.test
+++ b/lld/test/COFF/arm64ec-codemap.test
@@ -110,6 +110,40 @@ DISASMM-NEXT:                 ...
 DISASMM-NEXT: 180002ffe: 00 00                        addb    %al, (%rax)
 DISASMM-NEXT: 180003000: b8 06 00 00 00               movl    $0x6, %eax
 
+Merging data sections into code sections causes data to be separated from the code when sorting chunks.
+
+RUN: lld-link -out:testdm.dll -machine:arm64ec arm64ec-func-sym.obj x86_64-func-sym.obj codemap.obj \
+RUN:          data-sec.obj loadconfig-arm64ec.obj -dll -noentry -merge:.testdata=.text -merge:.rdata=test
+
+RUN: llvm-readobj --coff-load-config testdm.dll | FileCheck -check-prefix=CODEMAPDM %s
+CODEMAPDM:      CodeMap [
+CODEMAPDM-NEXT:     0x2000 - 0x2008  ARM64EC
+CODEMAPDM-NEXT:     0x3000 - 0x3006  X64
+CODEMAPDM-NEXT:     0x5200 - 0x5208  ARM64EC
+CODEMAPDM-NEXT:     0x6000 - 0x6006  X64
+CODEMAPDM-NEXT: ]
+
+RUN: llvm-objdump -d testdm.dll | FileCheck -check-prefix=DISASMDM %s
+DISASMDM:      Disassembly of section .text:
+DISASMDM-EMPTY:
+DISASMDM-NEXT: 0000000180001000 <.text>:
+DISASMDM-NEXT: 180001000: 00000001     udf     #0x1
+DISASMDM-NEXT:                 ...
+DISASMDM-NEXT: 180002000: 52800040     mov     w0, #0x2
+DISASMDM-NEXT: 180002004: d65f03c0     ret
+DISASMDM-NEXT:                 ...
+DISASMDM-NEXT: 180003000: b8 03 00 00 00               movl    $0x3, %eax
+DISASMDM-NEXT: 180003005: c3                           retq
+DISASMDM-EMPTY:
+DISASMDM-NEXT: Disassembly of section test:
+DISASMDM-EMPTY:
+DISASMDM-NEXT: 0000000180005000 <test>:
+DISASMDM:      180005200: 528000a0     mov     w0, #0x5
+DISASMDM-NEXT: 180005204: d65f03c0     ret
+DISASMDM-NEXT:                 ...
+DISASMDM-NEXT: 180006000: b8 06 00 00 00               movl    $0x6, %eax
+DISASMDM-NEXT: 180006005: c3                           retq
+
 #--- arm64-func-sym.s
     .text
     .globl arm64_func_sym
@@ -148,6 +182,10 @@ x86_64_func_sym2:
     movl $6, %eax
     retq
 
+#--- data-sec.s
+    .section .testdata, "rd"
+    .xword 1
+
 #--- codemap.s
     .section .rdata,"dr"
     .globl code_map