[LLD][COFF] Implement ARM64X relocations for the exception table

[cjacek]

No description provided.

[llvmbot]

@llvm/pr-subscribers-platform-windows
@llvm/pr-subscribers-lld-coff

@llvm/pr-subscribers-lld

Author: Jacek Caban (cjacek)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/123723.diff

4 Files Affected:

• (modified) lld/COFF/Chunks.cpp (+11-1)
• (modified) lld/COFF/Chunks.h (+6-1)
• (modified) lld/COFF/Writer.cpp (+20)
• (modified) lld/test/COFF/pdata-arm64ec.test (+16-5)

diff --git a/lld/COFF/Chunks.cpp b/lld/COFF/Chunks.cpp
index ff3c89884c24df..2ef74cb4ce6259 100644
--- a/lld/COFF/Chunks.cpp
+++ b/lld/COFF/Chunks.cpp
@@ -1167,7 +1167,7 @@ uint32_t ImportThunkChunkARM64EC::extendRanges() {
 }
 
 uint64_t Arm64XRelocVal::get() const {
-  return (sym ? sym->getRVA() : 0) + value;
+  return (sym ? sym->getRVA() : 0) + (chunk ? chunk->getRVA() : 0) + value;
 }
 
 size_t Arm64XDynamicRelocEntry::getSize() const {
@@ -1230,6 +1230,16 @@ void DynamicRelocsChunk::finalize() {
   size = alignTo(size, sizeof(uint32_t));
 }
 
+// Set the reloc value. The reloc entry must be allocated beforehand.
+void DynamicRelocsChunk::set(uint32_t rva, Arm64XRelocVal value) {
+  Arm64XDynamicRelocEntry &entry =
+      *llvm::find_if(arm64xRelocs, [rva](const Arm64XDynamicRelocEntry &e) {
+        return e.offset.get() == rva;
+      });
+  assert(!entry.value.get());
+  entry.value = value;
+}
+
 void DynamicRelocsChunk::writeTo(uint8_t *buf) const {
   auto table = reinterpret_cast<coff_dynamic_reloc_table *>(buf);
   table->Version = 1;
diff --git a/lld/COFF/Chunks.h b/lld/COFF/Chunks.h
index 7ba58e336451fc..d6216efdd90bdd 100644
--- a/lld/COFF/Chunks.h
+++ b/lld/COFF/Chunks.h
@@ -840,10 +840,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)
+      : chunk(chunk), value(offset) {}
   uint64_t get() const;
 
 private:
   Defined *sym = nullptr;
+  Chunk *chunk = nullptr;
   uint64_t value;
 };
 
@@ -874,10 +877,12 @@ class DynamicRelocsChunk : public NonSectionChunk {
   void finalize();
 
   void add(llvm::COFF::Arm64XFixupType type, uint8_t size,
-           Arm64XRelocVal offset, Arm64XRelocVal value) {
+           Arm64XRelocVal offset, Arm64XRelocVal value = Arm64XRelocVal()) {
     arm64xRelocs.emplace_back(type, size, offset, value);
   }
 
+  void set(uint32_t rva, Arm64XRelocVal value);
+
 private:
   std::vector<Arm64XDynamicRelocEntry> arm64xRelocs;
   size_t size;
diff --git a/lld/COFF/Writer.cpp b/lld/COFF/Writer.cpp
index 3d95d219a493cd..507621d1e95f80 100644
--- a/lld/COFF/Writer.cpp
+++ b/lld/COFF/Writer.cpp
@@ -2392,6 +2392,13 @@ void Writer::setECSymbols() {
       symtab->findUnderscore("__arm64x_native_entrypoint")
           ->replaceKeepingName(altEntrySym, sizeof(SymbolUnion));
     }
+
+    if (hybridPdata.first)
+      ctx.dynamicRelocs->set(
+          dataDirOffset64 + EXCEPTION_TABLE * sizeof(data_directory) +
+              offsetof(data_directory, Size),
+          hybridPdata.last->getRVA() - hybridPdata.first->getRVA() +
+              hybridPdata.last->getSize());
   }
 }
 
@@ -2644,6 +2651,19 @@ void Writer::createDynamicRelocs() {
       Warn(ctx) << "'__chpe_metadata' is missing for ARM64X target";
   }
 
