[llvm-ar][Object][COFF] Compute UseECMap in computeMemberData.

cjacek · cjacek · commit 6f151fb1e8d6 · 2024-03-14T15:37:12.000+01:00
Compute it instead of requiring caller to pass it to, which is inconvenient for llvm-ar.
diff --git a/llvm/include/llvm/Object/ArchiveWriter.h b/llvm/include/llvm/Object/ArchiveWriter.h
@@ -51,8 +51,7 @@ enum class SymtabWritingMode {
 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::unique_ptr<MemoryBuffer> OldArchiveBuf = nullptr);
 
 // 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 useECMap(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.getArch() == Triple::aarch64;
+  }
+
+  return false;
+}
+
 bool isImportDescriptor(StringRef Name) {
   return Name.starts_with(ImportDescriptorPrefix) ||
          Name == StringRef{NullImportDescriptorSymbolName} ||
@@ -803,6 +822,11 @@ computeMemberData(raw_ostream &StringTable, raw_ostream &SymNames,
           getSymbolicFile(M.Buf->getMemBufferRef(), Context);
       if (!SymFileOrErr)
         return createFileError(M.MemberName, SymFileOrErr.takeError());
+
+      // Use EC map if any member is ARM64 COFF file.
+      if (Kind == Archive::K_COFF && !SymMap->UseECMap && *SymFileOrErr)
+        SymMap->UseECMap = useECMap(**SymFileOrErr);
+
       SymFiles.push_back(std::move(*SymFileOrErr));
     }
   }
@@ -957,7 +981,7 @@ static Error writeArchiveToStream(raw_ostream &Out,
                                   ArrayRef<NewArchiveMember> NewMembers,
                                   SymtabWritingMode WriteSymtab,
                                   object::Archive::Kind Kind,
-                                  bool Deterministic, bool Thin, bool IsEC) {
+                                  bool Deterministic, bool Thin) {
   assert((!Thin || !isBSDLike(Kind)) && "Only the gnu format has a thin mode");
 
   SmallString<0> SymNamesBuf;
@@ -977,7 +1001,6 @@ 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);
@@ -1226,15 +1249,15 @@ 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) {
   Expected<sys::fs::TempFile> Temp =
       sys::fs::TempFile::create(ArcName + ".temp-archive-%%%%%%%.a");
   if (!Temp)
     return Temp.takeError();
   raw_fd_ostream Out(Temp->FD, false);
 
   if (Error E = writeArchiveToStream(Out, NewMembers, WriteSymtab, Kind,
-                                     Deterministic, Thin, IsEC)) {
+                                     Deterministic, Thin)) {
     if (Error DiscardError = Temp->discard())
       return joinErrors(std::move(E), std::move(DiscardError));
     return E;
@@ -1263,7 +1286,7 @@ writeArchiveToBuffer(ArrayRef<NewArchiveMember> NewMembers,
   raw_svector_ostream ArchiveStream(ArchiveBufferVector);
 
   if (Error E = writeArchiveToStream(ArchiveStream, NewMembers, WriteSymtab,
-                                     Kind, Deterministic, Thin, false))
+                                     Kind, Deterministic, Thin))
     return std::move(E);
 
   return std::make_unique<SmallVectorMemoryBuffer>(
diff --git a/llvm/lib/Object/COFFImportFile.cpp b/llvm/lib/Object/COFFImportFile.cpp
@@ -711,8 +711,7 @@ Error writeImportLibrary(StringRef ImportName, StringRef Path,
 
   return writeArchive(Path, Members, SymtabWritingMode::NormalSymtab,
                       object::Archive::K_COFF,
-                      /*Deterministic*/ true, /*Thin*/ false,
-                      /*OldArchiveBuf*/ nullptr, isArm64EC(Machine));
+                      /*Deterministic*/ true, /*Thin*/ false);
 }
 
 } // namespace object
diff --git a/llvm/lib/ToolDrivers/llvm-lib/LibDriver.cpp b/llvm/lib/ToolDrivers/llvm-lib/LibDriver.cpp
@@ -507,10 +507,10 @@ int llvm::libDriverMain(ArrayRef<const char *> ArgsArr) {
   std::reverse(Members.begin(), Members.end());
 
   bool Thin = Args.hasArg(OPT_llvmlibthin);
-  if (Error E = writeArchive(
-          OutputPath, Members, SymtabWritingMode::NormalSymtab,
-          Thin ? object::Archive::K_GNU : object::Archive::K_COFF,
-          /*Deterministic=*/true, Thin, nullptr, COFF::isArm64EC(LibMachine))) {
+  if (Error E =
+          writeArchive(OutputPath, Members, SymtabWritingMode::NormalSymtab,
+                       Thin ? object::Archive::K_GNU : object::Archive::K_COFF,
+                       /*Deterministic=*/true, Thin)) {
     handleAllErrors(std::move(E), [&](const ErrorInfoBase &EI) {
       llvm::errs() << OutputPath << ": " << EI.message() << "\n";
     });
diff --git a/llvm/test/tools/llvm-ar/ecsymbols.yaml b/llvm/test/tools/llvm-ar/ecsymbols.yaml
@@ -0,0 +1,83 @@
+## 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 crs %t.a %t.arm64ec.obj
+# RUN: llvm-nm --print-armap %t.a | FileCheck --check-prefixes=NOMAP,ECMAP %s
+
+## Add ARM64 object to the archive
+# RUN: llvm-ar rs %t.a %t.arm64.obj
+# RUN: llvm-nm --print-armap %t.a | FileCheck --check-prefixes=MAP,ECMAP %s
+
+## Create ARM64 archive
+# RUN: rm -f %t.a
+# RUN: llvm-ar crs %t.a %t.arm64.obj
+# RUN: llvm-nm --print-armap %t.a | FileCheck --check-prefixes=MAP,NOECMAP %s
+
+# Add ARM64EC object to the archive
+# RUN: llvm-ar rs %t.a %t.arm64ec.obj
+# RUN: llvm-nm --print-armap %t.a | FileCheck --check-prefixes=MAP,ECMAP %s
+
+## Create mixed archive
+# RUN: rm -f %t.a
+# RUN: llvm-ar crs %t.a %t.arm64ec.obj %t.arm64.obj
+# RUN: llvm-nm --print-armap %t.a | FileCheck --check-prefixes=MAP,ECMAP %s
+
+## Create mixed archive
+# RUN: rm -f %t.a
+# RUN: llvm-ar crs %t.a %t.amd64.obj %t.arm64.obj
+# RUN: llvm-nm --print-armap %t.a | FileCheck --check-prefixes=MAP,AMDECMAP %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
+# ECMAP-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
+...

Original file line number	Diff line number	Diff line change
`@@ -711,8 +711,7 @@ Error writeImportLibrary(StringRef ImportName, StringRef Path,`
`711`	`711`
`712`	`712`	`return writeArchive(Path, Members, SymtabWritingMode::NormalSymtab,`
`713`	`713`	`object::Archive::K_COFF,`
`714`		`- /Deterministic/ true, /Thin/ false,`
`715`		`- /OldArchiveBuf/ nullptr, isArm64EC(Machine));`
	`714`	`+ /Deterministic/ true, /Thin/ false);`
`716`	`715`	`}`
`717`	`716`
`718`	`717`	`} // namespace object`