[LLD][COFF] Add support for ARM64EC entry thunks.

For x86_64 callable functions, ARM64EC requires an entry thunk generated by
the compiler. The linker interprets .hybmp sections to associate function
chunks with their entry points and writes an offset to thunks preceding
function section contents.

Additionally, ICF needs to be aware of entry thunks to not consider chunks
to be equal when they have different entry thunks, and GC needs to mark
entry thunks together with function chunks.

I used a new SectionChunkEC class instead of storing entry thunks in
SectionChunk, following the guideline to keep SectionChunk as compact as
possible. This way, there is no memory usage increase on non-EC targets.

Author: cjacek

lld/COFF/Chunks.cpp

@@ -31,8 +31,8 @@ using llvm::support::ulittle32_t;
 
 namespace lld::coff {
 
-SectionChunk::SectionChunk(ObjFile *f, const coff_section *h)
-    : Chunk(SectionKind), file(f), header(h), repl(this) {
+SectionChunk::SectionChunk(ObjFile *f, const coff_section *h, Kind k)
+    : Chunk(k), file(f), header(h), repl(this) {
   // Initialize relocs.
   if (file)
     setRelocs(file->getCOFFObj()->getRelocations(header));
@@ -410,6 +410,10 @@ void SectionChunk::writeTo(uint8_t *buf) const {
 
     applyRelocation(buf + rel.VirtualAddress, rel);
   }
+
+  // Write the offset to EC entry thunk preceding section contents.
+  if (Defined *entryThunk = getEntryThunk())
+    write32le(buf - sizeof(uint32_t), entryThunk->getRVA() - rva + 1);
 }
 
 void SectionChunk::applyRelocation(uint8_t *off,

lld/COFF/Chunks.h

@@ -55,7 +55,12 @@ enum : unsigned { Log2MaxSectionAlignment = 13 };
 // doesn't even have actual data (if common or bss).
 class Chunk {
 public:
-  enum Kind : uint8_t { SectionKind, OtherKind, ImportThunkKind };
+  enum Kind : uint8_t {
+    SectionKind,
+    SectionECKind,
+    OtherKind,
+    ImportThunkKind
+  };
   Kind kind() const { return chunkKind; }
 
   // Returns the size of this chunk (even if this is a common or BSS.)
@@ -120,6 +125,10 @@ class Chunk {
   llvm::Triple::ArchType getArch() const;
   std::optional<chpe_range_type> getArm64ECRangeType() const;
 
+  // ARM64EC entry thunk associated with the chunk.
+  Defined *getEntryThunk() const;
+  void setEntryThunk(Defined *entryThunk);
+
 protected:
   Chunk(Kind k = OtherKind) : chunkKind(k), hasData(true), p2Align(0) {}
 
@@ -176,7 +185,7 @@ class NonSectionChunk : public Chunk {
   // bytes, so this is used only for logging or debugging.
   virtual StringRef getDebugName() const { return ""; }
 
-  static bool classof(const Chunk *c) { return c->kind() != SectionKind; }
+  static bool classof(const Chunk *c) { return c->kind() >= OtherKind; }
 
 protected:
   NonSectionChunk(Kind k = OtherKind) : Chunk(k) {}
@@ -210,7 +219,7 @@ class RuntimePseudoReloc {
 };
 
 // A chunk corresponding a section of an input file.
-class SectionChunk final : public Chunk {
+class SectionChunk : public Chunk {
   // Identical COMDAT Folding feature accesses section internal data.
   friend class ICF;
 
@@ -231,8 +240,8 @@ class SectionChunk final : public Chunk {
     Symbol *operator*() const { return file->getSymbol(I->SymbolTableIndex); }
   };
 
-  SectionChunk(ObjFile *file, const coff_section *header);
-  static bool classof(const Chunk *c) { return c->kind() == SectionKind; }
+  SectionChunk(ObjFile *file, const coff_section *header, Kind k = SectionKind);
+  static bool classof(const Chunk *c) { return c->kind() <= SectionECKind; }
   size_t getSize() const { return header->SizeOfRawData; }
   ArrayRef<uint8_t> getContents() const;
   void writeTo(uint8_t *buf) const;
@@ -393,6 +402,16 @@ class SectionChunk final : public Chunk {
   uint32_t sectionNameSize = 0;
 };
 
+// A section chunk corresponding a section of an EC input file.
+class SectionChunkEC final : public SectionChunk {
+public:
+  static bool classof(const Chunk *c) { return c->kind() == SectionECKind; }
+
+  SectionChunkEC(ObjFile *file, const coff_section *header)
+      : SectionChunk(file, header, SectionECKind) {}
+  Defined *entryThunk = nullptr;
+};
+
 // Inline methods to implement faux-virtual dispatch for SectionChunk.
 
 inline size_t Chunk::getSize() const {
@@ -775,6 +794,17 @@ inline bool Chunk::isHotPatchable() const {
   return false;
 }
 
+inline Defined *Chunk::getEntryThunk() const {
+  if (auto *c = dyn_cast<const SectionChunkEC>(this))
+    return c->entryThunk;
+  return nullptr;
+}
+
+inline void Chunk::setEntryThunk(Defined *entryThunk) {
+  if (auto c = dyn_cast<SectionChunkEC>(this))
+    c->entryThunk = entryThunk;
+}
+
 void applyMOV32T(uint8_t *off, uint32_t v);
 void applyBranch24T(uint8_t *off, int32_t v);
 

lld/COFF/Driver.cpp

@@ -2604,6 +2604,8 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   if (auto *arg = args.getLastArg(OPT_print_symbol_order))
     config->printSymbolOrder = arg->getValue();
 
+  ctx.symtab.initializeEntryThunks();
+
   // Identify unreferenced COMDAT sections.
   if (config->doGC) {
     if (config->mingw) {

lld/COFF/ICF.cpp

@@ -184,16 +184,25 @@ bool ICF::equalsConstant(const SectionChunk *a, const SectionChunk *b) {
 // Compare "moving" part of two sections, namely relocation targets.
 bool ICF::equalsVariable(const SectionChunk *a, const SectionChunk *b) {
   // Compare relocations.
-  auto eq = [&](const coff_relocation &r1, const coff_relocation &r2) {
-    Symbol *b1 = a->file->getSymbol(r1.SymbolTableIndex);
-    Symbol *b2 = b->file->getSymbol(r2.SymbolTableIndex);
+  auto eqSym = [&](Symbol *b1, Symbol *b2) {
     if (b1 == b2)
       return true;
     if (auto *d1 = dyn_cast<DefinedRegular>(b1))
       if (auto *d2 = dyn_cast<DefinedRegular>(b2))
         return d1->getChunk()->eqClass[cnt % 2] == d2->getChunk()->eqClass[cnt % 2];
     return false;
   };
+  auto eq = [&](const coff_relocation &r1, const coff_relocation &r2) {
+    Symbol *b1 = a->file->getSymbol(r1.SymbolTableIndex);
+    Symbol *b2 = b->file->getSymbol(r2.SymbolTableIndex);
+    return eqSym(b1, b2);
+  };
+
+  Symbol *e1 = a->getEntryThunk();
+  Symbol *e2 = b->getEntryThunk();
+  if ((e1 || e2) && (!e1 || !e2 || !eqSym(e1, e2)))
+    return false;
+
   return std::equal(a->getRelocs().begin(), a->getRelocs().end(),
                     b->getRelocs().begin(), eq) &&
          assocEquals(a, b);

lld/COFF/InputFiles.cpp

@@ -152,6 +152,38 @@ void ObjFile::parseLazy() {
   }
 }
 
+struct ECMapEntry {
+  ulittle32_t src;
+  ulittle32_t dst;
+  ulittle32_t type;
+};
+
+void ObjFile::initializeECThunks() {
+  for (SectionChunk *chunk : hybmpChunks) {
+    if (chunk->getContents().size() % sizeof(ECMapEntry)) {
+      error("Invalid .hybmp chunk size " + Twine(chunk->getContents().size()));
+      return;
+    }
+
+    const uint8_t *end =
+        chunk->getContents().data() + chunk->getContents().size();
+    for (const uint8_t *iter = chunk->getContents().data(); iter != end;
+         iter += sizeof(ECMapEntry)) {
+      auto entry = reinterpret_cast<const ECMapEntry *>(iter);
+      switch (entry->type) {
+      case Arm64ECThunkType::Entry:
+        ctx.symtab.addEntryThunk(getSymbol(entry->src), getSymbol(entry->dst));
+        break;
+      case Arm64ECThunkType::Exit:
+      case Arm64ECThunkType::GuestExit:
+        break;
+      default:
+        warn("Ignoring unknown EC thunk type " + Twine(entry->type));
+      }
+    }
+  }
+}
+
 void ObjFile::parse() {
   // Parse a memory buffer as a COFF file.
   std::unique_ptr<Binary> bin = CHECK(createBinary(mb), this);
@@ -168,6 +200,7 @@ void ObjFile::parse() {
   initializeSymbols();
   initializeFlags();
   initializeDependencies();
+  initializeECThunks();
 }
 
 const coff_section *ObjFile::getSection(uint32_t i) {
@@ -242,7 +275,11 @@ SectionChunk *ObjFile::readSection(uint32_t sectionNumber,
 
   if (sec->Characteristics & llvm::COFF::IMAGE_SCN_LNK_REMOVE)
     return nullptr;
-  auto *c = make<SectionChunk>(this, sec);
+  SectionChunk *c;
+  if (isArm64EC(getMachineType()))
+    c = make<SectionChunkEC>(this, sec);
+  else
+    c = make<SectionChunk>(this, sec);
   if (def)
     c->checksum = def->CheckSum;
 
@@ -260,6 +297,8 @@ SectionChunk *ObjFile::readSection(uint32_t sectionNumber,
     guardEHContChunks.push_back(c);
   else if (name == ".sxdata")
     sxDataChunks.push_back(c);
+  else if (isArm64EC(getMachineType()) && name == ".hybmp$x")
+    hybmpChunks.push_back(c);
   else if (ctx.config.tailMerge && sec->NumberOfRelocations == 0 &&
            name == ".rdata" && leaderName.starts_with("??_C@"))
     // COFF sections that look like string literal sections (i.e. no

lld/COFF/InputFiles.h

@@ -227,6 +227,7 @@ class ObjFile : public InputFile {
   void initializeSymbols();
   void initializeFlags();
   void initializeDependencies();
+  void initializeECThunks();
 
   SectionChunk *
   readSection(uint32_t sectionNumber,
@@ -292,6 +293,8 @@ class ObjFile : public InputFile {
   std::vector<SectionChunk *> guardLJmpChunks;
   std::vector<SectionChunk *> guardEHContChunks;
 
+  std::vector<SectionChunk *> hybmpChunks;
+
   // This vector contains a list of all symbols defined or referenced by this
   // file. They are indexed such that you can get a Symbol by symbol
   // index. Nonexistent indices (which are occupied by auxiliary

lld/COFF/MarkLive.cpp

@@ -68,6 +68,10 @@ void markLive(COFFLinkerContext &ctx) {
     // Mark associative sections if any.
     for (SectionChunk &c : sc->children())
       enqueue(&c);
+
+    // Mark EC entry thunks.
+    if (Defined *entryThunk = sc->getEntryThunk())
+      addSym(entryThunk);
   }
 }
 }

lld/COFF/SymbolTable.cpp

@@ -561,6 +561,25 @@ std::pair<Symbol *, bool> SymbolTable::insert(StringRef name, InputFile *file) {
   return result;
 }
 
+void SymbolTable::addEntryThunk(Symbol *from, Symbol *to) {
+  entryThunks.push_back({from, to});
+}
+
+void SymbolTable::initializeEntryThunks() {
+  for (auto it : entryThunks) {
+    auto *to = dyn_cast<Defined>(it.second);
+    if (!to)
+      continue;
+    auto *from = dyn_cast<DefinedRegular>(it.first);
+    if (!from || !from->getChunk()->isCOMDAT() ||
+        cast<DefinedRegular>(from)->getValue()) {
+      error("non COMDAT symbol '" + from->getName() + "' in hybrid map");
+      continue;
+    }
+    from->getChunk()->setEntryThunk(to);
+  }
+}
+
 Symbol *SymbolTable::addUndefined(StringRef name, InputFile *f,
                                   bool isWeakAlias) {
   auto [s, wasInserted] = insert(name, f);

lld/COFF/SymbolTable.h

@@ -106,6 +106,8 @@ class SymbolTable {
   Symbol *addImportThunk(StringRef name, DefinedImportData *s,
                          uint16_t machine);
   void addLibcall(StringRef name);
+  void addEntryThunk(Symbol *from, Symbol *to);
+  void initializeEntryThunks();
 
   void reportDuplicate(Symbol *existing, InputFile *newFile,
                        SectionChunk *newSc = nullptr,
@@ -134,6 +136,7 @@ class SymbolTable {
   llvm::DenseMap<llvm::CachedHashStringRef, Symbol *> symMap;
   std::unique_ptr<BitcodeCompiler> lto;
   bool ltoCompilationDone = false;
+  std::vector<std::pair<Symbol *, Symbol *>> entryThunks;
 
   COFFLinkerContext &ctx;
 };

lld/COFF/Writer.cpp

@@ -1532,6 +1532,10 @@ void Writer::assignAddresses() {
       }
       if (padding && c->isHotPatchable())
         virtualSize += padding;
+      // If chunk has EC entry thunk, reserve a space for an offset to the
+      // thunk.
+      if (c->getEntryThunk())
+        virtualSize += sizeof(uint32_t);
       virtualSize = alignTo(virtualSize, c->getAlignment());
       c->setRVA(rva + virtualSize);
       virtualSize += c->getSize();