+  if (pdata.first != hybridPdata.first) {
+    ctx.dynamicRelocs->add(IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE, sizeof(uint32_t),
+                           dataDirOffset64 +
+                               EXCEPTION_TABLE * sizeof(data_directory) +
+                               offsetof(data_directory, RelativeVirtualAddress),
+                           hybridPdata.first);
+    // The Size value is assigned after addresses are finalized.
+    ctx.dynamicRelocs->add(IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE, sizeof(uint32_t),
+                           dataDirOffset64 +
+                               EXCEPTION_TABLE * sizeof(data_directory) +
+                               offsetof(data_directory, Size));
+  }
+
   // Set the hybrid load config to the EC load config.
   ctx.dynamicRelocs->add(IMAGE_DVRT_ARM64X_FIXUP_TYPE_VALUE, sizeof(uint32_t),
                          dataDirOffset64 +
diff --git a/lld/test/COFF/pdata-arm64ec.test b/lld/test/COFF/pdata-arm64ec.test
index 7f20c460dc1099..fbec797525f7f8 100644
--- a/lld/test/COFF/pdata-arm64ec.test
+++ b/lld/test/COFF/pdata-arm64ec.test
@@ -6,6 +6,7 @@ Test handlign of hybrid .pdata section on ARM64EC target.
 RUN: llvm-mc -filetype=obj -triple=arm64-windows arm64-func-sym.s -o arm64-func-sym.obj
 RUN: llvm-mc -filetype=obj -triple=arm64ec-windows arm64ec-func-sym.s -o arm64ec-func-sym.obj
 RUN: llvm-mc -filetype=obj -triple=x86_64-windows x86_64-func-sym.s -o x86_64-func-sym.obj
+RUN: llvm-mc -filetype=obj -triple=aarch64-windows %S/Inputs/loadconfig-arm64.s -o loadconfig-arm64.obj
 RUN: llvm-mc -filetype=obj -triple=arm64ec-windows %p/Inputs/loadconfig-arm64ec.s -o loadconfig-arm64ec.obj
 
 Only arm64ec code:
@@ -55,11 +56,21 @@ DATA3: 180005000 00100000 11000001 00200000 0e200000
 Mixed arm64x code:
 
 RUN: lld-link -out:test4.dll -machine:arm64x arm64-func-sym.obj arm64ec-func-sym.obj \
-RUN:          x86_64-func-sym.obj loadconfig-arm64ec.obj -dll -noentry
+RUN:          x86_64-func-sym.obj loadconfig-arm64.obj loadconfig-arm64ec.obj -dll -noentry
 
 RUN: llvm-readobj --headers test4.dll | FileCheck -check-prefix=DIR3 %s
-DIR3:      ExceptionTableRVA: 0x6000
-DIR3-NEXT: ExceptionTableSize: 0x10
+DIR3:      ImageOptionalHeader {
+DIR3:        DataDirectory {
+DIR3:          ExceptionTableRVA: 0x6000
+DIR3-NEXT:     ExceptionTableSize: 0x10
+DIR3:        }
+DIR3:      }
+DIR3:      HybridObject {
+DIR3:        ImageOptionalHeader {
+DIR3:          ExceptionTableRVA: 0x6010
+DIR3-NEXT:     ExceptionTableSize: 0xC
+DIR3:        }
+DIR3:      }
 
 RUN: llvm-objdump -s --section=.pdata test4.dll | FileCheck -check-prefix=DATA4 %s
 DATA4: 180006000 00100000 11000001 00200000 11000001  ......... ......
@@ -74,12 +85,12 @@ RUN: llvm-readobj --headers test5.dll | FileCheck -check-prefix=DIR2 %s
 RUN: llvm-objdump -s --section=.pdata test5.dll | FileCheck -check-prefix=DATA3 %s
 
 RUN: lld-link -out:test6.dll -machine:arm64x arm64ec-func-sym.obj x86_64-func-sym.obj \
-RUN:          arm64-func-sym.obj loadconfig-arm64ec.obj -dll -noentry
+RUN:          arm64-func-sym.obj loadconfig-arm64.obj loadconfig-arm64ec.obj -dll -noentry
 RUN: llvm-readobj --headers test6.dll | FileCheck -check-prefix=DIR3 %s
 RUN: llvm-objdump -s --section=.pdata test6.dll | FileCheck -check-prefix=DATA4 %s
 
 RUN: lld-link -out:test7.dll -machine:arm64x x86_64-func-sym.obj arm64ec-func-sym.obj \
-RUN:          arm64-func-sym.obj loadconfig-arm64ec.obj -dll -noentry
+RUN:          arm64-func-sym.obj loadconfig-arm64.obj loadconfig-arm64ec.obj -dll -noentry
 RUN: llvm-readobj --headers test7.dll | FileCheck -check-prefix=DIR3 %s
 RUN: llvm-objdump -s --section=.pdata test7.dll | FileCheck -check-prefix=DATA4 %s
 

[cjacek]

This introduces a new mechanism to postpone setting relocation values until addresses are finalized, which is useful for data directory size fields. A similar approach will be needed for export tables.

The mechanism is based on RVAs, as they can be easily recalculated when setting the value, avoiding the need to store references (something that could be problematic in combination with sorting). Since this function will be called only twice in total (once in this PR), performance is not a concern.

[mstorsjo]

Looks mostly good, one minor nit.

[mstorsjo]

I guess this assumes that this does find an entry. If it doesn't, does it trigger UB or something else - can we receive the returned iterator(?) and check it with an assert, before dereferencing it?

[cjacek]

I added an assert and merged. Thanks!