[llvm-lib] Add support for ARM64EC libraries.

ARM64EC allows having both pure ARM64 objects and ARM64EC in the
same archive. This allows using single static library for linking
pure ARM64, pure ARM64EC or mixed modules (what MS calls ARM64X:
a single module that may be used in both modes). To achieve that,
such static libraries need two separated symbol maps. The usual map
contains only pure ARM64 symbols, while a new /<ECSYMBOLS>/ section
contains EC symbols. EC symbols map has very similar format to the
usual map, except it doesn't contain object offsets and uses offsets
from regular map instead. This is true even for pure ARM64EC static
library: it will simply have 0 symbols in the symbol map.

Reviewed By: efriedma

Differential Revision: https://reviews.llvm.org/D143541

1 files changed, 73 insertions, 10 deletions

diff --git a/llvm/lib/Object/ArchiveWriter.cpp b/llvm/lib/Object/ArchiveWriter.cpp
index 2d0f92e..e9b8811 100644
--- a/llvm/lib/Object/ArchiveWriter.cpp
+++ b/llvm/lib/Object/ArchiveWriter.cpp
@@ -45,7 +45,9 @@
 using namespace llvm;
 
 struct SymMap {
+  bool UseECMap;
   std::map<std::string, uint16_t> Map;
+  std::map<std::string, uint16_t> ECMap;
 };
 
 NewArchiveMember::NewArchiveMember(MemoryBufferRef BufRef)
@@ -414,6 +416,20 @@ static uint64_t computeSymbolMapSize(uint64_t NumObj, SymMap &SymMap,
   return Size;
 }
 
