[ORC] Refactor visit-members in StaticLibraryDefinitionGenerator. (#141546)

This refactor was motivated by two bugs identified in out-of-tree
builds:

1. Some implementations of the VisitMembersFunction type (often used to	
implement special loading semantics, e.g. -all_load or -ObjC) were assuming
that buffers for archive members were null-terminated, which they are not in
general. This was triggering occasional assertions.

2. Archives may include multiple members with the same file name, e.g.
when constructed by appending files with the same name:
  % llvm-ar crs libfoo.a foo.o
  % llvm-ar q libfoo.a foo.o
  % llvm-ar t libfoo.a foo.o
  foo.o

   While confusing, these members may be safe to link (provided that they're
   individually valid and don't define duplicate symbols). In ORC however, the
   archive member name may be used to construct an ORC initializer symbol,
   which must also be unique. In that case the duplicate member names lead to a
   duplicate definition error even if the members define unrelated symbols.

In addition to these bugs, StaticLibraryDefinitionGenerator had grown a
collection of all member buffers (ObjectFilesMap), a BumpPtrAllocator
that was redundantly storing synthesized archive member names (these are
copied into the MemoryBuffers created for each Object, but were never
freed in the allocator), and a set of COFF-specific import files.

To fix the bugs above and simplify StaticLibraryDefinitionGenerator this
patch makes the following changes:

1. StaticLibraryDefinitionGenerator::VisitMembersFunction is generalized
   to take a reference to the containing archive, and the index of the
   member within the archive. It now returns an Expected<bool> indicating
   whether the member visited should be treated as loadable, not loadable,
   or as invalidating the entire archive.
2. A static StaticLibraryDefinitionGenerator::createMemberBuffer method
   is added which creates MemoryBuffers with unique names of the form
   `<archive-name>[<index>](<member-name>)`. This defers construction of
   member names until they're loaded, allowing the BumpPtrAllocator (with
   its redundant name storage) to be removed.
3. The ObjectFilesMap (symbol name -> memory-buffer-ref) is replaced
   with a SymbolToMemberIndexMap (symbol name -> index) which should be
   smaller and faster to construct.
4. The 'loadability' result from VisitMemberFunctions is now taken into
   consideration when building the SymbolToMemberIndexMap so that members
   that have already been loaded / filtered out can be skipped, and do not
   take up any ongoing space.
5. The COFF ImportedDynamicLibraries member is moved out into the
   COFFImportFileScanner utility, which can be used as a
   VisitMemberFunction.

This fixes the bugs described above; and should lower memory consumption
slightly, especially for archives with many files and / or symbol where
most files are eventually loaded.

llvm/include/llvm/ExecutionEngine/Orc/COFF.h

@@ -0,0 +1,44 @@
+//===-------------- COFF.h - COFF format utilities --------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Contains utilities for load COFF relocatable object files.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_COFF_H
+#define LLVM_EXECUTIONENGINE_ORC_COFF_H
+
+#include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
+
+#include <set>
+#include <string>
+
+namespace llvm {
+
+namespace object {
+class Archive;
+} // namespace object
+
+namespace orc {
+
+class COFFImportFileScanner {
+public:
+  COFFImportFileScanner(std::set<std::string> &ImportedDynamicLibraries)
+      : ImportedDynamicLibraries(ImportedDynamicLibraries) {}
+  Expected<bool> operator()(object::Archive &A, MemoryBufferRef MemberBuf,
+                            size_t Index) const;
+
+private:
+  std::set<std::string> &ImportedDynamicLibraries;
+};
+
+} // namespace orc
+} // namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_MACHO_H

llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h

@@ -139,6 +139,7 @@ private:
   COFFPlatform(
       ObjectLinkingLayer &ObjLinkingLayer, JITDylib &PlatformJD,
       std::unique_ptr<StaticLibraryDefinitionGenerator> OrcRuntimeGenerator,
+      std::set<std::string> DylibsToPreload,
       std::unique_ptr<MemoryBuffer> OrcRuntimeArchiveBuffer,
       std::unique_ptr<object::Archive> OrcRuntimeArchive,
       LoadDynamicLibrary LoadDynLibrary, bool StaticVCRuntime,

llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h

@@ -273,9 +273,23 @@ public:
       unique_function<Expected<MaterializationUnit::Interface>(
           ExecutionSession &ES, MemoryBufferRef ObjBuffer)>;
 
-  /// Callback for visiting archive members at construction time.
-  /// Con be used to pre-load archive members.
-  using VisitMembersFunction = unique_function<Error(MemoryBufferRef)>;
+  /// Callback for visiting archive members at construction time. Can be used
+  /// to pre-load members.
+  ///
+  /// Callbacks are provided with a reference to the underlying archive, a
+  /// MemoryBufferRef covering the bytes for the given member, and the index of
+  /// the given member.
+  ///
+  /// Implementations should return true if the given member file should be
+  /// loadable via the generator, false if it should not, and an Error if the
+  /// member is malformed in a way that renders the archive itself invalid.
+  ///
+  /// Note: Linkers typically ignore invalid files within archives, so it's
+  ///       expected that implementations will usually return `false` (i.e.
+  ///       not-loadable) for malformed buffers, and will only return an
+  ///       Error in exceptional circumstances.
+  using VisitMembersFunction = unique_function<Expected<bool>(
+      object::Archive &, MemoryBufferRef, size_t)>;
 
   /// A VisitMembersFunction that unconditionally loads all object files from
   /// the archive.
@@ -283,6 +297,9 @@ public:
   static VisitMembersFunction loadAllObjectFileMembers(ObjectLayer &L,
                                                        JITDylib &JD);
 
+  static std::unique_ptr<MemoryBuffer>
+  createMemberBuffer(object::Archive &A, MemoryBufferRef BufRef, size_t Index);
+
   /// Try to create a StaticLibraryDefinitionGenerator from the given path.
   ///
   /// This call will succeed if the file at the given path is a static library
@@ -313,32 +330,22 @@ public:
          VisitMembersFunction VisitMembers = VisitMembersFunction(),
          GetObjectFileInterface GetObjFileInterface = GetObjectFileInterface());
 
-  /// Returns a list of filenames of dynamic libraries that this archive has
-  /// imported. This class does not load these libraries by itself. User is
-  /// responsible for making sure these libraries are available to the JITDylib.
-  const std::set<std::string> &getImportedDynamicLibraries() const {
-    return ImportedDynamicLibraries;
-  }
-
   Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                       JITDylibLookupFlags JDLookupFlags,
                       const SymbolLookupSet &Symbols) override;
 
 private:
