[lld] [LLD][COFF] Implement support for hybrid IAT on ARM64X (PR #124189)

[Jacek Caban via llvm-commits]

================
@@ -717,52 +717,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,
----------------
cjacek wrote:

> I'm trying to form an understanding of the expected end state of things after this loop.
> 
> At the end, we're going to have the regular IAT (`lookups` here) look like this:
> 
> ```
> [native]...[native][hybrid]...[hybrid][null chunk][second EC-only]...[more EC-only][lookupTerminator]
> ```
> 
> I presume the purpose of the null chunk in the middle of the regular IAT is to terminate it, so in native mode you only see the native-only and hybrid entries? After applying dynamic relocations, the null chunk is turned into the entry it was initially meant to be, so you see all of it - and we adjust the pointer to the start, so that it points at the first non-native entry.

Yes, it's all correct.

> What's the idea and purpose of the lookupTerminator? In the native view, it's not visible, and after relocations in EC mode, it's turned into a null terminator. Why not just do a plain null terminator to begin with?
> 
> Secondly, if I follow it correctly, the aux IAT is supposed to look like this:
> 
> ```
> [null chunks for native-only][aux entries for hybrid imports][aux entries for ec-only imports][null chunk]
> ```

I think you're correct, using a null chunk should work here. However, that's not what MSVC does. I'm not sure why; I can only speculate. My best guess is that a null chunk might confuse some applications that manually read the IAT (such as anti-cheat systems or DRM software). Following MSVC's approach ensures that the number of null chunks in `.idata$4` and `.idata$5` matches the number of different imported DLLs, and that all hint chunks are referenced from the IAT in both the native and EC views, even if they are referenced from an entry that is inactive in the native view.

Since this is just speculation, I'm open to changing it to null chunks if you think strictly following MSVC isn't necessary.

> I've forgotten the purpose/role of the AUX IAT.

On ARM64EC, the regular IAT contains pointers to imported functions as exposed by PE exports. These may be x86 export thunks that perform FFS jumps to EC code, or they may simply be x86 functions. The purpose of the auxiliary IAT is to provide `__imp` symbols in a form that can be directly called from EC code (possibly by pointing to a thunk that performs a runtime check on the regular IAT pointer).

> Doesn't it use null pointers as terminators - how does this work? Or does the delta for the start of the regular IAT also get applied to this one?

The delta from the start of the IAT needs to match. The auxiliary IAT is referenced in CHPE metadata as an RVA to the entire table. The loader may need to associate its entries with those in the regular IAT, potentially using the same offset from the table's start as in the regular IAT. To ensure this works, we need to populate the auxiliary IAT with native-only imports, even though they are not used at runtime.

https://github.com/llvm/llvm-project/pull/124189