+static uint64_t computeECSymbolsSize(SymMap &SymMap,
+                                     uint32_t *Padding = nullptr) {
+  uint64_t Size = sizeof(uint32_t); // Number of symbols
+
+  for (auto S : SymMap.ECMap)
+    Size += sizeof(uint16_t) + S.first.length() + 1;
+
+  uint32_t Pad = offsetToAlignment(Size, Align(2));
+  Size += Pad;
+  if (Padding)
+    *Padding = Pad;
+  return Size;
+}
+
 static void writeSymbolTableHeader(raw_ostream &Out, object::Archive::Kind Kind,
                                    bool Deterministic, uint64_t Size,
                                    uint64_t PrevMemberOffset = 0) {
@@ -446,8 +462,11 @@ static uint64_t computeHeadersSize(object::Archive::Kind Kind,
   uint32_t HeaderSize = computeSymbolTableHeaderSize();
   uint64_t Size = strlen("!<arch>\n") + HeaderSize + SymtabSize;
 
-  if (SymMap)
+  if (SymMap) {
     Size += HeaderSize + computeSymbolMapSize(NumMembers, *SymMap);
+    if (SymMap->ECMap.size())
+      Size += HeaderSize + computeECSymbolsSize(*SymMap);
+  }
 
   return Size + StringMemberSize;
 }
@@ -521,6 +540,41 @@ static void writeSymbolMap(raw_ostream &Out, object::Archive::Kind Kind,
     Out.write(uint8_t(0));
 }
 
+static void writeECSymbols(raw_ostream &Out, object::Archive::Kind Kind,
+                           bool Deterministic, ArrayRef<MemberData> Members,
+                           SymMap &SymMap) {
+  uint32_t Pad;
+  uint64_t Size = computeECSymbolsSize(SymMap, &Pad);
+  printGNUSmallMemberHeader(Out, "/<ECSYMBOLS>", now(Deterministic), 0, 0, 0,
+                            Size);
+
+  printLE<uint32_t>(Out, SymMap.ECMap.size());
+
+  for (auto S : SymMap.ECMap)
+    printLE(Out, S.second);
+  for (auto S : SymMap.ECMap)
+    Out << S.first << '\0';
+  while (Pad--)
+    Out.write(uint8_t(0));
+}
+
+static bool isECObject(object::SymbolicFile &Obj) {
+  if (Obj.isCOFF())
+    return cast<llvm::object::COFFObjectFile>(&Obj)->getMachine() !=
+           COFF::IMAGE_FILE_MACHINE_ARM64;
+
+  if (Obj.isIR()) {
+    Expected<std::string> TripleStr =
+        getBitcodeTargetTriple(Obj.getMemoryBufferRef());
+    if (!TripleStr)
+      return false;
+    Triple T(*TripleStr);
+    return T.isWindowsArm64EC() || T.getArch() == Triple::x86_64;
+  }
+
+  return false;
+}
+
 static Expected<std::vector<unsigned>>
 getSymbols(MemoryBufferRef Buf, uint16_t Index, raw_ostream &SymNames,
            SymMap *SymMap, bool &HasObject) {
@@ -548,20 +602,25 @@ getSymbols(MemoryBufferRef Buf, uint16_t Index, raw_ostream &SymNames,
     Obj = std::move(*ObjOrErr);
   }
 
+  std::map<std::string, uint16_t> *Map = nullptr;
+  if (SymMap)
+    Map = SymMap->UseECMap && isECObject(*Obj) ? &SymMap->ECMap : &SymMap->Map;
   HasObject = true;
   for (const object::BasicSymbolRef &S : Obj->symbols()) {
     if (!isArchiveSymbol(S))
       continue;
-    if (SymMap) {
+    if (Map) {
       std::string Name;
       raw_string_ostream NameStream(Name);
       if (Error E = S.printName(NameStream))
         return std::move(E);
-      if (SymMap->Map.find(Name) != SymMap->Map.end())
+      if (Map->find(Name) != Map->end())
         continue; // ignore duplicated symbol
-      SymMap->Map[Name] = Index;
-      Ret.push_back(SymNames.tell());
-      SymNames << Name << '\0';
+      (*Map)[Name] = Index;
+      if (Map == &SymMap->Map) {
+        Ret.push_back(SymNames.tell());
+        SymNames << Name << '\0';
+      }
     } else {
       Ret.push_back(SymNames.tell());
       if (Error E = S.printName(SymNames))
@@ -755,7 +814,7 @@ Expected<std::string> computeArchiveRelativePath(StringRef From, StringRef To) {
 static Error writeArchiveToStream(raw_ostream &Out,
                                   ArrayRef<NewArchiveMember> NewMembers,
                                   bool WriteSymtab, object::Archive::Kind Kind,
-                                  bool Deterministic, bool Thin) {
+                                  bool Deterministic, bool Thin, bool IsEC) {
   assert((!Thin || !isBSDLike(Kind)) && "Only the gnu format has a thin mode");
 
   SmallString<0> SymNamesBuf;
@@ -769,6 +828,7 @@ static Error writeArchiveToStream(raw_ostream &Out,
   if (isCOFFArchive(Kind) && NewMembers.size() > 0xfffe)
     Kind = object::Archive::K_GNU;
 
+  SymMap.UseECMap = IsEC;
   Expected<std::vector<MemberData>> DataOrErr = computeMemberData(
       StringTable, SymNames, Kind, Thin, Deterministic, WriteSymtab,
       isCOFFArchive(Kind) ? &SymMap : nullptr, NewMembers);
@@ -855,6 +915,9 @@ static Error writeArchiveToStream(raw_ostream &Out,
       Out << StringTableMember.Header << StringTableMember.Data
           << StringTableMember.Padding;
 
+    if (WriteSymtab && SymMap.ECMap.size())
+      writeECSymbols(Out, Kind, Deterministic, Data, SymMap);
+
     for (const MemberData &M : Data)
       Out << M.Header << M.Data << M.Padding;
   } else {
@@ -944,7 +1007,7 @@ static Error writeArchiveToStream(raw_ostream &Out,
 Error writeArchive(StringRef ArcName, ArrayRef<NewArchiveMember> NewMembers,
                    bool WriteSymtab, object::Archive::Kind Kind,
                    bool Deterministic, bool Thin,
-                   std::unique_ptr<MemoryBuffer> OldArchiveBuf) {
+                   std::unique_ptr<MemoryBuffer> OldArchiveBuf, bool IsEC) {
   Expected<sys::fs::TempFile> Temp =
       sys::fs::TempFile::create(ArcName + ".temp-archive-%%%%%%%.a");
   if (!Temp)
@@ -952,7 +1015,7 @@ Error writeArchive(StringRef ArcName, ArrayRef<NewArchiveMember> NewMembers,
   raw_fd_ostream Out(Temp->FD, false);
 
   if (Error E = writeArchiveToStream(Out, NewMembers, WriteSymtab, Kind,
-                                     Deterministic, Thin)) {
+                                     Deterministic, Thin, IsEC)) {
     if (Error DiscardError = Temp->discard())
       return joinErrors(std::move(E), std::move(DiscardError));
     return E;
@@ -981,7 +1044,7 @@ writeArchiveToBuffer(ArrayRef<NewArchiveMember> NewMembers, bool WriteSymtab,
   raw_svector_ostream ArchiveStream(ArchiveBufferVector);
 
   if (Error E = writeArchiveToStream(ArchiveStream, NewMembers, WriteSymtab,
-                                     Kind, Deterministic, Thin))
+                                     Kind, Deterministic, Thin, false))
     return std::move(E);
 
   return std::make_unique<SmallVectorMemoryBuffer>(