[ORC] Generalize loadRelocatableObject to loadLinkableFile, add archive support.

This allows us to rewrite part of StaticLibraryDefinitionGenerator in terms of
loadLinkableFile.

It's also useful for clients who may not know (either from file extensions or
context) whether a given path will be an object file, an archive, or a
universal binary.

rdar://134638070

Author: lhames

llvm/include/llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h

@@ -25,6 +25,7 @@
 #include "llvm/Support/Memory.h"
 #include "llvm/Support/RecyclingAllocator.h"
 
+#include <cassert>
 #include <cstdint>
 #include <future>
 #include <mutex>
@@ -363,7 +364,9 @@ class InProcessMemoryManager : public JITLinkMemoryManager {
   static Expected<std::unique_ptr<InProcessMemoryManager>> Create();
 
   /// Create an instance using the given page size.
-  InProcessMemoryManager(uint64_t PageSize) : PageSize(PageSize) {}
+  InProcessMemoryManager(uint64_t PageSize) : PageSize(PageSize) {
+    assert(isPowerOf2_64(PageSize) && "PageSize must be a power of 2");
+  }
 
   void allocate(const JITLinkDylib *JD, LinkGraph &G,
                 OnAllocatedFunction OnAllocated) override;

llvm/include/llvm/ExecutionEngine/Orc/LoadLinkableFile.h

@@ -0,0 +1,64 @@
+//===--- LoadLinkableFile.h -- Load relocatables and archives ---*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// A wrapper for `MemoryBuffer::getFile` / `MemoryBuffer::getFileSlice` that:
+//
+//   1. Handles relocatable object files, archives, and macho universal
+//      binaries.
+//   2. Adds file paths to errors by default.
+//   3. Checks architecture compatibility up-front.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_LOADLINKABLEFILE_H
+#define LLVM_EXECUTIONENGINE_ORC_LOADLINKABLEFILE_H
+
+#include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/TargetParser/Triple.h"
+
+namespace llvm {
+namespace orc {
+
+enum class LinkableFileKind { Archive, RelocatableObject };
+
+enum LoadArchives {
+  Never,   // Linkable file must not be an archive.
+  Allowed, // Linkable file is allowed to be an archive.
+  Required // Linkable file is required to be an archive.
+};
+
+/// Create a MemoryBuffer covering the "linkable" part of the given path.
+///
+/// The path must contain a relocatable object file or universal binary, or
+/// (if AllowArchives is true) an archive.
+///
+/// If the path is a universal binary then it must contain a slice whose
+/// architecture matches the architecture in the triple (an error will be
+/// returned if there is no such slice, or if the triple does not specify an
+/// architectur).
+///
+/// If the path (or universal binary slice) is a relocatable object file then
+/// its architecture must match the architecture in the triple (if given).
+///
+/// If the path (or universal binary slice) is a relocatable object file then
+/// its format must match the format in the triple (if given).
+///
+/// No verification (e.g. architecture or format) is performed on the contents
+/// of archives.
+///
+/// If IdentifierOverride is provided then it will be used as the name of the
+/// resulting buffer, rather than Path.
+Expected<std::pair<std::unique_ptr<MemoryBuffer>, LinkableFileKind>>
+loadLinkableFile(StringRef Path, const Triple &TT, LoadArchives LA,
+                 std::optional<StringRef> IdentifierOverride = std::nullopt);
+
+} // End namespace orc
+} // End namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_LOADLINKABLEFILE_H

llvm/include/llvm/ExecutionEngine/Orc/LoadRelocatableObject.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_EXECUTIONENGINE_ORC_MACHO_H
 #define LLVM_EXECUTIONENGINE_ORC_MACHO_H
 
+#include "llvm/ExecutionEngine/Orc/LoadLinkableFile.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/TargetParser/Triple.h"
@@ -45,16 +46,21 @@ checkMachORelocatableObject(std::unique_ptr<MemoryBuffer> Obj, const Triple &TT,
 /// Load a relocatable object compatible with TT from Path.
 /// If Path is a universal binary, this function will return a buffer for the
 /// slice compatible with Triple (if one is present).
-Expected<std::unique_ptr<MemoryBuffer>> loadMachORelocatableObject(
-    StringRef Path, const Triple &TT,
+Expected<std::pair<std::unique_ptr<MemoryBuffer>, LinkableFileKind>>
+loadMachOLinkableFile(
+    StringRef Path, const Triple &TT, LoadArchives LA,
     std::optional<StringRef> IdentifierOverride = std::nullopt);
 
 /// Load a compatible relocatable object (if available) from a MachO universal
 /// binary.
-Expected<std::unique_ptr<MemoryBuffer>>
-loadMachORelocatableObjectFromUniversalBinary(
-    StringRef UBPath, std::unique_ptr<MemoryBuffer> UBBuf, const Triple &TT,
-    std::optional<StringRef> IdentifierOverride = std::nullopt);
+/// Path is only used for error reporting. Identifier will be used to name the
+/// resulting buffer.
+Expected<std::pair<std::unique_ptr<MemoryBuffer>, LinkableFileKind>>
+loadLinkableSliceFromMachOUniversalBinary(sys::fs::file_t FD,
+                                          std::unique_ptr<MemoryBuffer> UBBuf,
+                                          const Triple &TT, LoadArchives LA,
+                                          StringRef UBPath,
+                                          StringRef Identifier);
 
 /// Utility for identifying the file-slice compatible with TT in a universal
 /// binary.

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

@@ -326,22 +326,21 @@ class InProcessMemoryManager::IPInFlightAlloc
 
 Expected<std::unique_ptr<InProcessMemoryManager>>
 InProcessMemoryManager::Create() {
-  if (auto PageSize = sys::Process::getPageSize())
+  if (auto PageSize = sys::Process::getPageSize()) {
+    // FIXME: Just check this once on startup.
+    if (!isPowerOf2_64((uint64_t)*PageSize))
+      return make_error<StringError>(
+          "Could not create InProcessMemoryManager: Page size " +
+              Twine(*PageSize) + " is not a power of 2",
+          inconvertibleErrorCode());
+
     return std::make_unique<InProcessMemoryManager>(*PageSize);
-  else
+  } else
     return PageSize.takeError();
 }
 
 void InProcessMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
                                       OnAllocatedFunction OnAllocated) {
-
-  // FIXME: Just check this once on startup.
-  if (!isPowerOf2_64((uint64_t)PageSize)) {
-    OnAllocated(make_error<StringError>("Page size is not a power of 2",
-                                        inconvertibleErrorCode()));
-    return;
-  }
-
   BasicLayout BL(G);
 
   /// Scan the request and calculate the group and total sizes.

llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp

@@ -29,7 +29,7 @@ add_llvm_component_library(LLVMOrcJIT
   JITTargetMachineBuilder.cpp
   LazyReexports.cpp
   Layer.cpp
-  LoadRelocatableObject.cpp
+  LoadLinkableFile.cpp
   LookupAndRecordAddrs.cpp
   LLJIT.cpp
   MachO.cpp

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

@@ -9,6 +9,7 @@
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/JITLink/x86_64.h"
 #include "llvm/ExecutionEngine/Orc/Layer.h"
+#include "llvm/ExecutionEngine/Orc/LoadLinkableFile.h"
 #include "llvm/ExecutionEngine/Orc/MachO.h"
 #include "llvm/ExecutionEngine/Orc/ObjectFileInterface.h"
 #include "llvm/IR/Constants.h"
@@ -277,45 +278,12 @@ StaticLibraryDefinitionGenerator::Load(
     ObjectLayer &L, const char *FileName,
     GetObjectFileInterface GetObjFileInterface) {
 
-  auto B = object::createBinary(FileName);
-  if (!B)
-    return createFileError(FileName, B.takeError());
-
-  // If this is a regular archive then create an instance from it.
-  if (isa<object::Archive>(B->getBinary())) {
-    auto [Archive, ArchiveBuffer] = B->takeBinary();
-    return Create(L, std::move(ArchiveBuffer),
-                  std::unique_ptr<object::Archive>(
-                      static_cast<object::Archive *>(Archive.release())),
-                  std::move(GetObjFileInterface));
-  }
-
-  // If this is a universal binary then search for a slice matching the given
-  // Triple.
-  if (auto *UB = dyn_cast<object::MachOUniversalBinary>(B->getBinary())) {
+  const auto &TT = L.getExecutionSession().getTargetTriple();
+  auto Linkable = loadLinkableFile(FileName, TT, LoadArchives::Required);
+  if (!Linkable)
+    return Linkable.takeError();
 
-    const auto &TT = L.getExecutionSession().getTargetTriple();
-
-    auto SliceRange = getMachOSliceRangeForTriple(*UB, TT);
-    if (!SliceRange)
-      return SliceRange.takeError();
-
-    auto SliceBuffer = MemoryBuffer::getFileSlice(FileName, SliceRange->second,
-                                                  SliceRange->first);
-    if (!SliceBuffer)
-      return make_error<StringError>(
-          Twine("Could not create buffer for ") + TT.str() + " slice of " +
-              FileName + ": [ " + formatv("{0:x}", SliceRange->first) + " .. " +
-              formatv("{0:x}", SliceRange->first + SliceRange->second) + ": " +
-              SliceBuffer.getError().message(),
-          SliceBuffer.getError());
-
-    return Create(L, std::move(*SliceBuffer), std::move(GetObjFileInterface));
-  }
-
-  return make_error<StringError>(Twine("Unrecognized file type for ") +
-                                     FileName,
-                                 inconvertibleErrorCode());
+  return Create(L, std::move(Linkable->first), std::move(GetObjFileInterface));
 }
 
 Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>

llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp

@@ -0,0 +1,121 @@
+//===------- LoadLinkableFile.cpp -- Load relocatables and archives -------===//
+//
+// 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/LoadLinkableFile.h"
+
+#include "llvm/ADT/ScopeExit.h"
+#include "llvm/BinaryFormat/Magic.h"
+#include "llvm/ExecutionEngine/Orc/MachO.h"
+#include "llvm/Support/FileSystem.h"
+
+#define DEBUG_TYPE "orc"
+
+namespace llvm {
+namespace orc {
+
+static Expected<std::unique_ptr<MemoryBuffer>>
+checkCOFFRelocatableObject(std::unique_ptr<MemoryBuffer> Obj,
+                           const Triple &TT) {
+  // TODO: Actually check the architecture of the file.
+  return std::move(Obj);
+}
+
+static Expected<std::unique_ptr<MemoryBuffer>>
+checkELFRelocatableObject(std::unique_ptr<MemoryBuffer> Obj, const Triple &TT) {
+  // TODO: Actually check the architecture of the file.
+  return std::move(Obj);
+}
+
+Expected<std::pair<std::unique_ptr<MemoryBuffer>, LinkableFileKind>>
+loadLinkableFile(StringRef Path, const Triple &TT, LoadArchives LA,
+                 std::optional<StringRef> IdentifierOverride) {
+  if (!IdentifierOverride)
+    IdentifierOverride = Path;
+
+  Expected<sys::fs::file_t> FDOrErr =
+      sys::fs::openNativeFileForRead(Path, sys::fs::OF_None);
+  if (!FDOrErr)
+    return createFileError(Path, FDOrErr.takeError());
+  sys::fs::file_t FD = *FDOrErr;
+  auto CloseFile = make_scope_exit([&]() { sys::fs::closeFile(FD); });
+
+  auto Buf =
+      MemoryBuffer::getOpenFile(FD, *IdentifierOverride, /*FileSize=*/-1);
+  if (!Buf)
+    return make_error<StringError>(
+        StringRef("Could not load object at path ") + Path, Buf.getError());
+
+  std::optional<Triple::ObjectFormatType> RequireFormat;
+  if (TT.getObjectFormat() != Triple::UnknownObjectFormat)
+    RequireFormat = TT.getObjectFormat();
+
+  switch (identify_magic((*Buf)->getBuffer())) {
+  case file_magic::archive:
+    if (LA != LoadArchives::Never)
+      return std::make_pair(std::move(*Buf), LinkableFileKind::Archive);
+    return make_error<StringError>(
+        Path + " does not contain a relocatable object file",
+        inconvertibleErrorCode());
+  case file_magic::coff_object:
+    if (LA == LoadArchives::Required)
+      return make_error<StringError>(Path + " does not contain an archive",
+                                     inconvertibleErrorCode());
+
+    if (!RequireFormat || *RequireFormat == Triple::COFF) {
+      auto CheckedBuf = checkCOFFRelocatableObject(std::move(*Buf), TT);
+      if (!CheckedBuf)
+        return CheckedBuf.takeError();
+      return std::make_pair(std::move(*CheckedBuf),
+                            LinkableFileKind::RelocatableObject);
+    }
+    break;
+  case file_magic::elf_relocatable:
+    if (LA == LoadArchives::Required)
+      return make_error<StringError>(Path + " does not contain an archive",
+                                     inconvertibleErrorCode());
+
+    if (!RequireFormat || *RequireFormat == Triple::ELF) {
+      auto CheckedBuf = checkELFRelocatableObject(std::move(*Buf), TT);
+      if (!CheckedBuf)
+        return CheckedBuf.takeError();
+      return std::make_pair(std::move(*CheckedBuf),
+                            LinkableFileKind::RelocatableObject);
+    }
+    break;
+  case file_magic::macho_object:
+    if (LA == LoadArchives::Required)
+      return make_error<StringError>(Path + " does not contain an archive",
+                                     inconvertibleErrorCode());
+
+    if (!RequireFormat || *RequireFormat == Triple::MachO) {
+      auto CheckedBuf = checkMachORelocatableObject(std::move(*Buf), TT, false);
+      if (!CheckedBuf)
+        return CheckedBuf.takeError();
+      return std::make_pair(std::move(*CheckedBuf),
+                            LinkableFileKind::RelocatableObject);
+    }
+    break;
+  case file_magic::macho_universal_binary:
+    if (!RequireFormat || *RequireFormat == Triple::MachO)
+      return loadLinkableSliceFromMachOUniversalBinary(
+          FD, std::move(*Buf), TT, LA, Path, *IdentifierOverride);
+    break;
+  default:
+    break;
+  }
+
+  return make_error<StringError>(
+      Path +
+          " does not contain a relocatable object file or archive compatible "
+          "with " +
+          TT.str(),
+      inconvertibleErrorCode());
+}
+
+} // End namespace orc.
+} // End namespace llvm.