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

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,12 @@ void SectionChunk::writeTo(uint8_t *buf) const {
 
     applyRelocation(buf + rel.VirtualAddress, rel);
   }
+
+  // Write the offset to EC entry thunk preceding section contents. The low bit
+  // is always set, so it's effectively an offset from the last byte of the
+  // offset.
+  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()));
+      continue;
+    }
+
+    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,28 @@ 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);
+    // We need to be able to add padding to the function and fill it with an
+    // offset to its entry thunks. To ensure that padding the function is
+    // feasible, functions are required to be COMDAT symbols with no offset.
+    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();