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

via llvm-commits llvm-commits at lists.llvm.org
Tue May 27 03:58:56 PDT 2025


Author: Lang Hames
Date: 2025-05-27T20:58:53+10:00
New Revision: 6fa8657a622173c177d863b763550de4d388f73c

URL: https://github.com/llvm/llvm-project/commit/6fa8657a622173c177d863b763550de4d388f73c
DIFF: https://github.com/llvm/llvm-project/commit/6fa8657a622173c177d863b763550de4d388f73c.diff

LOG: [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.

Added: 
    llvm/include/llvm/ExecutionEngine/Orc/COFF.h
    llvm/lib/ExecutionEngine/Orc/COFF.cpp
    llvm/test/ExecutionEngine/JITLink/Generic/Inputs/bar-initializer.ll
    llvm/test/ExecutionEngine/JITLink/Generic/Inputs/bar-ret-0.ll
    llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-initializer.ll
    llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-ret-0.ll
    llvm/test/ExecutionEngine/JITLink/Generic/all-load-multifile-archive-with-duplicate-member-filenames.test
    llvm/test/ExecutionEngine/JITLink/Generic/all-load-multifile-archive.test

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h
    llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
    llvm/include/llvm/ExecutionEngine/Orc/MachO.h
    llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
    llvm/lib/ExecutionEngine/Orc/COFFPlatform.cpp
    llvm/lib/ExecutionEngine/Orc/COFFVCRuntimeSupport.cpp
    llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
    llvm/lib/ExecutionEngine/Orc/MachO.cpp
    llvm/tools/llvm-jitlink/llvm-jitlink.cpp
    llvm/tools/llvm-jitlink/llvm-jitlink.h

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/COFF.h b/llvm/include/llvm/ExecutionEngine/Orc/COFF.h
new file mode 100644
index 0000000000000..adc9e9e171165
--- /dev/null
+++ b/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

diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h
index acfefbcddc49e..d9a6c41baa3db 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h
@@ -139,6 +139,7 @@ class COFFPlatform : public Platform {
   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,

diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
index 67347e851ca1b..64ac8cb80fc2b 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
@@ -273,9 +273,23 @@ class StaticLibraryDefinitionGenerator : public DefinitionGenerator {
       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 @@ class StaticLibraryDefinitionGenerator : public DefinitionGenerator {
   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 @@ class StaticLibraryDefinitionGenerator : public DefinitionGenerator {
          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

diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/MachO.h b/llvm/include/llvm/ExecutionEngine/Orc/MachO.h
index a9d34a82d53d2..81f889c4933db 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/MachO.h
+++ b/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 @@ class ForceLoadMachOArchiveMembers {
   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;

diff  --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
index 431f9590a874e..6c86f252bf7ea 100644
--- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/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

diff  --git a/llvm/lib/ExecutionEngine/Orc/COFF.cpp b/llvm/lib/ExecutionEngine/Orc/COFF.cpp
new file mode 100644
index 0000000000000..35502f4907e50
--- /dev/null
+++ b/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

diff  --git a/llvm/lib/ExecutionEngine/Orc/COFFPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/COFFPlatform.cpp
index ce04062598b52..d3c5761fa9731 100644
--- a/llvm/lib/ExecutionEngine/Orc/COFFPlatform.cpp
+++ b/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);

diff  --git a/llvm/lib/ExecutionEngine/Orc/COFFVCRuntimeSupport.cpp b/llvm/lib/ExecutionEngine/Orc/COFFVCRuntimeSupport.cpp
index c785381175284..cccb0b996aacb 100644
--- a/llvm/lib/ExecutionEngine/Orc/COFFVCRuntimeSupport.cpp
+++ b/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));
 

diff  --git a/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp b/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
index cf35c904d273a..47764ceffed86 100644
--- a/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
+++ b/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(
-      new StaticLibraryDefinitionGenerator(
-          L, std::move(ArchiveBuffer), std::move(Archive),
-          std::move(GetObjFileInterface), Err));
+  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;
 
-  if (Err)
-    return std::move(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;
 
-  return std::move(ADG);
+      SymbolToMemberIndexMap[ES.intern(Sym.getName())] = EntryItr->second;
+    }
+  }
+
+  return std::unique_ptr<StaticLibraryDefinitionGenerator>(
+      new StaticLibraryDefinitionGenerator(
+          L, std::move(ArchiveBuffer), std::move(Archive),
+          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;
+    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 ChildBuffer = It->second;
-    ChildBufferInfos.insert(
-        {ChildBuffer.getBuffer(), ChildBuffer.getBufferIdentifier()});
-  }
 
-  for (auto ChildBufferInfo : ChildBufferInfos) {
-    MemoryBufferRef ChildBufferRef(ChildBufferInfo.first,
-                                   ChildBufferInfo.second);
+    auto Member = Archive->findSym(*Name);
+    if (!Member)
+      return Member.takeError();
+    if (!*Member) // Skip "none" children.
+      continue;
 
-    auto I = GetObjFileInterface(L.getExecutionSession(), ChildBufferRef);
-    if (!I)
-      return I.takeError();
+    auto MemberBuf = (*Member)->getMemoryBufferRef();
+    if (!MemberBuf)
+      return MemberBuf.takeError();
 
-    if (auto Err = L.add(JD, MemoryBuffer::getMemBuffer(ChildBufferRef, false),
-                         std::move(*I)))
-      return Err;
+    ToLoad[Index] = *MemberBuf;
   }
 
-  return Error::success();
-}
+  // 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);
+  }
 
-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;
-      }
+  // Add loaded files to JITDylib.
+  for (auto &[Index, Buf] : ToLoad) {
+    auto MemberBuf = createMemberBuffer(*Archive, Buf, Index);
 
-      // Give members of the archive a name that contains the archive path so
-      // that they can be 
diff erentiated from a member with the same name in a
-      // 
diff erent 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);
+    auto Interface = GetObjFileInterface(L.getExecutionSession(),
+                                         MemberBuf->getMemBufferRef());
+    if (!Interface)
+      return Interface.takeError();
 
