[LLD][COFF] Add support for ARM64X same-address thunks (#151255)

Fixes MSVC CRT thread-local constructors support on hybrid ARM64X
targets.

`-arm64xsameaddress` is an undocumented option that ensures the
specified function has the same address in both native and EC views of
hybrid images. To achieve this, the linker emits additional thunks and
replaces the symbols
of those functions with the thunk symbol (the same thunk is used in both
views). The thunk code jumps to the native function (similar to range
extension thunks), but additional ARM64X relocations are emitted to
replace the target with the EC function in the EC view.

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. Because these thunks do not appear to be useful
in that context, we limit this behavior to actual hybrid targets. This
may change if compatibility requires it.

The tricky part is that thunks should be skipped if the symbol is not
live in either view, and symbol replacement must be reflected in weak
aliases. This requires thunk generation to happen before resolving weak
aliases but after the GC pass. To enable this, the `markLive` call was
moved earlier, and the final weak alias resolution was postponed until
afterward. This requires more code to be aware of weak aliases, which
previously could assume they were already resolved.

Author: cjacek

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);