[LLD][COFF] Implement support for hybrid IAT on ARM64X

In hybrid images, the PE header references a single IAT for both native and EC views,
merging entries where possible. When merging isn't feasible, different imports
are grouped together, and ARM64X relocations are emitted as needed.

Author: cjacek

lld/COFF/Chunks.cpp

@@ -1172,11 +1172,12 @@ uint64_t Arm64XRelocVal::get() const {
 
 size_t Arm64XDynamicRelocEntry::getSize() const {
   switch (type) {
+  case IMAGE_DVRT_ARM64X_FIXUP_TYPE_ZEROFILL:
+    return sizeof(uint16_t); // Just a header.
   case IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE:
     return sizeof(uint16_t) + size; // A header and a payload.
   case IMAGE_DVRT_ARM64X_FIXUP_TYPE_DELTA:
-  case IMAGE_DVRT_ARM64X_FIXUP_TYPE_ZEROFILL:
-    llvm_unreachable("unsupported type");
+    return 2 * sizeof(uint16_t); // A header and a delta.
   }
   llvm_unreachable("invalid type");
 }
@@ -1186,6 +1187,9 @@ void Arm64XDynamicRelocEntry::writeTo(uint8_t *buf) const {
   *out = (offset.get() & 0xfff) | (type << 12);
 
   switch (type) {
+  case IMAGE_DVRT_ARM64X_FIXUP_TYPE_ZEROFILL:
+    *out |= ((bit_width(size) - 1) << 14); // Encode the size.
+    break;
   case IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE:
     *out |= ((bit_width(size) - 1) << 14); // Encode the size.
     switch (size) {
@@ -1203,8 +1207,23 @@ void Arm64XDynamicRelocEntry::writeTo(uint8_t *buf) const {
     }
     break;
   case IMAGE_DVRT_ARM64X_FIXUP_TYPE_DELTA:
-  case IMAGE_DVRT_ARM64X_FIXUP_TYPE_ZEROFILL:
-    llvm_unreachable("unsupported type");
+    int delta = value.get();
+    // Negative offsets use a sign bit in the header.
+    if (delta < 0) {
+      *out |= 1 << 14;
+      delta = -delta;
+    }
+    // Depending on the value, the delta is encoded with a shift of 2 or 3 bits.
+    if (delta & 7) {
+      assert(!(delta & 3));
+      delta >>= 2;
+    } else {
+      *out |= (1 << 15);
+      delta >>= 3;
+    }
+    out[1] = delta;
+    assert(!(delta & ~0xffff));
+    break;
   }
 }
 

lld/COFF/DLL.cpp

@@ -716,52 +716,155 @@ class ExportOrdinalChunk : public NonSectionChunk {
 void IdataContents::create(COFFLinkerContext &ctx) {
   std::vector<std::vector<DefinedImportData *>> v = binImports(ctx, imports);
 
+  // Merge compatible EC and native import files in hybrid images.
+  if (ctx.hybridSymtab) {
+    for (std::vector<DefinedImportData *> &syms : v) {
+      // At this point, symbols are sorted by base name, ensuring that
+      // compatible import files, if present, are adjacent.
+      std::vector<DefinedImportData *> hybridSyms;
+      ImportFile *prev = nullptr;
+      for (DefinedImportData *sym : syms) {
+        ImportFile *file = sym->file;
+        if (!prev || file->isEC() == prev->isEC() ||
+            !file->isSameImport(prev)) {
+          hybridSyms.push_back(sym);
+          prev = file;
+          continue;
+        }
+
+        // The native variant exposes a subset of EC symbols and chunks. Use the
+        // EC variant to represent both.
+        if (file->isEC()) {
+          hybridSyms.pop_back();
+          hybridSyms.push_back(sym);
+        }
+
+        prev->hybridFile = file;
+        file->hybridFile = prev;
+        prev = nullptr;
+      }
+
+      // Sort symbols by type: native-only files first, followed by merged
+      // hybrid files, and then EC-only files.
+      llvm::stable_sort(hybridSyms,
+                        [](DefinedImportData *a, DefinedImportData *b) {
+                          if (a->file->hybridFile)
+                            return !b->file->hybridFile && b->file->isEC();
+                          return !a->file->isEC() && b->file->isEC();
+                        });
+      syms = std::move(hybridSyms);
+    }
+  }
+
   // Create .idata contents for each DLL.
   for (std::vector<DefinedImportData *> &syms : v) {
     // Create lookup and address tables. If they have external names,
     // we need to create hintName chunks to store the names.
     // If they don't (if they are import-by-ordinals), we store only
     // ordinal values to the table.
     size_t base = lookups.size();
+    Chunk *lookupsTerminator = nullptr, *addressesTerminator = nullptr;
     for (DefinedImportData *s : syms) {
       uint16_t ord = s->getOrdinal();
+      HintNameChunk *hintChunk = nullptr;
+      Chunk *lookupsChunk, *addressesChunk;
+
       if (s->getExternalName().empty()) {
-        lookups.push_back(make<OrdinalOnlyChunk>(ctx, ord));
-        addresses.push_back(make<OrdinalOnlyChunk>(ctx, ord));
+        lookupsChunk = make<OrdinalOnlyChunk>(ctx, ord);
+        addressesChunk = make<OrdinalOnlyChunk>(ctx, ord);
       } else {
-        auto *c = make<HintNameChunk>(s->getExternalName(), ord);
-        lookups.push_back(make<LookupChunk>(ctx, c));
-        addresses.push_back(make<LookupChunk>(ctx, c));
-        hints.push_back(c);
+        hintChunk = make<HintNameChunk>(s->getExternalName(), ord);
+        lookupsChunk = make<LookupChunk>(ctx, hintChunk);
+        addressesChunk = make<LookupChunk>(ctx, hintChunk);
+        hints.push_back(hintChunk);
       }
 
-      if (s->file->impECSym) {
+      // Detect the first EC-only import in the hybrid IAT. Emit null chunks
+      // and add an ARM64X relocation to replace it with the import for the EC
+      // view. Additionally, use the original chunks as import terminators
+      // and zero them with ARM64X relocations. Since these chunks appear
+      // after the null terminator in the native view, they are always ignored
+      // by the loader. However, MSVC emits them for some reason.
+      if (ctx.hybridSymtab && !lookupsTerminator && s->file->isEC() &&
+          !s->file->hybridFile) {
+        lookupsTerminator = lookupsChunk;
+        addressesTerminator = addressesChunk;
+        lookupsChunk = make<NullChunk>(ctx);
+        addressesChunk = make<NullChunk>(ctx);
+
+        Arm64XRelocVal relocVal = hintChunk;
+        if (!hintChunk)
+          relocVal = (1ULL << 63) | ord;
+        ctx.dynamicRelocs->add(IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE,
+                               sizeof(uint64_t), lookupsChunk, relocVal);
+        ctx.dynamicRelocs->add(IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE,
+                               sizeof(uint64_t), addressesChunk, relocVal);
+        ctx.dynamicRelocs->add(IMAGE_DVRT_ARM64X_FIXUP_TYPE_ZEROFILL,
+                               sizeof(uint64_t), lookupsTerminator);
+        ctx.dynamicRelocs->add(IMAGE_DVRT_ARM64X_FIXUP_TYPE_ZEROFILL,
+                               sizeof(uint64_t), addressesTerminator);
+      }
+
+      lookups.push_back(lookupsChunk);
+      addresses.push_back(addressesChunk);
+
+      if (s->file->isEC()) {
         auto chunk = make<AuxImportChunk>(s->file);
         auxIat.push_back(chunk);
         s->file->impECSym->setLocation(chunk);
 
         chunk = make<AuxImportChunk>(s->file);
         auxIatCopy.push_back(chunk);
         s->file->auxImpCopySym->setLocation(chunk);
+      } else if (ctx.hybridSymtab) {
+        // Fill the auxiliary IAT with null chunks for native-only imports.
+        auxIat.push_back(make<NullChunk>(ctx));
+        auxIatCopy.push_back(make<NullChunk>(ctx));
       }
     }
     // Terminate with null values.
-    lookups.push_back(make<NullChunk>(ctx));
-    addresses.push_back(make<NullChunk>(ctx));
-    if (ctx.config.machine == ARM64EC) {
+    lookups.push_back(lookupsTerminator ? lookupsTerminator
+                                        : make<NullChunk>(ctx));
+    addresses.push_back(addressesTerminator ? addressesTerminator
+                                            : make<NullChunk>(ctx));
+    if (ctx.symtabEC) {
       auxIat.push_back(make<NullChunk>(ctx));
       auxIatCopy.push_back(make<NullChunk>(ctx));
     }
 
-    for (int i = 0, e = syms.size(); i < e; ++i)
+    for (int i = 0, e = syms.size(); i < e; ++i) {
       syms[i]->setLocation(addresses[base + i]);
+      if (syms[i]->file->hybridFile)
+        syms[i]->file->hybridFile->impSym->setLocation(addresses[base + i]);
+    }
 
     // Create the import table header.
     dllNames.push_back(make<StringChunk>(syms[0]->getDLLName()));
     auto *dir = make<ImportDirectoryChunk>(dllNames.back());
     dir->lookupTab = lookups[base];
     dir->addressTab = addresses[base];
     dirs.push_back(dir);
+
+    if (ctx.hybridSymtab) {
+      // If native-only imports exist, emit ARM64X relocations, skipping them in
+      // the EC view.
+      uint32_t nativeOnly =
+          llvm::find_if(syms,
+                        [](DefinedImportData *s) { return s->file->isEC(); }) -
+          syms.begin();
+      if (nativeOnly) {
+        ctx.dynamicRelocs->add(
+            IMAGE_DVRT_ARM64X_FIXUP_TYPE_DELTA, 0,
+            Arm64XRelocVal(
+                dir, offsetof(ImportDirectoryTableEntry, ImportLookupTableRVA)),
+            nativeOnly * sizeof(uint64_t));
+        ctx.dynamicRelocs->add(
+            IMAGE_DVRT_ARM64X_FIXUP_TYPE_DELTA, 0,
+            Arm64XRelocVal(dir, offsetof(ImportDirectoryTableEntry,
+                                         ImportAddressTableRVA)),
+            nativeOnly * sizeof(uint64_t));
+      }
+    }
   }
   // Add null terminator.
   dirs.push_back(make<NullChunk>(sizeof(ImportDirectoryTableEntry), 4));

lld/COFF/InputFiles.cpp

@@ -1100,15 +1100,21 @@ void ObjFile::enqueuePdbFile(StringRef path, ObjFile *fromFile) {
 }
 
 ImportFile::ImportFile(COFFLinkerContext &ctx, MemoryBufferRef m)
-    : InputFile(ctx.symtab, ImportKind, m), live(!ctx.config.doGC) {}
+    : InputFile(ctx.getSymtab(getMachineType(m)), ImportKind, m),
+      live(!ctx.config.doGC) {}
 
-MachineTypes ImportFile::getMachineType() const {
+MachineTypes ImportFile::getMachineType(MemoryBufferRef m) {
   uint16_t machine =
-      reinterpret_cast<const coff_import_header *>(mb.getBufferStart())
-          ->Machine;
+      reinterpret_cast<const coff_import_header *>(m.getBufferStart())->Machine;
   return MachineTypes(machine);
 }
 
+bool ImportFile::isSameImport(const ImportFile *other) const {
+  if (!externalName.empty())
+    return other->externalName == externalName;
+  return hdr->OrdinalHint == other->hdr->OrdinalHint;
+}
+
 ImportThunkChunk *ImportFile::makeImportThunk() {
   switch (hdr->Machine) {
   case AMD64:

lld/COFF/InputFiles.h

@@ -351,11 +351,15 @@ class ImportFile : public InputFile {
   explicit ImportFile(COFFLinkerContext &ctx, MemoryBufferRef m);
 
   static bool classof(const InputFile *f) { return f->kind() == ImportKind; }
-  MachineTypes getMachineType() const override;
+  MachineTypes getMachineType() const override { return getMachineType(mb); }
+  static MachineTypes getMachineType(MemoryBufferRef m);
+  bool isSameImport(const ImportFile *other) const;
+  bool isEC() const { return impECSym != nullptr; }
 
   DefinedImportData *impSym = nullptr;
   Defined *thunkSym = nullptr;
   ImportThunkChunkARM64EC *impchkThunk = nullptr;
+  ImportFile *hybridFile = nullptr;
   std::string dllName;
 
 private: