[llvm-lib][Object][COFF] Use ARM64 machine type for import library descriptor objects. (#78537)

1 files changed, 35 insertions, 26 deletions

diff --git a/llvm/lib/Object/COFFImportFile.cpp b/llvm/lib/Object/COFFImportFile.cpp
index b60e32f..60556c1 100644
--- a/llvm/lib/Object/COFFImportFile.cpp
+++ b/llvm/lib/Object/COFFImportFile.cpp
@@ -150,7 +150,7 @@ namespace {
 class ObjectFactory {
   using u16 = support::ulittle16_t;
   using u32 = support::ulittle32_t;
-  MachineTypes Machine;
+  MachineTypes NativeMachine;
   BumpPtrAllocator Alloc;
   StringRef ImportName;
   StringRef Library;
@@ -159,7 +159,7 @@ class ObjectFactory {
 
 public:
   ObjectFactory(StringRef S, MachineTypes M)
-      : Machine(M), ImportName(S), Library(llvm::sys::path::stem(S)),
+      : NativeMachine(M), ImportName(S), Library(llvm::sys::path::stem(S)),
         ImportDescriptorSymbolName(("__IMPORT_DESCRIPTOR_" + Library).str()),
         NullThunkSymbolName(("\x7f" + Library + "_NULL_THUNK_DATA").str()) {}
 
@@ -182,10 +182,14 @@ public:
   // Create a short import file which is described in PE/COFF spec 7. Import
   // Library Format.
   NewArchiveMember createShortImport(StringRef Sym, uint16_t Ordinal,
-                                     ImportType Type, ImportNameType NameType);
+                                     ImportType Type, ImportNameType NameType,
+                                     MachineTypes Machine);
 
   // Create a weak external file which is described in PE/COFF Aux Format 3.
-  NewArchiveMember createWeakExternal(StringRef Sym, StringRef Weak, bool Imp);
+  NewArchiveMember createWeakExternal(StringRef Sym, StringRef Weak, bool Imp,
+                                      MachineTypes Machine);
+
+  bool is64Bit() const { return COFF::is64Bit(NativeMachine); }
 };
 } // namespace
 
@@ -197,7 +201,7 @@ ObjectFactory::createImportDescriptor(std::vector<uint8_t> &Buffer) {
 
   // COFF Header
   coff_file_header Header{
-      u16(Machine),
+      u16(NativeMachine),
       u16(NumberOfSections),
       u32(0),
       u32(sizeof(Header) + (NumberOfSections * sizeof(coff_section)) +
@@ -208,7 +212,7 @@ ObjectFactory::createImportDescriptor(std::vector<uint8_t> &Buffer) {
           (ImportName.size() + 1)),
       u32(NumberOfSymbols),
       u16(0),
-      u16(is64Bit(Machine) ? C_Invalid : IMAGE_FILE_32BIT_MACHINE),
+      u16(is64Bit() ? C_Invalid : IMAGE_FILE_32BIT_MACHINE),
   };
   append(Buffer, Header);
 
@@ -250,11 +254,11 @@ ObjectFactory::createImportDescriptor(std::vector<uint8_t> &Buffer) {
 
   const coff_relocation RelocationTable[NumberOfRelocations] = {
       {u32(offsetof(coff_import_directory_table_entry, NameRVA)), u32(2),
-       u16(getImgRelRelocation(Machine))},
+       u16(getImgRelRelocation(NativeMachine))},
       {u32(offsetof(coff_import_directory_table_entry, ImportLookupTableRVA)),
-       u32(3), u16(getImgRelRelocation(Machine))},
+       u32(3), u16(getImgRelRelocation(NativeMachine))},
       {u32(offsetof(coff_import_directory_table_entry, ImportAddressTableRVA)),
-       u32(4), u16(getImgRelRelocation(Machine))},
+       u32(4), u16(getImgRelRelocation(NativeMachine))},
   };
   append(Buffer, RelocationTable);
 
@@ -336,7 +340,7 @@ ObjectFactory::createNullImportDescriptor(std::vector<uint8_t> &Buffer) {
 
   // COFF Header
   coff_file_header Header{
-      u16(Machine),
+      u16(NativeMachine),
       u16(NumberOfSections),
       u32(0),
       u32(sizeof(Header) + (NumberOfSections * sizeof(coff_section)) +
@@ -344,7 +348,7 @@ ObjectFactory::createNullImportDescriptor(std::vector<uint8_t> &Buffer) {
           sizeof(coff_import_directory_table_entry)),
       u32(NumberOfSymbols),
       u16(0),
-      u16(is64Bit(Machine) ? C_Invalid : IMAGE_FILE_32BIT_MACHINE),
+      u16(is64Bit() ? C_Invalid : IMAGE_FILE_32BIT_MACHINE),
   };
   append(Buffer, Header);
 
@@ -393,11 +397,11 @@ ObjectFactory::createNullImportDescriptor(std::vector<uint8_t> &Buffer) {
 NewArchiveMember ObjectFactory::createNullThunk(std::vector<uint8_t> &Buffer) {
   const uint32_t NumberOfSections = 2;
   const uint32_t NumberOfSymbols = 1;
-  uint32_t VASize = is64Bit(Machine) ? 8 : 4;
+  uint32_t VASize = is64Bit() ? 8 : 4;
 
   // COFF Header
   coff_file_header Header{
-      u16(Machine),
+      u16(NativeMachine),
       u16(NumberOfSections),
       u32(0),
       u32(sizeof(Header) + (NumberOfSections * sizeof(coff_section)) +
@@ -407,7 +411,7 @@ NewArchiveMember ObjectFactory::createNullThunk(std::vector<uint8_t> &Buffer) {
           VASize),
       u32(NumberOfSymbols),
       u16(0),
-      u16(is64Bit(Machine) ? C_Invalid : IMAGE_FILE_32BIT_MACHINE),
+      u16(is64Bit() ? C_Invalid : IMAGE_FILE_32BIT_MACHINE),
   };
   append(Buffer, Header);
 
@@ -422,8 +426,7 @@ NewArchiveMember ObjectFactory::createNullThunk(std::vector<uint8_t> &Buffer) {
        u32(0),
        u16(0),
        u16(0),
-       u32((is64Bit(Machine) ? IMAGE_SCN_ALIGN_8BYTES
-                             : IMAGE_SCN_ALIGN_4BYTES) |
+       u32((is64Bit() ? IMAGE_SCN_ALIGN_8BYTES : IMAGE_SCN_ALIGN_4BYTES) |
            IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ |
            IMAGE_SCN_MEM_WRITE)},
       {{'.', 'i', 'd', 'a', 't', 'a', '$', '4'},
@@ -436,8 +439,7 @@ NewArchiveMember ObjectFactory::createNullThunk(std::vector<uint8_t> &Buffer) {
        u32(0),
        u16(0),
        u16(0),
-       u32((is64Bit(Machine) ? IMAGE_SCN_ALIGN_8BYTES
-                             : IMAGE_SCN_ALIGN_4BYTES) |
+       u32((is64Bit() ? IMAGE_SCN_ALIGN_8BYTES : IMAGE_SCN_ALIGN_4BYTES) |
            IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ |
            IMAGE_SCN_MEM_WRITE)},
   };
@@ -445,12 +447,12 @@ NewArchiveMember ObjectFactory::createNullThunk(std::vector<uint8_t> &Buffer) {
 
   // .idata$5, ILT
   append(Buffer, u32(0));
-  if (is64Bit(Machine))
+  if (is64Bit())
     append(Buffer, u32(0));
 
   // .idata$4, IAT
   append(Buffer, u32(0));
-  if (is64Bit(Machine))
+  if (is64Bit())
     append(Buffer, u32(0));
 
   // Symbol Table
@@ -475,7 +477,8 @@ NewArchiveMember ObjectFactory::createNullThunk(std::vector<uint8_t> &Buffer) {
 NewArchiveMember ObjectFactory::createShortImport(StringRef Sym,
                                                   uint16_t Ordinal,
                                                   ImportType ImportType,
-                                                  ImportNameType NameType) {
+                                                  ImportNameType NameType,
+                                                  MachineTypes Machine) {
   size_t ImpSize = ImportName.size() + Sym.size() + 2; // +2 for NULs
   size_t Size = sizeof(coff_import_header) + ImpSize;
   char *Buf = Alloc.Allocate<char>(Size);
@@ -501,7 +504,8 @@ NewArchiveMember ObjectFactory::createShortImport(StringRef Sym,
 }
 
 NewArchiveMember ObjectFactory::createWeakExternal(StringRef Sym,
-                                                   StringRef Weak, bool Imp) {
+                                                   StringRef Weak, bool Imp,
+                                                   MachineTypes Machine) {
   std::vector<uint8_t> Buffer;
   const uint32_t NumberOfSections = 1;
   const uint32_t NumberOfSymbols = 5;
@@ -585,8 +589,11 @@ Error writeImportLibrary(StringRef ImportName, StringRef Path,
                          ArrayRef<COFFShortExport> Exports,
                          MachineTypes Machine, bool MinGW) {
 
+  MachineTypes NativeMachine =
+      isArm64EC(Machine) ? IMAGE_FILE_MACHINE_ARM64 : Machine;
+
   std::vector<NewArchiveMember> Members;
-  ObjectFactory OF(llvm::sys::path::filename(ImportName), Machine);
+  ObjectFactory OF(llvm::sys::path::filename(ImportName), NativeMachine);
 
   std::vector<uint8_t> ImportDescriptor;
   Members.push_back(OF.createImportDescriptor(ImportDescriptor));
@@ -620,13 +627,15 @@ Error writeImportLibrary(StringRef ImportName, StringRef Path,
       return Name.takeError();
 
     if (!E.AliasTarget.empty() && *Name != E.AliasTarget) {
-      Members.push_back(OF.createWeakExternal(E.AliasTarget, *Name, false));
-      Members.push_back(OF.createWeakExternal(E.AliasTarget, *Name, true));
+      Members.push_back(
+          OF.createWeakExternal(E.AliasTarget, *Name, false, Machine));
+      Members.push_back(
+          OF.createWeakExternal(E.AliasTarget, *Name, true, Machine));
       continue;
     }
 
     Members.push_back(
-        OF.createShortImport(*Name, E.Ordinal, ImportType, NameType));
+        OF.createShortImport(*Name, E.Ordinal, ImportType, NameType, Machine));
   }
 
   return writeArchive(Path, Members, SymtabWritingMode::NormalSymtab,