[lld] Add support for EC code map.

[jacek]

Summary

The patch differs from MSVC's in two aspects that I know about:

• Empty chunks in empty sections (that are discarded anyway) can cause weird empty regions on MSVC (a test for such conditions is in the patch). I skip empty chunks in createECCodeMap to avoid that, but we could potentially remove those skips.
• Range extension thunks may slip out of map ranges on MSVC if they are added after a chunk that is otherwise the last ARM64EC chunk in a range. I tested it locally (the test needs a patch to support ARM64EC to LLD first). To imitate MSVC behaviour, we could move createECCodeMap before finalizeAddresses.

Both of above seem like a bug in MSVC, but they should be easy to imitate if desired.

[jacek]

Moving to https://github.com/llvm/llvm-project/pull/69101

lld/COFF/Chunks.h

 // of an empty chunk isn't MinGW specific.
 class EmptyChunk : public NonSectionChunk {
 public:
   EmptyChunk() {}
   size_t getSize() const override { return 0; }
   void writeTo(uint8_t *buf) const override {}
 };
 
+class ECCodeMapEntry {
+public:
+  ECCodeMapEntry(Chunk *first, Chunk *last, chpe_range_type type)
+      : first(first), last(last), type(type) {}
+  Chunk *first;
+  Chunk *last;
+  chpe_range_type type;
+};
+
+// This is a chunk containing CHPE code map on EC targets. It's a table
+// of address ranges and their types.
+class ECCodeMapChunk : public NonSectionChunk {
+public:
+  ECCodeMapChunk(std::vector<ECCodeMapEntry> &map) : map(map) {}
+  size_t getSize() const override;
+  void writeTo(uint8_t *buf) const override;
+
+private:
+  std::vector<ECCodeMapEntry> &map;
+};
+
 // MinGW specific, for the "automatic import of variables from DLLs" feature.
 // This provides the table of runtime pseudo relocations, for variable
 // references that turned out to need to be imported from a DLL even though
 // the reference didn't use the dllimport attribute. The MinGW runtime will
 // process this table after loading, before handling control over to user
 // code.
 class PseudoRelocTableChunk : public NonSectionChunk {
 public:

lld/COFF/Chunks.cpp

void RVAFlagTableChunk::writeTo(uint8_t *buf) const {
   llvm::sort(flags,
              [](const RVAFlag &a, const RVAFlag &b) { return a.rva < b.rva; });
   assert(llvm::unique(flags, [](const RVAFlag &a,
                                 const RVAFlag &b) { return a.rva == b.rva; }) ==
              flags.end() &&
          "RVA tables should be de-duplicated");
 }
 
+size_t ECCodeMapChunk::getSize() const {
+  return map.size() * sizeof(chpe_range_entry);
+}
+
+void ECCodeMapChunk::writeTo(uint8_t *buf) const {
+  auto table = reinterpret_cast<chpe_range_entry *>(buf);
+  for (uint32_t i = 0; i < map.size(); i++) {
+    const ECCodeMapEntry &entry = map[i];
+    uint32_t start = entry.first->getRVA();
+    table[i].StartOffset = start | entry.type;
+    table[i].Length = entry.last->getRVA() + entry.last->getSize() - start;
+  }
+}
+
 // MinGW specific, for the "automatic import of variables from DLLs" feature.
 size_t PseudoRelocTableChunk::getSize() const {
   if (relocs.empty())
     return 0;
   return 12 + 12 * relocs.size();
 }
 
 // MinGW specific.

lld/COFF/Driver.cpp

void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   ctx.symtab.addAbsolute(mangle("__guard_longjmp_count"), 0);
   ctx.symtab.addAbsolute(mangle("__guard_longjmp_table"), 0);
   // Needed for MSVC 2017 15.5 CRT.
   ctx.symtab.addAbsolute(mangle("__enclave_config"), 0);
   // Needed for MSVC 2019 16.8 CRT.
   ctx.symtab.addAbsolute(mangle("__guard_eh_cont_count"), 0);
   ctx.symtab.addAbsolute(mangle("__guard_eh_cont_table"), 0);
 
+  if (COFF::isArm64EC(config->machine)) {
+    ctx.symtab.addAbsolute("__hybrid_code_map", 0);
+    ctx.symtab.addAbsolute("__hybrid_code_map_count", 0);
+  }
+
   if (config->pseudoRelocs) {
     ctx.symtab.addAbsolute(mangle("__RUNTIME_PSEUDO_RELOC_LIST__"), 0);
     ctx.symtab.addAbsolute(mangle("__RUNTIME_PSEUDO_RELOC_LIST_END__"), 0);
   }
   if (config->mingw) {
     ctx.symtab.addAbsolute(mangle("__CTOR_LIST__"), 0);
     ctx.symtab.addAbsolute(mangle("__DTOR_LIST__"), 0);
   }

lld/COFF/Writer.cpp

private:
   void removeUnusedSections();
   void assignAddresses();
   bool isInRange(uint16_t relType, uint64_t s, uint64_t p, int margin);
   std::pair<Defined *, bool> getThunk(DenseMap<uint64_t, Defined *> &lastThunks,
                                       Defined *target, uint64_t p,
                                       uint16_t type, int margin);
   bool createThunks(OutputSection *os, int margin);
   bool verifyRanges(const std::vector<Chunk *> chunks);
+  void createECCodeMap();
   void finalizeAddresses();
   void removeEmptySections();
   void assignOutputSectionIndices();
   void createSymbolAndStringTable();
   void openFile(StringRef outputPath);
   template <typename PEHeaderTy> void writeHeader();
   void createSEHTable();
   void createRuntimePseudoRelocs();
+  void createECChunks();
   void insertCtorDtorSymbols();
   void markSymbolsWithRelocations(ObjFile *file, SymbolRVASet &usedSymbols);
   void createGuardCFTables();
   void markSymbolsForRVATable(ObjFile *file,
                               ArrayRef<SectionChunk *> symIdxChunks,
                               SymbolRVASet &tableSymbols);
   void getSymbolsFromSections(ObjFile *file,
                               ArrayRef<SectionChunk *> symIdxChunks,
private:
 
   void checkLoadConfig();
   template <typename T> void checkLoadConfigGuardData(const T *loadConfig);
 
   std::unique_ptr<FileOutputBuffer> &buffer;
   std::map<PartialSectionKey, PartialSection *> partialSections;
   std::vector<char> strtab;
   std::vector<llvm::object::coff_symbol16> outputSymtab;
+  std::vector<ECCodeMapEntry> codeMap;
   IdataContents idata;
   Chunk *importTableStart = nullptr;
   uint64_t importTableSize = 0;
   Chunk *edataStart = nullptr;
   Chunk *edataEnd = nullptr;
   Chunk *iatStart = nullptr;
   uint64_t iatSize = 0;
   DelayLoadContents delayIdata;
for (size_t i = 0, e = originalRelocs.size(); i != e; ++i) {
       }
     }
 
     sc->setRelocs(newRelocs);
   }
   return addressesChanged;
 }
 