-  StaticLibraryDefinitionGenerator(ObjectLayer &L,
-                                   std::unique_ptr<MemoryBuffer> ArchiveBuffer,
-                                   std::unique_ptr<object::Archive> Archive,
-                                   GetObjectFileInterface GetObjFileInterface,
-                                   Error &Err);
-  Error buildObjectFilesMap();
+  StaticLibraryDefinitionGenerator(
+      ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer,
+      std::unique_ptr<object::Archive> Archive,
+      GetObjectFileInterface GetObjFileInterface,
+      DenseMap<SymbolStringPtr, size_t> SymbolToMemberIndexMap);
 
   ObjectLayer &L;
   GetObjectFileInterface GetObjFileInterface;
-  std::set<std::string> ImportedDynamicLibraries;
   std::unique_ptr<MemoryBuffer> ArchiveBuffer;
   std::unique_ptr<object::Archive> Archive;
-  DenseMap<SymbolStringPtr, MemoryBufferRef> ObjectFilesMap;
-  BumpPtrAllocator ObjFileNameStorage;
+  DenseMap<SymbolStringPtr, size_t> SymbolToMemberIndexMap;
 };
 
 /// A utility class to create COFF dllimport GOT symbols (__imp_*) and PLT

llvm/include/llvm/ExecutionEngine/Orc/MachO.h

@@ -14,6 +14,7 @@
 #define LLVM_EXECUTIONENGINE_ORC_MACHO_H
 
 #include "llvm/ExecutionEngine/Orc/LoadLinkableFile.h"
