[llvm-ar][Object][COFF] Add support for EC symbols to llvm-ar. (llvm#85230)

cjacek · chencha3 · commit 2a9e13405172 · 2024-03-22T14:49:45.000-05:00
Make writeArchive IsEC argument optional and use EC symbol map when indicated by input object files.
diff --git a/llvm/include/llvm/Object/ArchiveWriter.h b/llvm/include/llvm/Object/ArchiveWriter.h
@@ -52,7 +52,7 @@ Error writeArchive(StringRef ArcName, ArrayRef<NewArchiveMember> NewMembers,
                    SymtabWritingMode WriteSymtab, object::Archive::Kind Kind,
                    bool Deterministic, bool Thin,
                    std::unique_ptr<MemoryBuffer> OldArchiveBuf = nullptr,
-                   bool IsEC = false);
+                   std::optional<bool> IsEC = std::nullopt);
 
 // writeArchiveToBuffer is similar to writeArchive but returns the Archive in a
 // buffer instead of writing it out to a file.
diff --git a/llvm/lib/Object/ArchiveWriter.cpp b/llvm/lib/Object/ArchiveWriter.cpp
@@ -48,7 +48,7 @@ using namespace llvm;
 using namespace llvm::object;
 
 struct SymMap {
-  bool UseECMap;
+  bool UseECMap = false;
   std::map<std::string, uint16_t> Map;
   std::map<std::string, uint16_t> ECMap;
 };
@@ -678,6 +678,25 @@ static bool isECObject(object::SymbolicFile &Obj) {
   return false;
 }
 
