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

Author: lhames

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 @@ 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,

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

@@ -273,16 +273,33 @@ 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.
   /// Archive members that are not valid object files will be skipped.
   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

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;

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));