[LLD][COFF] Add support for ARM64X same-address thunks

lld/COFF/Chunks.cpp

@@ -875,6 +875,19 @@ void RangeExtensionThunkARM64::writeTo(uint8_t *buf) const {
   applyArm64Imm(buf + 4, target->getRVA() & 0xfff, 0);
 }
 
+void SameAddressThunkARM64EC::setDynamicRelocs(COFFLinkerContext &ctx) const {
+  // Add ARM64X relocations replacing adrp/add instructions with a version using
+  // the hybrid target.
+  RangeExtensionThunkARM64 hybridView(ARM64EC, hybridTarget);
+  uint8_t buf[sizeof(arm64Thunk)];
+  hybridView.setRVA(rva);
+  hybridView.writeTo(buf);
+  uint32_t addrp = *reinterpret_cast<ulittle32_t *>(buf);
+  uint32_t add = *reinterpret_cast<ulittle32_t *>(buf + sizeof(uint32_t));
+  ctx.dynamicRelocs->set(this, addrp);
+  ctx.dynamicRelocs->set(Arm64XRelocVal(this, sizeof(uint32_t)), add);
+}
+
 LocalImportChunk::LocalImportChunk(COFFLinkerContext &c, Defined *s)
     : sym(s), ctx(c) {
   setAlignment(ctx.config.wordsize);
@@ -1258,7 +1271,8 @@ void DynamicRelocsChunk::finalize() {
 }
 
 // Set the reloc value. The reloc entry must be allocated beforehand.
-void DynamicRelocsChunk::set(uint32_t rva, Arm64XRelocVal value) {
+void DynamicRelocsChunk::set(Arm64XRelocVal offset, Arm64XRelocVal value) {
+  uint32_t rva = offset.get();
   auto entry =
       llvm::find_if(arm64xRelocs, [rva](const Arm64XDynamicRelocEntry &e) {
         return e.offset.get() == rva;

lld/COFF/Chunks.h

@@ -193,6 +193,8 @@ class NonSectionChunk : public Chunk {
   // allowed ranges. Return the additional space required for the extension.
   virtual uint32_t extendRanges() { return 0; };
 
+  virtual Defined *getEntryThunk() const { return nullptr; };
+
   static bool classof(const Chunk *c) { return c->kind() >= OtherKind; }
 
 protected:
@@ -633,7 +635,7 @@ class ImportThunkChunkARM64EC : public ImportThunkChunk {
   bool verifyRanges() override;
   uint32_t extendRanges() override;
 
-  Defined *exitThunk;
+  Defined *exitThunk = nullptr;
   Defined *sym = nullptr;
   bool extended = false;
 
@@ -675,6 +677,26 @@ class RangeExtensionThunkARM64 : public NonSectionCodeChunk {
   MachineTypes machine;
 };
 
+// A chunk used to guarantee the same address for a function in both views of
+// a hybrid image. Similar to RangeExtensionThunkARM64 chunks, it calls the
+// target symbol using a BR instruction. It also contains an entry thunk for EC
+// compatibility and additional ARM64X relocations that swap targets between
+// views.
+class SameAddressThunkARM64EC : public RangeExtensionThunkARM64 {
+public:
+  explicit SameAddressThunkARM64EC(Defined *t, Defined *hybridTarget,
+                                   Defined *entryThunk)
+      : RangeExtensionThunkARM64(ARM64EC, t), hybridTarget(hybridTarget),
+        entryThunk(entryThunk) {}
+
+  Defined *getEntryThunk() const override { return entryThunk; }
+  void setDynamicRelocs(COFFLinkerContext &ctx) const;
+
+private:
+  Defined *hybridTarget;
+  Defined *entryThunk;
+};
+
 // Windows-specific.
 // See comments for DefinedLocalImport class.
 class LocalImportChunk : public NonSectionChunk {
@@ -843,13 +865,13 @@ class Arm64XRelocVal {
 public:
   Arm64XRelocVal(uint64_t value = 0) : value(value) {}
   Arm64XRelocVal(Defined *sym, int32_t offset = 0) : sym(sym), value(offset) {}
-  Arm64XRelocVal(Chunk *chunk, int32_t offset = 0)
+  Arm64XRelocVal(const Chunk *chunk, int32_t offset = 0)
       : chunk(chunk), value(offset) {}
   uint64_t get() const;
 
 private:
   Defined *sym = nullptr;
-  Chunk *chunk = nullptr;
+  const Chunk *chunk = nullptr;
   uint64_t value;
 };
 
@@ -884,7 +906,7 @@ class DynamicRelocsChunk : public NonSectionChunk {
     arm64xRelocs.emplace_back(type, size, offset, value);
   }
 
-  void set(uint32_t rva, Arm64XRelocVal value);
+  void set(Arm64XRelocVal offset, Arm64XRelocVal value);
 
 private:
   std::vector<Arm64XDynamicRelocEntry> arm64xRelocs;
@@ -940,6 +962,8 @@ inline bool Chunk::isHotPatchable() const {
 inline Defined *Chunk::getEntryThunk() const {
   if (auto *c = dyn_cast<const SectionChunkEC>(this))
     return c->entryThunk;
+  if (auto *c = dyn_cast<const NonSectionChunk>(this))
+    return c->getEntryThunk();
   return nullptr;
 }
 

lld/COFF/Config.h

@@ -223,6 +223,9 @@ struct Configuration {
   StringRef manifestUIAccess = "'false'";
   StringRef manifestFile;
 
+  // used for /arm64xsameaddress
+  std::vector<std::pair<Symbol *, Symbol *>> sameAddresses;
+
   // used for /dwodir
   StringRef dwoDir;
 

lld/COFF/Driver.cpp

@@ -500,7 +500,9 @@ void LinkerDriver::parseDirectives(InputFile *file) {
       file->symtab.parseAlternateName(arg->getValue());
       break;
     case OPT_arm64xsameaddress:
-      if (!file->symtab.isEC())
+      if (file->symtab.isEC())
+        parseSameAddress(arg->getValue());
+      else
         Warn(ctx) << arg->getSpelling()
                   << " is not allowed in non-ARM64EC files (" << toString(file)
                   << ")";
@@ -2295,6 +2297,13 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
        args.filtered(OPT_dependentloadflag, OPT_dependentloadflag_opt))
     parseDependentLoadFlags(arg);
 
+  for (auto *arg : args.filtered(OPT_arm64xsameaddress)) {
+    if (ctx.hybridSymtab)
+      parseSameAddress(arg->getValue());
+    else
+      Warn(ctx) << arg->getSpelling() << " is allowed only on EC targets";
+  }
+
   if (tar) {
     llvm::TimeTraceScope timeScope("Reproducer: response file");
     tar->append(
@@ -2668,12 +2677,46 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
     createECExportThunks();
 
   // Resolve remaining undefined symbols and warn about imported locals.
+  std::vector<Undefined *> aliases;
   ctx.forEachSymtab(
-      [&](SymbolTable &symtab) { symtab.resolveRemainingUndefines(); });
+      [&](SymbolTable &symtab) { symtab.resolveRemainingUndefines(aliases); });
 
   if (errorCount())
     return;
 
+  ctx.forEachActiveSymtab([](SymbolTable &symtab) {
+    symtab.initializeECThunks();
+    symtab.initializeLoadConfig();
+  });
+
+  // Identify unreferenced COMDAT sections.
+  if (config->doGC) {
+    if (config->mingw) {
+      // markLive doesn't traverse .eh_frame, but the personality function is
+      // only reached that way. The proper solution would be to parse and
+      // traverse the .eh_frame section, like the ELF linker does.
+      // For now, just manually try to retain the known possible personality
+      // functions. This doesn't bring in more object files, but only marks
+      // functions that already have been included to be retained.
+      ctx.forEachSymtab([&](SymbolTable &symtab) {
+        for (const char *n : {"__gxx_personality_v0", "__gcc_personality_v0",
+                              "rust_eh_personality"}) {
+          Defined *d = dyn_cast_or_null<Defined>(symtab.findUnderscore(n));
+          if (d && !d->isGCRoot) {
+            d->isGCRoot = true;
+            config->gcroot.push_back(d);
+          }
+        }
+      });
+    }
+
+    markLive(ctx);
+  }
+
+  ctx.symtab.initializeSameAddressThunks();
+  for (auto alias : aliases)
+    alias->resolveWeakAlias();
+
   if (config->mingw) {
     // Make sure the crtend.o object is the last object file. This object
     // file can contain terminating section chunks that need to be placed
@@ -2765,35 +2808,6 @@ void LinkerDriver::linkerMain(ArrayRef<const char *> argsArr) {
   if (auto *arg = args.getLastArg(OPT_print_symbol_order))
     config->printSymbolOrder = arg->getValue();
 
-  if (ctx.symtab.isEC())
-    ctx.symtab.initializeECThunks();
-  ctx.forEachActiveSymtab(
-      [](SymbolTable &symtab) { symtab.initializeLoadConfig(); });
-
-  // Identify unreferenced COMDAT sections.
-  if (config->doGC) {
-    if (config->mingw) {
-      // markLive doesn't traverse .eh_frame, but the personality function is
-      // only reached that way. The proper solution would be to parse and
-      // traverse the .eh_frame section, like the ELF linker does.
-      // For now, just manually try to retain the known possible personality
-      // functions. This doesn't bring in more object files, but only marks
-      // functions that already have been included to be retained.
-      ctx.forEachSymtab([&](SymbolTable &symtab) {
-        for (const char *n : {"__gxx_personality_v0", "__gcc_personality_v0",
-                              "rust_eh_personality"}) {
-          Defined *d = dyn_cast_or_null<Defined>(symtab.findUnderscore(n));
-          if (d && !d->isGCRoot) {
-            d->isGCRoot = true;
-            config->gcroot.push_back(d);
-          }
-        }
-      });
-    }
-
-    markLive(ctx);
-  }
-
   // Needs to happen after the last call to addFile().
   convertResources();
 

lld/COFF/Driver.h

@@ -214,6 +214,8 @@ class LinkerDriver {
   void parsePDBPageSize(StringRef);
   void parseSection(StringRef);
 
+  void parseSameAddress(StringRef);
+
   // Parses a MS-DOS stub file
   void parseDosStub(StringRef path);
 

lld/COFF/DriverUtils.cpp

@@ -328,6 +328,22 @@ void LinkerDriver::parseSwaprun(StringRef arg) {
   } while (!arg.empty());
 }
 
+void LinkerDriver::parseSameAddress(StringRef arg) {
+  auto mangledName = getArm64ECMangledFunctionName(arg);
+  Symbol *sym = ctx.symtab.addUndefined(mangledName ? *mangledName : arg);
+
+  // MSVC appears to generate thunks even for non-hybrid ARM64EC images.
+  // As a side effect, the native symbol is pulled in. Since this is used
+  // in the CRT for thread-local constructors, it results in the image
+  // containing unnecessary native code. As these thunks don't appear to
+  // be useful, we limit this behavior to actual hybrid targets. This may
+  // change if compatibility becomes necessary.
+  if (ctx.config.machine != ARM64X)
+    return;
+  Symbol *nativeSym = ctx.hybridSymtab->addUndefined(arg);
+  ctx.config.sameAddresses.emplace_back(sym, nativeSym);
+}
+
 // An RAII temporary file class that automatically removes a temporary file.
 namespace {
 class TemporaryFile {

lld/COFF/MarkLive.cpp

@@ -49,7 +49,10 @@ void markLive(COFFLinkerContext &ctx) {
       addSym(file->impchkThunk->exitThunk);
   };
 
-  addSym = [&](Symbol *b) {
+  addSym = [&](Symbol *s) {
+    Defined *b = s->getDefined();
+    if (!b)
+      return;
     if (auto *sym = dyn_cast<DefinedRegular>(b)) {
       enqueue(sym->getChunk());
     } else if (auto *sym = dyn_cast<DefinedImportData>(b)) {

lld/COFF/Options.td

@@ -31,6 +31,9 @@ multiclass B_priv<string name> {
 def align   : P<"align", "Section alignment">;
 def aligncomm : P<"aligncomm", "Set common symbol alignment">;
 def alternatename : P<"alternatename", "Define weak alias">;
+def arm64xsameaddress
+    : P<"arm64xsameaddress", "Generate a thunk for the symbol with the same "
+                             "address in both native and EC views on ARM64X.">;
 def base    : P<"base", "Base address of the program">;
 def color_diagnostics: Flag<["--"], "color-diagnostics">,
     HelpText<"Alias for --color-diagnostics=always">;
@@ -373,4 +376,3 @@ def tlbid : P_priv<"tlbid">;
 def tlbout : P_priv<"tlbout">;
 def verbose_all : P_priv<"verbose">;
 def guardsym : P_priv<"guardsym">;
-def arm64xsameaddress : P_priv<"arm64xsameaddress">;

lld/COFF/SymbolTable.cpp

@@ -452,7 +452,7 @@ void SymbolTable::reportUnresolvable() {
   reportProblemSymbols(undefs, /*localImports=*/nullptr, true);
 }
 
-void SymbolTable::resolveRemainingUndefines() {
+void SymbolTable::resolveRemainingUndefines(std::vector<Undefined *> &aliases) {
   llvm::TimeTraceScope timeScope("Resolve remaining undefined symbols");
   SmallPtrSet<Symbol *, 8> undefs;
   DenseMap<Symbol *, Symbol *> localImports;
@@ -468,8 +468,10 @@ void SymbolTable::resolveRemainingUndefines() {
     StringRef name = undef->getName();
 
     // A weak alias may have been resolved, so check for that.
-    if (undef->resolveWeakAlias())
+    if (undef->getWeakAlias()) {
+      aliases.push_back(undef);
       continue;
+    }
 
     // If we can resolve a symbol by removing __imp_ prefix, do that.
     // This odd rule is for compatibility with MSVC linker.
@@ -620,10 +622,10 @@ void SymbolTable::initializeECThunks() {
     return;
 
   for (auto it : entryThunks) {
-    auto *to = dyn_cast<Defined>(it.second);
+    Defined *to = it.second->getDefined();
     if (!to)
       continue;
-    auto *from = dyn_cast<DefinedRegular>(it.first);
+    auto *from = dyn_cast_or_null<DefinedRegular>(it.first->getDefined());
     // 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.
@@ -642,7 +644,8 @@ void SymbolTable::initializeECThunks() {
     Symbol *sym = exitThunks.lookup(file->thunkSym);
     if (!sym)
       sym = exitThunks.lookup(file->impECSym);
-    file->impchkThunk->exitThunk = dyn_cast_or_null<Defined>(sym);
+    if (sym)
+      file->impchkThunk->exitThunk = sym->getDefined();
   }
 
   // On ARM64EC, the __imp_ symbol references the auxiliary IAT, while the
@@ -659,6 +662,35 @@ void SymbolTable::initializeECThunks() {
   });
 }
 
+void SymbolTable::initializeSameAddressThunks() {
+  for (auto iter : ctx.config.sameAddresses) {
+    auto sym = dyn_cast_or_null<DefinedRegular>(iter.first->getDefined());
+    if (!sym || !sym->isLive())
+      continue;
+    auto nativeSym =
+        dyn_cast_or_null<DefinedRegular>(iter.second->getDefined());
+    if (!nativeSym || !nativeSym->isLive())
+      continue;
+    Defined *entryThunk = sym->getChunk()->getEntryThunk();
+    if (!entryThunk)
+      continue;
+
+    // Replace symbols with symbols referencing the thunk. Store the original
+    // symbol as equivalent DefinedSynthetic instances for use in the thunk
+    // itself.
+    auto symClone = make<DefinedSynthetic>(sym->getName(), sym->getChunk(),
+                                           sym->getValue());
+    auto nativeSymClone = make<DefinedSynthetic>(
+        nativeSym->getName(), nativeSym->getChunk(), nativeSym->getValue());
+    SameAddressThunkARM64EC *thunk =
+        make<SameAddressThunkARM64EC>(nativeSymClone, symClone, entryThunk);
+    sameAddressThunks.push_back(thunk);
+
+    replaceSymbol<DefinedSynthetic>(sym, sym->getName(), thunk);
+    replaceSymbol<DefinedSynthetic>(nativeSym, nativeSym->getName(), thunk);
+  }
+}
+
 Symbol *SymbolTable::addUndefined(StringRef name, InputFile *f,
                                   bool overrideLazy) {
   auto [s, wasInserted] = insert(name, f);