+static bool isAnyArm64COFF(object::SymbolicFile &Obj) {
+  if (Obj.isCOFF())
+    return COFF::isAnyArm64(cast<COFFObjectFile>(&Obj)->getMachine());
+
+  if (Obj.isCOFFImportFile())
+    return COFF::isAnyArm64(cast<COFFImportFile>(&Obj)->getMachine());
+
+  if (Obj.isIR()) {
+    Expected<std::string> TripleStr =
+        getBitcodeTargetTriple(Obj.getMemoryBufferRef());
+    if (!TripleStr)
+      return false;
+    Triple T(*TripleStr);
+    return T.isOSWindows() && T.getArch() == Triple::aarch64;
+  }
+
+  return false;
+}
+
 bool isImportDescriptor(StringRef Name) {
   return Name.starts_with(ImportDescriptorPrefix) ||
          Name == StringRef{NullImportDescriptorSymbolName} ||
@@ -731,7 +750,8 @@ static Expected<std::vector<MemberData>>
 computeMemberData(raw_ostream &StringTable, raw_ostream &SymNames,
                   object::Archive::Kind Kind, bool Thin, bool Deterministic,
                   SymtabWritingMode NeedSymbols, SymMap *SymMap,
-                  LLVMContext &Context, ArrayRef<NewArchiveMember> NewMembers) {
+                  LLVMContext &Context, ArrayRef<NewArchiveMember> NewMembers,
+                  std::optional<bool> IsEC) {
   static char PaddingData[8] = {'\n', '\n', '\n', '\n', '\n', '\n', '\n', '\n'};
   uint64_t MemHeadPadSize = 0;
   uint64_t Pos =
@@ -807,6 +827,30 @@ computeMemberData(raw_ostream &StringTable, raw_ostream &SymNames,
     }
   }
 
+  if (SymMap) {
+    if (IsEC) {
+      SymMap->UseECMap = *IsEC;
+    } else {
+      // When IsEC is not specified by the caller, use it when we have both
+      // any ARM64 object (ARM64 or ARM64EC) and any EC object (ARM64EC or
+      // AMD64). This may be a single ARM64EC object, but may also be separate
+      // ARM64 and AMD64 objects.
+      bool HaveArm64 = false, HaveEC = false;
+      for (std::unique_ptr<SymbolicFile> &SymFile : SymFiles) {
+        if (!SymFile)
+          continue;
+        if (!HaveArm64)
+          HaveArm64 = isAnyArm64COFF(*SymFile);
+        if (!HaveEC)
+          HaveEC = isECObject(*SymFile);
+        if (HaveArm64 && HaveEC) {
+          SymMap->UseECMap = true;
+          break;
+        }
+      }
+    }
+  }
+
   // The big archive format needs to know the offset of the previous member
   // header.
   uint64_t PrevOffset = 0;
@@ -953,11 +997,10 @@ Expected<std::string> computeArchiveRelativePath(StringRef From, StringRef To) {
   return std::string(Relative);
 }
 
-static Error writeArchiveToStream(raw_ostream &Out,
-                                  ArrayRef<NewArchiveMember> NewMembers,
-                                  SymtabWritingMode WriteSymtab,
-                                  object::Archive::Kind Kind,
-                                  bool Deterministic, bool Thin, bool IsEC) {
+static Error
+writeArchiveToStream(raw_ostream &Out, ArrayRef<NewArchiveMember> NewMembers,
+                     SymtabWritingMode WriteSymtab, object::Archive::Kind Kind,
+                     bool Deterministic, bool Thin, std::optional<bool> IsEC) {
   assert((!Thin || !isBSDLike(Kind)) && "Only the gnu format has a thin mode");
 
   SmallString<0> SymNamesBuf;
@@ -977,10 +1020,9 @@ static Error writeArchiveToStream(raw_ostream &Out,
   // reference to it, thus SymbolicFile should be destroyed first.
   LLVMContext Context;
 
-  SymMap.UseECMap = IsEC;
   Expected<std::vector<MemberData>> DataOrErr = computeMemberData(
       StringTable, SymNames, Kind, Thin, Deterministic, WriteSymtab,
-      isCOFFArchive(Kind) ? &SymMap : nullptr, Context, NewMembers);
+      isCOFFArchive(Kind) ? &SymMap : nullptr, Context, NewMembers, IsEC);
   if (Error E = DataOrErr.takeError())
     return E;
   std::vector<MemberData> &Data = *DataOrErr;
@@ -1226,7 +1268,8 @@ static Error writeArchiveToStream(raw_ostream &Out,
 Error writeArchive(StringRef ArcName, ArrayRef<NewArchiveMember> NewMembers,
                    SymtabWritingMode WriteSymtab, object::Archive::Kind Kind,
                    bool Deterministic, bool Thin,
-                   std::unique_ptr<MemoryBuffer> OldArchiveBuf, bool IsEC) {
+                   std::unique_ptr<MemoryBuffer> OldArchiveBuf,
+                   std::optional<bool> IsEC) {
   Expected<sys::fs::TempFile> Temp =
       sys::fs::TempFile::create(ArcName + ".temp-archive-%%%%%%%.a");
   if (!Temp)
@@ -1263,7 +1306,7 @@ writeArchiveToBuffer(ArrayRef<NewArchiveMember> NewMembers,
   raw_svector_ostream ArchiveStream(ArchiveBufferVector);
 
   if (Error E = writeArchiveToStream(ArchiveStream, NewMembers, WriteSymtab,
-                                     Kind, Deterministic, Thin, false))
+                                     Kind, Deterministic, Thin, std::nullopt))
     return std::move(E);
 
   return std::make_unique<SmallVectorMemoryBuffer>(
diff --git a/llvm/test/tools/llvm-ar/ecsymbols.ll b/llvm/test/tools/llvm-ar/ecsymbols.ll
@@ -0,0 +1,19 @@
+;; Test that ECSYMBOLS section is created when ARM64EC bitcode is used.
+
+; RUN: llvm-as %s -o %t.bc
+; RUN: rm -f %t.a
+; RUN: llvm-ar cr %t.a %t.bc
+; RUN: llvm-nm --print-armap %t.a | FileCheck %s
+
+; CHECK-NOT: Archive map
+; CHECK: Archive EC map
+; CHECK-NEXT: a in ecsymbols.ll.tmp.bc
+; CHECK-NEXT: b in ecsymbols.ll.tmp.bc
+; CHECK-NEXT: c in ecsymbols.ll.tmp.bc
+; CHECK-EMPTY:
+
+target triple = "arm64ec-unknown-windows-msvc"
+
+define void @b() { ret void }
+define void @c() { ret void }
+define void @a() { ret void }
diff --git a/llvm/test/tools/llvm-ar/ecsymbols.yaml b/llvm/test/tools/llvm-ar/ecsymbols.yaml
@@ -0,0 +1,84 @@
+## Test that ECSYMBOLS section is created when ARM64EC is used.
+
+# RUN: yaml2obj %s -o %t.arm64ec.obj -DMACHINE=IMAGE_FILE_MACHINE_ARM64EC
+# RUN: yaml2obj %s -o %t.arm64.obj -DMACHINE=IMAGE_FILE_MACHINE_ARM64
+# RUN: yaml2obj %s -o %t.amd64.obj -DMACHINE=IMAGE_FILE_MACHINE_AMD64
+
+## Create ARM64EC archive.
+# RUN: rm -f %t*.a
+# RUN: llvm-ar cr %t1.a %t.arm64ec.obj
+# RUN: llvm-nm --print-armap %t1.a | FileCheck --check-prefixes=NOMAP,ECMAP %s
+
+## Add ARM64 object to the archive.
+# RUN: llvm-ar r %t1.a %t.arm64.obj
+# RUN: llvm-nm --print-armap %t1.a | FileCheck --check-prefixes=MAP,ECMAP %s
+
+## Create ARM64 archive.
+# RUN: llvm-ar cr %t2.a %t.arm64.obj
+# RUN: llvm-nm --print-armap %t2.a | FileCheck --check-prefixes=MAP,NOECMAP %s
+
+## Add ARM64EC object to the archive.
+# RUN: llvm-ar r %t2.a %t.arm64ec.obj
+# RUN: llvm-nm --print-armap %t2.a | FileCheck --check-prefixes=MAP,ECMAP %s
+
+## Create mixed archive with ARM64 and ARM64EC members.
+# RUN: llvm-ar cr %t3.a %t.arm64ec.obj %t.arm64.obj
+# RUN: llvm-nm --print-armap %t3.a | FileCheck --check-prefixes=MAP,ECMAP %s
+
+## Create mixed archive with ARM64 and AMD64 members.
+# RUN: llvm-ar cr %t4.a %t.amd64.obj %t.arm64.obj
+# RUN: llvm-nm --print-armap %t4.a | FileCheck --check-prefixes=MAP,AMDECMAP %s
+
+## Create an archive with no symbol table.
+# RUN: llvm-ar crS %t5.a %t.amd64.obj %t.arm64.obj
+# RUN: llvm-nm --print-armap %t5.a | FileCheck --check-prefixes=NOMAP,NOECMAP %s
+
+# MAP: Archive map
+# MAP-NEXT: a in ecsymbols.yaml.tmp.arm64.obj
+# MAP-NEXT: b in ecsymbols.yaml.tmp.arm64.obj
+# MAP-NEXT: c in ecsymbols.yaml.tmp.arm64.obj
+# MAP-EMPTY:
+# NOMAP-NOT: Archive map
+
+# ECMAP: Archive EC map
+# ECMAP-NEXT: a in ecsymbols.yaml.tmp.arm64ec.obj
+# ECMAP-NEXT: b in ecsymbols.yaml.tmp.arm64ec.obj
+# ECMAP-NEXT: c in ecsymbols.yaml.tmp.arm64ec.obj
+# ECMAP-EMPTY:
+# NOECMAP-NOT: Archive EC map
+
+# AMDECMAP: Archive EC map
+# AMDECMAP-NEXT: a in ecsymbols.yaml.tmp.amd64.obj
+# AMDECMAP-NEXT: b in ecsymbols.yaml.tmp.amd64.obj
+# AMDECMAP-NEXT: c in ecsymbols.yaml.tmp.amd64.obj
+# AMDECMAP-EMPTY:
+
+--- !COFF
+header:
+  Machine:         [[MACHINE]]
+  Characteristics: [  ]
+sections:
+  - Name:            .text
+    Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ]
+    Alignment:       4
+    SectionData:     ''
+symbols:
+  - Name:            b
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_FUNCTION
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+  - Name:            c
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_FUNCTION
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+  - Name:            a
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_FUNCTION
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...
diff --git a/llvm/test/tools/llvm-lib/empty.test b/llvm/test/tools/llvm-lib/empty.test
@@ -0,0 +1,27 @@
+Create import libraries with empty exports and make sure that archive symbols
+are properly populated.
+
+RUN: split-file %s %t.dir && cd %t.dir
+
+RUN: llvm-lib -machine:arm64 -out:arm64.lib -def:test.def
+RUN: llvm-nm --print-armap arm64.lib | FileCheck --check-prefixes=CHECK,NOECMAP %s
+
+RUN: llvm-lib -machine:arm64ec -out:arm64ec.lib -def:test.def
+RUN: llvm-nm --print-armap arm64ec.lib | FileCheck --check-prefixes=CHECK,ECMAP %s
+
+CHECK:      Archive map
+CHECK-NEXT: __IMPORT_DESCRIPTOR_test in test.dll
+CHECK-NEXT: __NULL_IMPORT_DESCRIPTOR in test.dll
+CHECK-NEXT: test_NULL_THUNK_DATA in test.dll
+CHECK-EMPTY:
+
+ECMAP:      Archive EC map
+ECMAP-NEXT: __IMPORT_DESCRIPTOR_test in test.dll
+ECMAP-NEXT: __NULL_IMPORT_DESCRIPTOR in test.dll
+ECMAP-NEXT: test_NULL_THUNK_DATA in test.dll
+
+NOECMAP-NOT: Archive EC map
+
+#--- test.def
+LIBRARY test.dll
+EXPORTS