-      MemoryBuffers[DataOffset] = MemBuffer;
-    }
-    if (!Excluded.count(DataOffset))
-      ObjectFilesMap[L.getExecutionSession().intern(SymName)] =
-          MemoryBuffers[DataOffset];
+    if (auto Err = L.add(JD, std::move(MemberBuf), std::move(*Interface)))
+      return Err;
   }
 
   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>

diff  --git a/llvm/lib/ExecutionEngine/Orc/MachO.cpp b/llvm/lib/ExecutionEngine/Orc/MachO.cpp
index 784c3487d64fc..89721d16930c0 100644
--- a/llvm/lib/ExecutionEngine/Orc/MachO.cpp
+++ b/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.

diff  --git a/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/bar-initializer.ll b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/bar-initializer.ll
new file mode 100644
index 0000000000000..14ffb6b8aeb20
--- /dev/null
+++ b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/bar-initializer.ll
@@ -0,0 +1,7 @@
+ at initializer_bar_ran = global i32 0
+ at 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
+}

diff  --git a/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/bar-ret-0.ll b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/bar-ret-0.ll
new file mode 100644
index 0000000000000..4afd29158008c
--- /dev/null
+++ b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/bar-ret-0.ll
@@ -0,0 +1,3 @@
+define i32 @bar()  {
+  ret i32 0
+}

diff  --git a/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-initializer.ll b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-initializer.ll
new file mode 100644
index 0000000000000..2a5ae6f2c5ad7
--- /dev/null
+++ b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-initializer.ll
@@ -0,0 +1,7 @@
+ at initializer_foo_ran = global i32 0
+ at 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
+}

diff  --git a/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-ret-0.ll b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-ret-0.ll
new file mode 100644
index 0000000000000..00969edd3e5d4
--- /dev/null
+++ b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-ret-0.ll
@@ -0,0 +1,3 @@
+define i32 @foo()  {
+  ret i32 0
+}

diff  --git a/llvm/test/ExecutionEngine/JITLink/Generic/all-load-multifile-archive-with-duplicate-member-filenames.test b/llvm/test/ExecutionEngine/JITLink/Generic/all-load-multifile-archive-with-duplicate-member-filenames.test
new file mode 100644
index 0000000000000..edcae5205d8b8
--- /dev/null
+++ b/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 
diff erent '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)

diff  --git a/llvm/test/ExecutionEngine/JITLink/Generic/all-load-multifile-archive.test b/llvm/test/ExecutionEngine/JITLink/Generic/all-load-multifile-archive.test
new file mode 100644
index 0000000000000..576f69804652b
--- /dev/null
+++ b/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)
\ No newline at end of file

diff  --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index ab2f685b4fc1d..c0d5a78e505e8 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/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"))

diff  --git a/llvm/tools/llvm-jitlink/llvm-jitlink.h b/llvm/tools/llvm-jitlink/llvm-jitlink.h
index ca5162f1a2d52..b3faf9fbe9f5f 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.h
+++ b/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"


        


More information about the llvm-commits mailing list