+#include "llvm/Object/Archive.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/TargetParser/Triple.h"
@@ -22,6 +23,7 @@ namespace llvm {
 
 namespace object {
 
+class Archive;
 class MachOUniversalBinary;
 
 } // namespace object
@@ -81,7 +83,8 @@ public:
   ForceLoadMachOArchiveMembers(ObjectLayer &L, JITDylib &JD, bool ObjCOnly)
       : L(L), JD(JD), ObjCOnly(ObjCOnly) {}
 
-  Error operator()(MemoryBufferRef MemberBuf);
+  Expected<bool> operator()(object::Archive &A, MemoryBufferRef MemberBuf,
+                            size_t Index);
 
 private:
   ObjectLayer &L;

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

@@ -8,6 +8,7 @@ endif()
 
 add_llvm_component_library(LLVMOrcJIT
   AbsoluteSymbols.cpp
+  COFF.cpp
   COFFVCRuntimeSupport.cpp
   COFFPlatform.cpp
   CompileOnDemandLayer.cpp

llvm/lib/ExecutionEngine/Orc/COFF.cpp

@@ -0,0 +1,39 @@
+//===------------------ COFF.cpp - COFF format utilities ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/COFF.h"
+#include "llvm/Object/Binary.h"
+
+#define DEBUG_TYPE "orc"
+
+namespace llvm::orc {
+
+Expected<bool> COFFImportFileScanner::operator()(object::Archive &A,
+                                                 MemoryBufferRef MemberBuf,
+                                                 size_t Index) const {
+  // Try to build a binary for the member.
+  auto Bin = object::createBinary(MemberBuf);
+  if (!Bin) {
+    // If we can't then consume the error and return false (i.e. not loadable).
+    consumeError(Bin.takeError());
+    return false;
+  }
+
+  // If this is a COFF import file then handle it and return false (not
+  // loadable).
+  if ((*Bin)->isCOFFImportFile()) {
+    ImportedDynamicLibraries.insert((*Bin)->getFileName().str());
+    return false;
+  }
+
+  // Otherwise the member is loadable (at least as far as COFFImportFileScanner
+  // is concerned), so return true;
+  return true;
+}
+
+} // namespace llvm::orc

llvm/lib/ExecutionEngine/Orc/COFFPlatform.cpp

@@ -7,7 +7,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/COFFPlatform.h"
+
 #include "llvm/ExecutionEngine/Orc/AbsoluteSymbols.h"
+#include "llvm/ExecutionEngine/Orc/COFF.h"
 #include "llvm/ExecutionEngine/Orc/DebugUtils.h"
 #include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
 #include "llvm/ExecutionEngine/Orc/ObjectFileInterface.h"
@@ -170,8 +172,10 @@ COFFPlatform::Create(ObjectLinkingLayer &ObjLinkingLayer, JITDylib &PlatformJD,
   if (!GeneratorArchive)
     return GeneratorArchive.takeError();
 
+  std::set<std::string> DylibsToPreload;
   auto OrcRuntimeArchiveGenerator = StaticLibraryDefinitionGenerator::Create(
-      ObjLinkingLayer, nullptr, std::move(*GeneratorArchive));
+      ObjLinkingLayer, nullptr, std::move(*GeneratorArchive),
+      COFFImportFileScanner(DylibsToPreload));
   if (!OrcRuntimeArchiveGenerator)
     return OrcRuntimeArchiveGenerator.takeError();
 
@@ -207,8 +211,9 @@ COFFPlatform::Create(ObjectLinkingLayer &ObjLinkingLayer, JITDylib &PlatformJD,
   Error Err = Error::success();
   auto P = std::unique_ptr<COFFPlatform>(new COFFPlatform(
       ObjLinkingLayer, PlatformJD, std::move(*OrcRuntimeArchiveGenerator),
-      std::move(OrcRuntimeArchiveBuffer), std::move(RuntimeArchive),
-      std::move(LoadDynLibrary), StaticVCRuntime, VCRuntimePath, Err));
+      std::move(DylibsToPreload), std::move(OrcRuntimeArchiveBuffer),
+      std::move(RuntimeArchive), std::move(LoadDynLibrary), StaticVCRuntime,
+      VCRuntimePath, Err));
   if (Err)
     return std::move(Err);
   return std::move(P);
@@ -376,6 +381,7 @@ bool COFFPlatform::supportedTarget(const Triple &TT) {
 COFFPlatform::COFFPlatform(
     ObjectLinkingLayer &ObjLinkingLayer, JITDylib &PlatformJD,
     std::unique_ptr<StaticLibraryDefinitionGenerator> OrcRuntimeGenerator,
+    std::set<std::string> DylibsToPreload,
     std::unique_ptr<MemoryBuffer> OrcRuntimeArchiveBuffer,
     std::unique_ptr<object::Archive> OrcRuntimeArchive,
     LoadDynamicLibrary LoadDynLibrary, bool StaticVCRuntime,
@@ -401,10 +407,6 @@ COFFPlatform::COFFPlatform(
   }
   VCRuntimeBootstrap = std::move(*VCRT);
 
-  std::set<std::string> DylibsToPreload;
-  for (auto &Lib : OrcRuntimeGenerator->getImportedDynamicLibraries())
-    DylibsToPreload.insert(Lib);
-
   auto ImportedLibs =
       StaticVCRuntime ? VCRuntimeBootstrap->loadStaticVCRuntime(PlatformJD)
                       : VCRuntimeBootstrap->loadDynamicVCRuntime(PlatformJD);

llvm/lib/ExecutionEngine/Orc/COFFVCRuntimeSupport.cpp

@@ -8,6 +8,7 @@
 
 #include "llvm/ExecutionEngine/Orc/COFFVCRuntimeSupport.h"
 
+#include "llvm/ExecutionEngine/Orc/COFF.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
 #include "llvm/Support/VirtualFileSystem.h"
@@ -81,12 +82,14 @@ Error COFFVCRuntimeBootstrapper::loadVCRuntime(
   auto LoadLibrary = [&](SmallString<256> LibPath, StringRef LibName) -> Error {
     sys::path::append(LibPath, LibName);
 
-    auto G = StaticLibraryDefinitionGenerator::Load(ObjLinkingLayer,
-                                                    LibPath.c_str());
+    std::set<std::string> NewImportedLibraries;
+    auto G = StaticLibraryDefinitionGenerator::Load(
+        ObjLinkingLayer, LibPath.c_str(),
+        COFFImportFileScanner(NewImportedLibraries));
     if (!G)
       return G.takeError();
 
-    llvm::append_range(ImportedLibraries, (*G)->getImportedDynamicLibraries());
+    llvm::append_range(ImportedLibraries, NewImportedLibraries);
 
     JD.addGenerator(std::move(*G));
 

llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp

@@ -20,7 +20,6 @@
 #include "llvm/IR/Module.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Object/MachOUniversal.h"
-#include "llvm/Support/StringSaver.h"
 #include "llvm/Target/TargetMachine.h"
 #include <string>
 
@@ -275,14 +274,19 @@ Error DynamicLibrarySearchGenerator::tryToGenerate(
 StaticLibraryDefinitionGenerator::VisitMembersFunction
 StaticLibraryDefinitionGenerator::loadAllObjectFileMembers(ObjectLayer &L,
                                                            JITDylib &JD) {
-  return [&](MemoryBufferRef Buf) -> Error {
+  return [&](object::Archive &A, MemoryBufferRef Buf,
+             size_t Index) -> Expected<bool> {
     switch (identify_magic(Buf.getBuffer())) {
     case file_magic::elf_relocatable:
     case file_magic::macho_object:
     case file_magic::coff_object:
-      return L.add(JD, MemoryBuffer::getMemBuffer(Buf));
+      if (auto Err = L.add(JD, createMemberBuffer(A, Buf, Index)))
+        return std::move(Err);
+      // Since we've loaded it already, mark this as not loadable.
+      return false;
     default:
-      return Error::success();
+      // Non-object-file members are not loadable.
+      return false;
     }
   };
 }
@@ -307,13 +311,18 @@ StaticLibraryDefinitionGenerator::Create(
     std::unique_ptr<object::Archive> Archive, VisitMembersFunction VisitMembers,
     GetObjectFileInterface GetObjFileInterface) {
 
-  Error Err = Error::success();
+  DenseSet<uint64_t> Excluded;
 
   if (VisitMembers) {
+    size_t Index = 0;
+    Error Err = Error::success();
     for (auto Child : Archive->children(Err)) {
       if (auto ChildBuf = Child.getMemoryBufferRef()) {
-        if (auto Err2 = VisitMembers(*ChildBuf))
-          return std::move(Err2);
+        if (auto Loadable = VisitMembers(*Archive, *ChildBuf, Index++)) {
+          if (!*Loadable)
+            Excluded.insert(Child.getDataOffset());
+        } else
+          return Loadable.takeError();
       } else {
         // We silently allow non-object archive members. This matches the
         // behavior of ld.
@@ -324,15 +333,39 @@ StaticLibraryDefinitionGenerator::Create(
       return std::move(Err);
   }
 
-  std::unique_ptr<StaticLibraryDefinitionGenerator> ADG(
+  DenseMap<SymbolStringPtr, size_t> SymbolToMemberIndexMap;
+  {
+    DenseMap<uint64_t, size_t> OffsetToIndex;
+    size_t Index = 0;
+    Error Err = Error::success();
+    for (auto &Child : Archive->children(Err)) {
+      // For all members not excluded above, add them to the OffsetToIndex map.
+      if (!Excluded.count(Child.getDataOffset()))
+        OffsetToIndex[Child.getDataOffset()] = Index;
+      ++Index;
+    }
+    if (Err)
+      return Err;
+
+    auto &ES = L.getExecutionSession();
+    for (auto &Sym : Archive->symbols()) {
+      auto Member = Sym.getMember();
+      if (!Member)
+        return Member.takeError();
+      auto EntryItr = OffsetToIndex.find(Member->getDataOffset());
+
+      // Missing entry means this member should be ignored.
+      if (EntryItr == OffsetToIndex.end())
+        continue;
+
+      SymbolToMemberIndexMap[ES.intern(Sym.getName())] = EntryItr->second;
+    }
+  }
+
+  return std::unique_ptr<StaticLibraryDefinitionGenerator>(
       new StaticLibraryDefinitionGenerator(
           L, std::move(ArchiveBuffer), std::move(Archive),
-          std::move(GetObjFileInterface), Err));
-
-  if (Err)
-    return std::move(Err);
-
-  return std::move(ADG);
+          std::move(GetObjFileInterface), std::move(SymbolToMemberIndexMap)));
 }
 
 Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
@@ -392,86 +425,81 @@ Error StaticLibraryDefinitionGenerator::tryToGenerate(
   if (!Archive)
     return Error::success();
 
-  DenseSet<std::pair<StringRef, StringRef>> ChildBufferInfos;
+  DenseMap<size_t, MemoryBufferRef> ToLoad;
 
-  for (const auto &KV : Symbols) {
-    const auto &Name = KV.first;
-    auto It = ObjectFilesMap.find(Name);
-    if (It == ObjectFilesMap.end())
+  for (const auto &[Name, _] : Symbols) {
+    // Check whehter the archive contains this symbol.
+    auto It = SymbolToMemberIndexMap.find(Name);
+    if (It == SymbolToMemberIndexMap.end())
       continue;
-    auto ChildBuffer = It->second;
-    ChildBufferInfos.insert(
-        {ChildBuffer.getBuffer(), ChildBuffer.getBufferIdentifier()});
+    size_t Index = It->second;
+
+    // If we're already loading the member containing this symbol then we're
+    // done.
+    if (ToLoad.count(Index))
+      continue;
+
+    auto Member = Archive->findSym(*Name);
+    if (!Member)
+      return Member.takeError();
+    if (!*Member) // Skip "none" children.
+      continue;
+
+    auto MemberBuf = (*Member)->getMemoryBufferRef();
+    if (!MemberBuf)
+      return MemberBuf.takeError();
+
+    ToLoad[Index] = *MemberBuf;
   }
 
-  for (auto ChildBufferInfo : ChildBufferInfos) {
-    MemoryBufferRef ChildBufferRef(ChildBufferInfo.first,
-                                   ChildBufferInfo.second);
+  // Remove symbols to be loaded.
+  {
+    // FIXME: Enable DenseMap removal using NonOwningSymbolStringPtr?
+    std::vector<SymbolStringPtr> ToRemove;
+    for (auto &[Name, Index] : SymbolToMemberIndexMap)
+      if (ToLoad.count(Index))
+        ToRemove.push_back(Name);
+    for (auto &Name : ToRemove)
+      SymbolToMemberIndexMap.erase(Name);
+  }
 
-    auto I = GetObjFileInterface(L.getExecutionSession(), ChildBufferRef);
-    if (!I)
-      return I.takeError();
+  // Add loaded files to JITDylib.
+  for (auto &[Index, Buf] : ToLoad) {
+    auto MemberBuf = createMemberBuffer(*Archive, Buf, Index);
 
-    if (auto Err = L.add(JD, MemoryBuffer::getMemBuffer(ChildBufferRef, false),
-                         std::move(*I)))
+    auto Interface = GetObjFileInterface(L.getExecutionSession(),
+                                         MemberBuf->getMemBufferRef());
+    if (!Interface)
+      return Interface.takeError();
+
+    if (auto Err = L.add(JD, std::move(MemberBuf), std::move(*Interface)))
       return Err;
   }
 
   return Error::success();
 }
 
-Error StaticLibraryDefinitionGenerator::buildObjectFilesMap() {
-  DenseMap<uint64_t, MemoryBufferRef> MemoryBuffers;
-  DenseSet<uint64_t> Visited;
-  DenseSet<uint64_t> Excluded;
-  StringSaver FileNames(ObjFileNameStorage);
-  for (auto &S : Archive->symbols()) {
-    StringRef SymName = S.getName();
-    auto Member = S.getMember();
-    if (!Member)
-      return Member.takeError();
-    auto DataOffset = Member->getDataOffset();
-    if (!Visited.count(DataOffset)) {
-      Visited.insert(DataOffset);
-      auto Child = Member->getAsBinary();
-      if (!Child)
-        return Child.takeError();
-      if ((*Child)->isCOFFImportFile()) {
-        ImportedDynamicLibraries.insert((*Child)->getFileName().str());
-        Excluded.insert(DataOffset);
-        continue;
-      }
-
-      // Give members of the archive a name that contains the archive path so
-      // that they can be differentiated from a member with the same name in a
-      // different archive. This also ensure initializer symbols names will be
-      // unique within a JITDylib.
-      StringRef FullName = FileNames.save(Archive->getFileName() + "(" +
-                                          (*Child)->getFileName() + ")");
-      MemoryBufferRef MemBuffer((*Child)->getMemoryBufferRef().getBuffer(),
-                                FullName);
-
-      MemoryBuffers[DataOffset] = MemBuffer;
-    }
-    if (!Excluded.count(DataOffset))
-      ObjectFilesMap[L.getExecutionSession().intern(SymName)] =
-          MemoryBuffers[DataOffset];
-  }
-
-  return Error::success();
+std::unique_ptr<MemoryBuffer>
+StaticLibraryDefinitionGenerator::createMemberBuffer(object::Archive &A,
+                                                     MemoryBufferRef BufRef,
+                                                     size_t Index) {
+  return MemoryBuffer::getMemBuffer(BufRef.getBuffer(),
+                                    (A.getFileName() + "[" + Twine(Index) +
+                                     "](" + BufRef.getBufferIdentifier() + ")")
+                                        .str(),
+                                    false);
 }
 
 StaticLibraryDefinitionGenerator::StaticLibraryDefinitionGenerator(
     ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer,
     std::unique_ptr<object::Archive> Archive,
-    GetObjectFileInterface GetObjFileInterface, Error &Err)
+    GetObjectFileInterface GetObjFileInterface,
+    DenseMap<SymbolStringPtr, size_t> SymbolToMemberIndexMap)
     : L(L), GetObjFileInterface(std::move(GetObjFileInterface)),
-      ArchiveBuffer(std::move(ArchiveBuffer)), Archive(std::move(Archive)) {
-  ErrorAsOutParameter _(Err);
+      ArchiveBuffer(std::move(ArchiveBuffer)), Archive(std::move(Archive)),
+      SymbolToMemberIndexMap(std::move(SymbolToMemberIndexMap)) {
   if (!this->GetObjFileInterface)
     this->GetObjFileInterface = getObjectFileInterface;
-  if (!Err)
-    Err = buildObjectFilesMap();
 }
 
 std::unique_ptr<DLLImportDefinitionGenerator>

llvm/lib/ExecutionEngine/Orc/MachO.cpp

@@ -10,6 +10,7 @@
 
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/BinaryFormat/MachO.h"
+#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/Layer.h"
 #include "llvm/Object/MachOUniversal.h"
 #include "llvm/Support/FileSystem.h"
@@ -229,18 +230,29 @@ getMachOSliceRangeForTriple(MemoryBufferRef UBBuf, const Triple &TT) {
   return getMachOSliceRangeForTriple(**UB, TT);
 }
 
-Error ForceLoadMachOArchiveMembers::operator()(MemoryBufferRef MemberBuf) {
-  if (!ObjCOnly)
-    return L.add(JD, MemoryBuffer::getMemBuffer(MemberBuf));
+Expected<bool> ForceLoadMachOArchiveMembers::operator()(
+    object::Archive &A, MemoryBufferRef MemberBuf, size_t Index) {
+
+  auto LoadMember = [&]() {
+    return StaticLibraryDefinitionGenerator::createMemberBuffer(A, MemberBuf,
+                                                                Index);
+  };
+
+  if (!ObjCOnly) {
+    // If we're loading all files then just load the buffer immediately. Return
+    // false to indicate that there's no further loading to do here.
+    if (auto Err = L.add(JD, LoadMember()))
+      return Err;
+    return false;
+  }
 
   // We need to check whether this archive member contains any Objective-C
   // or Swift metadata.
-
   auto Obj = object::ObjectFile::createObjectFile(MemberBuf);
   if (!Obj) {
-    // We silently ignore invalid files.
+    // Invalid files are not loadable, but don't invalidate the archive.
     consumeError(Obj.takeError());
-    return Error::success();
+    return false;
   }
 
   if (auto *MachOObj = dyn_cast<object::MachOObjectFile>(&**Obj)) {
@@ -253,14 +265,19 @@ Error ForceLoadMachOArchiveMembers::operator()(MemoryBufferRef MemberBuf) {
             *SecName == "__objc_clsrolist" || *SecName == "__objc_catlist" ||
             *SecName == "__objc_catlist2" || *SecName == "__objc_nlcatlist" ||
             (SegName == "__TEXT" && (*SecName).starts_with("__swift") &&
-             *SecName != "__swift_modhash"))
-          return L.add(JD, MemoryBuffer::getMemBuffer(MemberBuf));
+             *SecName != "__swift_modhash")) {
+          if (auto Err = L.add(JD, LoadMember()))
+            return Err;
+          return false;
+        }
       } else
         return SecName.takeError();
     }
   }
 
-  return Error::success();
+  // This is an object file but we didn't load it, so return true to indicate
+  // that it's still loadable.
+  return true;
 }
 
 } // End namespace orc.

llvm/test/ExecutionEngine/JITLink/Generic/Inputs/bar-initializer.ll

@@ -0,0 +1,7 @@
+@initializer_bar_ran = global i32 0
+@llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @initializer_bar, ptr null }]
+
+define internal void @initializer_bar() {
+  store i32 1, ptr @initializer_bar_ran
+  ret void
+}

llvm/test/ExecutionEngine/JITLink/Generic/Inputs/bar-ret-0.ll

@@ -0,0 +1,3 @@
+define i32 @bar()  {
+  ret i32 0
+}

llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-initializer.ll

@@ -0,0 +1,7 @@
+@initializer_foo_ran = global i32 0
+@llvm.global_ctors = appending global [1 x { i32, ptr, ptr }] [{ i32, ptr, ptr } { i32 65535, ptr @initializer_foo, ptr null }]
+
+define internal void @initializer_foo() {
+  store i32 1, ptr @initializer_foo_ran
+  ret void
+}

llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-ret-0.ll

@@ -0,0 +1,3 @@
+define i32 @foo()  {
+  ret i32 0
+}

llvm/test/ExecutionEngine/JITLink/Generic/all-load-multifile-archive-with-duplicate-member-filenames.test

@@ -0,0 +1,32 @@
+# RUN: rm -rf %t && mkdir -p %t
+# RUN: llc -filetype=obj -o %t/foo.o %S/Inputs/foo-initializer.ll
+# RUN: llvm-ar crs %t/libFoo.a %t/foo.o
+# RUN: llc -filetype=obj -o %t/foo.o %S/Inputs/bar-initializer.ll
+# RUN: llvm-ar q %t/libFoo.a %t/foo.o
+# RUN: llc -filetype=obj -o %t/main.o %S/Inputs/main-ret-0.ll
+# RUN: llvm-jitlink -all_load -show-init-es -noexec %t/main.o -L%t -lFoo \
+# RUN:     | FileCheck %s
+#
+# Check that synthesized archive member names are unambiguous, even if an
+# archive contains multiple files with the same name.
+#
+# Background: Static achives may contain duplicate member filenames. E.g. we
+# can create an archive with initially contaning a single object file 'foo.o',
+# then append another copy of 'foo.o' to produce an archive containing two
+# consecutive copies:
+#
+#   % llvm-ar crs libfoo.a foo.o
+#   % llvm-ar q libfoo.a foo.o
+#   % llvm-ar t libfoo.a
+#   foo.o
+#   foo.o
+#
+# In this test we create two different 'foo.o' files, each containing a static
+# initializer. Since initializer names are based on the full archive member
+# names, failure to give the members unique names will result in a duplicate
+# definition error. The `-all_load` option is used to force all member files in
+# the archive to be loaded, despite the members not being directly referenced.
+#
+# CHECK-DAG: main{{.*}}main.o
+# CHECK-DAG: initializer_foo_ran{{.*}}libFoo.a[0](foo.o)
+# CHECK-DAG: initializer_bar_ran{{.*}}libFoo.a[1](foo.o)

llvm/test/ExecutionEngine/JITLink/Generic/all-load-multifile-archive.test

@@ -0,0 +1,14 @@
+# RUN: rm -rf %t && mkdir -p %t
+# RUN: llc -filetype=obj -o %t/foo.o %S/Inputs/foo-ret-0.ll
+# RUN: llc -filetype=obj -o %t/bar.o %S/Inputs/bar-ret-0.ll
+# RUN: llvm-ar crs %t/libFoo.a %t/foo.o %t/bar.o
+# RUN: llc -filetype=obj -o %t/main.o %S/Inputs/main-ret-0.ll
+# RUN: llvm-jitlink -noexec -all_load -show-init-es %t/main.o -L%t -lFoo \
+# RUN:     | FileCheck %s
+#
+# Check that the llvm-jitlink -all-load option loads all members of
+# multi-file archives.
+#
+# CHECK-DAG: main{{.*}}main.o
+# CHECK-DAG: foo{{.*}}libFoo.a[0](foo.o)
+# CHECK-DAG: bar{{.*}}libFoo.a[1](bar.o)

llvm/tools/llvm-jitlink/llvm-jitlink.cpp

@@ -448,7 +448,7 @@ static Error applyLibraryLinkModifiers(Session &S, LinkGraph &G) {
   if (!S.HiddenArchives.empty()) {
     StringRef ObjName(G.getName());
     if (ObjName.ends_with(')')) {
-      auto LibName = ObjName.split('(').first;
+      auto LibName = ObjName.split('[').first;
       if (S.HiddenArchives.count(LibName)) {
         for (auto *Sym : G.defined_symbols())
           Sym->setScope(std::max(Sym->getScope(), Scope::Hidden));
@@ -2279,8 +2279,26 @@ static Error addLibraries(Session &S,
 
     auto &LinkLayer = S.getLinkLayer(LazyLinkIdxs.count(LL.Position));
 
+    std::set<std::string> ImportedDynamicLibraries;
     StaticLibraryDefinitionGenerator::VisitMembersFunction VisitMembers;
-    if (AllLoad)
+
+    // COFF gets special handling due to import libraries.
+    if (S.ES.getTargetTriple().isOSBinFormatCOFF()) {
+      if (AllLoad) {
+        VisitMembers =
+            [ImportScanner = COFFImportFileScanner(ImportedDynamicLibraries),
+             LoadAll =
+                 StaticLibraryDefinitionGenerator::loadAllObjectFileMembers(
+                     LinkLayer, JD)](object::Archive &A,
+                                     MemoryBufferRef MemberBuf,
+                                     size_t Index) mutable -> Expected<bool> {
+          if (!ImportScanner(A, MemberBuf, Index))
+            return false;
+          return LoadAll(A, MemberBuf, Index);
+        };
+      } else
+        VisitMembers = COFFImportFileScanner(ImportedDynamicLibraries);
+    } else if (AllLoad)
       VisitMembers = StaticLibraryDefinitionGenerator::loadAllObjectFileMembers(
           LinkLayer, JD);
     else if (S.ES.getTargetTriple().isOSBinFormatMachO() && ForceLoadObjC)
@@ -2294,7 +2312,7 @@ static Error addLibraries(Session &S,
 
     // Push additional dynamic libraries to search.
     // Note that this mechanism only happens in COFF.
-    for (auto FileName : (*G)->getImportedDynamicLibraries()) {
+    for (auto FileName : ImportedDynamicLibraries) {
       LibraryLoad NewLL;
       auto FileNameRef = StringRef(FileName);
       if (!FileNameRef.ends_with_insensitive(".dll"))

llvm/tools/llvm-jitlink/llvm-jitlink.h

@@ -14,6 +14,7 @@
 #define LLVM_TOOLS_LLVM_JITLINK_LLVM_JITLINK_H
 
 #include "llvm/ADT/StringSet.h"
+#include "llvm/ExecutionEngine/Orc/COFF.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
 #include "llvm/ExecutionEngine/Orc/LazyObjectLinkingLayer.h"