+// Create a code map for CHPE metadata.
+void Writer::createECCodeMap() {
+  if (!isArm64EC(ctx.config.machine))
+    return;
+
+  // Clear the map in case we were're recomputing the map after adding
+  // a range extension thunk.
+  codeMap.clear();
+
+  Chunk *first = nullptr, *last;
+  chpe_range_type lastType;
+
+  auto closeRange = [&]() {
+    if (first) {
+      codeMap.push_back({first, last, lastType});
+      first = nullptr;
+    }
+  };
+
+  for (OutputSection *sec : ctx.outputSections) {
+    if (!sec->isCodeSection()) {
+      closeRange();
+      continue;
+    }
+
+    for (Chunk *c : sec->chunks) {
+      // Skip empty section chunks. MSVC does not seem to do that and
+      // generates empty code ranges in some cases.
+      if (isa<SectionChunk>(c) && !c->getSize())
+        continue;
+
+      chpe_range_type chunkType = c->getECRangeType();
+      if (!first) {
+        first = c;
+      } else if (lastType != chunkType) {
+        closeRange();
+        first = c;
+      }
+      last = c;
+      lastType = chunkType;
+    }
+  }
+
+  closeRange();
+
+  Symbol *tableCountSym = ctx.symtab.findUnderscore("__hybrid_code_map_count");
+  cast<DefinedAbsolute>(tableCountSym)->setVA(codeMap.size());
+}
+
 // Verify that all relocations are in range, with no extra margin requirements.
 bool Writer::verifyRanges(const std::vector<Chunk *> chunks) {
   for (Chunk *c : chunks) {
     SectionChunk *sc = dyn_cast_or_null<SectionChunk>(c);
     if (!sc)
       continue;
 
     ArrayRef<coff_relocation> relocs = sc->getRelocs();
void Writer::createMiscChunks() {
   // Create SEH table. x86-only.
   if (config->safeSEH)
     createSEHTable();
 
   // Create /guard:cf tables if requested.
   if (config->guardCF != GuardCFLevel::Off)
     createGuardCFTables();
 
+  if (isArm64EC(config->machine))
+    createECChunks();
+
   if (config->autoImport)
     createRuntimePseudoRelocs();
 
   if (config->mingw)
     insertCtorDtorSymbols();
 }
 
 // Create .idata section for the DLL-imported symbol table.
void Writer::sortECChunks() {
   }
 }
 
 // Visits all sections to assign incremental, non-overlapping RVAs and
 // file offsets.
 void Writer::assignAddresses() {
   Configuration *config = &ctx.config;
 
+  // We need to create EC code map so that ECCodeMapChunk knows its size.
+  // We do it here to make sure that we account for range extension chunks.
+  createECCodeMap();
+
   sizeOfHeaders = dosStubSize + sizeof(PEMagic) + sizeof(coff_file_header) +
                   sizeof(data_directory) * numberOfDataDirectory +
                   sizeof(coff_section) * ctx.outputSections.size();
   sizeOfHeaders +=
       config->is64() ? sizeof(pe32plus_header) : sizeof(pe32_header);
   sizeOfHeaders = alignTo(sizeOfHeaders, config->fileAlign);
   fileSize = sizeOfHeaders;
 
void Writer::maybeAddRVATable(SymbolRVASet tableSymbols, StringRef tableSym,
   rdataSec->addChunk(tableChunk);
 
   Symbol *t = ctx.symtab.findUnderscore(tableSym);
   Symbol *c = ctx.symtab.findUnderscore(countSym);
   replaceSymbol<DefinedSynthetic>(t, t->getName(), tableChunk);
   cast<DefinedAbsolute>(c)->setVA(tableChunk->getSize() / (hasFlag ? 5 : 4));
 }
 
+// Create CHPE metadata chunks.
+void Writer::createECChunks() {
+  auto codeMapChunk = make<ECCodeMapChunk>(codeMap);
+  rdataSec->addChunk(codeMapChunk);
+  Symbol *codeMapSym = ctx.symtab.findUnderscore("__hybrid_code_map");
+  replaceSymbol<DefinedSynthetic>(codeMapSym, codeMapSym->getName(),
+                                  codeMapChunk);
+}
+
 // MinGW specific. Gather all relocations that are imported from a DLL even
 // though the code didn't expect it to, produce the table that the runtime
 // uses for fixing them up, and provide the synthetic symbols that the
 // runtime uses for finding the table.
 void Writer::createRuntimePseudoRelocs() {
   std::vector<RuntimePseudoReloc> rels;
 
   for (Chunk *c : ctx.symtab.getChunks()) {

lld/test/COFF/Inputs/loadconfig-arm64ec.s

 __os_arm64x_dispatch_fptr:
         .xword 0
 
         .data
         .globl __chpe_metadata
         .p2align 2, 0
 __chpe_metadata:
         .word 1
-        .rva code_map
-        .word code_map_count
+        .rva __hybrid_code_map
+        .word __hybrid_code_map_count
         .word 0 // __x64_code_ranges_to_entry_points
         .word 0 //__arm64x_redirection_metadata
         .rva __os_arm64x_dispatch_call_no_redirect
         .rva __os_arm64x_dispatch_ret
         .rva __os_arm64x_check_call
         .rva __os_arm64x_check_icall
         .rva __os_arm64x_check_icall_cfg
         .word 0 // __arm64x_native_entrypoint

lld/test/COFF/arm64ec-codemap.test

 REQUIRES: aarch64, x86
 RUN: split-file %s %t.dir && cd %t.dir
 
 RUN: llvm-mc -filetype=obj -triple=arm64-windows arm64-func-sym.s -o arm64-func-sym.obj
 RUN: llvm-mc -filetype=obj -triple=arm64ec-windows arm64ec-func-sym.s -o arm64ec-func-sym.obj
+RUN: llvm-mc -filetype=obj -triple=arm64ec-windows empty-sec.s -o arm64ec-empty-sec.obj
 RUN: llvm-mc -filetype=obj -triple=x86_64-windows x86_64-func-sym.s -o x86_64-func-sym.obj
-RUN: llvm-mc -filetype=obj -triple=arm64ec-windows codemap.s -o codemap.obj
-RUN: llvm-mc -filetype=obj -triple=arm64ec-windows codemap2.s -o codemap2.obj
-RUN: llvm-mc -filetype=obj -triple=arm64ec-windows codemap3.s -o codemap3.obj
+RUN: llvm-mc -filetype=obj -triple=x86_64-windows empty-sec.s -o x86_64-empty-sec.obj
 RUN: llvm-mc -filetype=obj -triple=arm64ec-windows %S/Inputs/loadconfig-arm64ec.s -o loadconfig-arm64ec.obj
 
 Link ARM64EC DLL and verify that the code is arranged as expected.
 
 RUN: lld-link -out:test.dll -machine:arm64ec arm64ec-func-sym.obj x86_64-func-sym.obj \
-RUN:          codemap.obj loadconfig-arm64ec.obj -dll -noentry
+RUN:          loadconfig-arm64ec.obj -dll -noentry
 
 RUN: llvm-readobj --coff-load-config test.dll | FileCheck -check-prefix=CODEMAP %s
 CODEMAP:       CodeMap [
 CODEMAP-NEXT:    0x1000 - 0x1008  ARM64EC
 CODEMAP-NEXT:    0x2000 - 0x2006  X64
 CODEMAP-NEXT:    0x5000 - 0x5008  ARM64EC
 CODEMAP-NEXT:    0x6000 - 0x6006  X64
 CODEMAP-NEXT:  ]

 DISASM-NEXT: 180005004: d65f03c0     ret
 DISASM-NEXT:                 ...
 DISASM-NEXT: 180006000: b8 06 00 00 00               movl    $0x6, %eax
 DISASM-NEXT: 180006005: c3                           retq
 
 Order of arguments doesn't matter in this case, chunks are sorted by target type anyway.
 
 RUN: lld-link -out:test2.dll -machine:arm64ec x86_64-func-sym.obj arm64ec-func-sym.obj \
-RUN:          codemap.obj loadconfig-arm64ec.obj -dll -noentry
+RUN:          loadconfig-arm64ec.obj -dll -noentry
 RUN: llvm-readobj --coff-load-config test2.dll | FileCheck -check-prefix=CODEMAP %s
 RUN: llvm-objdump -d test2.dll | FileCheck -check-prefix=DISASM %s
 
 RUN: lld-link -out:testx.dll -machine:arm64x arm64-func-sym.obj arm64ec-func-sym.obj \
-RUN:          x86_64-func-sym.obj codemap2.obj loadconfig-arm64ec.obj -dll -noentry
+RUN:          x86_64-func-sym.obj loadconfig-arm64ec.obj -dll -noentry
+
+Adding empty chunks does not affect code map ranges.
+
+RUN: lld-link -out:test3.dll -machine:arm64ec x86_64-empty-sec.obj arm64ec-empty-sec.obj \
+RUN:          arm64ec-func-sym.obj  x86_64-func-sym.obj loadconfig-arm64ec.obj -dll -noentry
+RUN: llvm-readobj --coff-load-config test3.dll | FileCheck -check-prefix=CODEMAP %s
 
 Do the same with ARM64X target.
 
 RUN: llvm-readobj --coff-load-config testx.dll | FileCheck -check-prefix=CODEMAPX %s
 CODEMAPX:       CodeMap [
 CODEMAPX-NEXT:    0x1000 - 0x1008  ARM64
 CODEMAPX-NEXT:    0x2000 - 0x2008  ARM64EC
 CODEMAPX-NEXT:    0x3000 - 0x3006  X64

 DISASMX-NEXT: 180006004: d65f03c0     ret
 DISASMX-NEXT:                 ...
 DISASMX-NEXT: 180007000: b8 06 00 00 00               movl    $0x6, %eax
 DISASMX-NEXT: 180007005: c3                           retq
 
 Test merged sections.
 
 RUN: lld-link -out:testm.dll -machine:arm64ec arm64ec-func-sym.obj x86_64-func-sym.obj \
-RUN:          codemap3.obj loadconfig-arm64ec.obj -dll -noentry -merge:test=.text
+RUN:          loadconfig-arm64ec.obj -dll -noentry -merge:test=.text
 
 RUN: llvm-readobj --coff-load-config testm.dll | FileCheck -check-prefix=CODEMAPM %s
 CODEMAPM:      CodeMap [
 CODEMAPM-NEXT:   0x1000 - 0x1010  ARM64EC
 CODEMAPM-NEXT:   0x2000 - 0x200E  X64
 CODEMAPM-NEXT: ]
 
 RUN: llvm-objdump -d testm.dll | FileCheck -check-prefix=DISASMM %s
x86_64_func_sym:
 
     .section test, "xr"
     .globl x86_64_func_sym2
     .p2align 2, 0x0
 x86_64_func_sym2:
     movl $6, %eax
     retq
 
-#--- codemap.s
-    .section .rdata,"dr"
-    .globl code_map
-code_map:
-    .rva arm64ec_func_sym + 1
-    .word 8
-    .rva x86_64_func_sym + 2
-    .word 6
-    .rva arm64ec_func_sym2 + 1
-    .word 8
-    .rva x86_64_func_sym2 + 2
-    .word 6
-
-    .globl code_map_count
-code_map_count = 4
-
-#--- codemap2.s
-    .section .rdata,"dr"
-    .globl code_map
-code_map:
-    .rva arm64_func_sym
-    .word 8
-    .rva arm64ec_func_sym + 1
-    .word 8
-    .rva x86_64_func_sym + 2
-    .word 6
-    .rva arm64ec_func_sym2 + 1
-    .word 8
-    .rva x86_64_func_sym2 + 2
-    .word 6
-
-    .globl code_map_count
-code_map_count = 5
-
-#--- codemap3.s
-    .section .rdata,"dr"
-    .globl code_map
-code_map:
-    .rva arm64ec_func_sym + 1
-    .word 16
-    .rva x86_64_func_sym + 2
-    .word 14
-
-    .globl code_map_count
-code_map_count = 2
+#--- empty-sec.s
+    .section .empty1, "xr"
+    .section .empty2, "xr"
+    .section .empty3, "xr"