[ORC][llvm-jitlink] Add support for emulating ld64 -weak-lx / -weak_library.

Linking libraries in ld64 with -weak-lx / -weak_library causes all references
to symbols in those libraries to be made weak, allowing the librarie to be
missing at runtime.

This patch extends EPCDynamicLibrarySearchGenerator with support for emulating
this behavior: If an instance is constructed with an Allow predicate but no
dylib handle then all symbols matching the predicate are immediately resolved
to null.

The llvm-jitlink tool is updated with -weak-lx / -weak_library options for
testing. Unlike their ld64 counterparts these options take a TBD file as input,
and always resolve all exports in the TBD file to null.

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h

@@ -43,6 +43,19 @@ class EPCDynamicLibrarySearchGenerator : public DefinitionGenerator {
       : EPC(ES.getExecutorProcessControl()), H(H), Allow(std::move(Allow)),
         AddAbsoluteSymbols(std::move(AddAbsoluteSymbols)) {}
 
+  /// Create an EPCDynamicLibrarySearchGenerator that resolves all symbols
+  /// matching the Allow predicate to null. This can be used to emulate linker
+  /// options like -weak-l / -weak_library where the library is missing at
+  /// runtime. (Note: here we're explicitly returning null for these symbols,
+  /// rather than returning no value at all for them, which is the usual
+  /// "missing symbol" behavior in ORC. This distinction shouldn't matter for
+  /// most use-cases).
+  EPCDynamicLibrarySearchGenerator(
+      ExecutionSession &ES, SymbolPredicate Allow,
+      AddAbsoluteSymbolsFn AddAbsoluteSymbols = nullptr)
+      : EPC(ES.getExecutorProcessControl()), Allow(std::move(Allow)),
+        AddAbsoluteSymbols(std::move(AddAbsoluteSymbols)) {}
+
   /// Permanently loads the library at the given path and, on success, returns
   /// an EPCDynamicLibrarySearchGenerator that will search it for symbol
   /// definitions in the library. On failure returns the reason the library
@@ -66,8 +79,10 @@ class EPCDynamicLibrarySearchGenerator : public DefinitionGenerator {
                       const SymbolLookupSet &Symbols) override;
 
 private:
+  Error addAbsolutes(JITDylib &JD, SymbolMap Symbols);
+
   ExecutorProcessControl &EPC;
-  tpctypes::DylibHandle H;
+  std::optional<tpctypes::DylibHandle> H;
   SymbolPredicate Allow;
   AddAbsoluteSymbolsFn AddAbsoluteSymbols;
 };

llvm/include/llvm/ExecutionEngine/Orc/GetTapiInterface.h

@@ -0,0 +1,28 @@
+//===---- GetTapiInterface.h -- Get interface from TAPI file ----*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Get symbol interface from TAPI file.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_GETTAPIINTERFACE_H
+#define LLVM_EXECUTIONENGINE_ORC_GETTAPIINTERFACE_H
+
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/Object/TapiUniversal.h"
+
+namespace llvm::orc {
+
+/// Returns a SymbolNameSet containing the exported symbols defined in the
+/// relevant slice of the TapiUniversal file.
+Expected<SymbolNameSet> getInterfaceFromTapiFile(ExecutionSession &ES,
+                                                 object::TapiUniversal &TU);
+
+} // namespace llvm::orc
+
+#endif // LLVM_EXECUTIONENGINE_ORC_GETTAPIINTERFACE_H

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

@@ -25,6 +25,7 @@ add_llvm_component_library(LLVMOrcJIT
   EPCIndirectionUtils.cpp
   ExecutionUtils.cpp
   ObjectFileInterface.cpp
+  GetTapiInterface.cpp
   IndirectionUtils.cpp
   IRCompileLayer.cpp
   IRTransformLayer.cpp

llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp

@@ -42,6 +42,18 @@ Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
              << Symbols << "\n";
     });
 
+  // If there's no handle then resolve all requested symbols to null.
+  if (!H) {
+    assert(Allow && "No handle or filter?");
+    SymbolMap Nulls;
+    for (auto &[Name, LookupFlags] : Symbols) {
+      if (Allow(Name))
+        Nulls[Name] = {};
+    }
+    return addAbsolutes(JD, std::move(Nulls));
+  }
+
+  // Otherwise proceed with lookup in the remote.
   SymbolLookupSet LookupSymbols;
 
   for (auto &KV : Symbols) {
@@ -51,7 +63,7 @@ Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
     LookupSymbols.add(KV.first, SymbolLookupFlags::WeaklyReferencedSymbol);
   }
 
-  DylibManager::LookupRequest Request(H, LookupSymbols);
+  DylibManager::LookupRequest Request(*H, LookupSymbols);
   // Copy-capture LookupSymbols, since LookupRequest keeps a reference.
   EPC.getDylibMgr().lookupSymbolsAsync(Request, [this, &JD, LS = std::move(LS),
                                                  LookupSymbols](
@@ -85,15 +97,19 @@ Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
       return LS.continueLookup(Error::success());
 
     // Define resolved symbols.
-    Error Err = AddAbsoluteSymbols
-                    ? AddAbsoluteSymbols(JD, std::move(NewSymbols))
-                    : JD.define(absoluteSymbols(std::move(NewSymbols)));
+    Error Err = addAbsolutes(JD, std::move(NewSymbols));
 
     LS.continueLookup(std::move(Err));
   });
 
   return Error::success();
 }
 
+Error EPCDynamicLibrarySearchGenerator::addAbsolutes(JITDylib &JD,
+                                                     SymbolMap Symbols) {
+  return AddAbsoluteSymbols ? AddAbsoluteSymbols(JD, std::move(Symbols))
+                            : JD.define(absoluteSymbols(std::move(Symbols)));
+}
+
 } // end namespace orc
 } // end namespace llvm

llvm/lib/ExecutionEngine/Orc/GetTapiInterface.cpp

@@ -0,0 +1,39 @@
+//===--------- GetTapiInterface.cpp - Get interface from TAPI file --------===//
+//
+// 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/GetTapiInterface.h"
+
+#define DEBUG_TYPE "orc"
+
+namespace llvm::orc {
+
+Expected<SymbolNameSet> getInterfaceFromTapiFile(ExecutionSession &ES,
+                                                 object::TapiUniversal &TU) {
+  SymbolNameSet Symbols;
+
+  auto CPUType = MachO::getCPUType(ES.getTargetTriple());
+  if (!CPUType)
+    return CPUType.takeError();
+
+  auto CPUSubType = MachO::getCPUSubType(ES.getTargetTriple());
+  if (!CPUSubType)
+    return CPUSubType.takeError();
+
+  auto &TUIF = TU.getInterfaceFile();
+  auto ArchInterface =
+      TUIF.extract(MachO::getArchitectureFromCpuType(*CPUType, *CPUSubType));
+  if (!ArchInterface)
+    return ArchInterface.takeError();
+
+  for (auto *Sym : (*ArchInterface)->exports())
+    Symbols.insert(ES.intern(Sym->getName()));
+
+  return Symbols;
+}
+
+} // namespace llvm::orc

llvm/test/ExecutionEngine/JITLink/AArch64/Inputs/MachO_Foo.tbd

@@ -0,0 +1,23 @@
+--- !tapi-tbd
+tbd-version: 4
+targets:  [ arm64-macos ]
+uuids:
+  - target: arm64-macos
+    value: 00000000-0000-0000-0000-000000000000
+flags: [ installapi ]
+install-name: Foo.framework/Foo
+current-version: 1.2.3
+compatibility-version: 1.2
+swift-abi-version: 5
+parent-umbrella:
+  - targets: [ arm64-macos ]
+    umbrella: System
+exports:
+  - targets: [ arm64-macos ]
+    symbols: [ _foo ]
+    objc-classes: []
+    objc-eh-types: []
+    objc-ivars: []
+    weak-symbols: []
+    thread-local-symbols: []
+...

llvm/test/ExecutionEngine/JITLink/AArch64/Inputs/MachO_main_ret_foo.s

@@ -0,0 +1,15 @@
+        .section        __TEXT,__text,regular,pure_instructions
+        .globl  _main
+        .p2align        2
+_main:
+        .cfi_startproc
+Lloh0:
+        adrp    x0, _foo@GOTPAGE
+Lloh1:
+        ldr     x0, [x0, _foo@GOTPAGEOFF]
+
+        ret
+        .loh AdrpLdrGot Lloh0, Lloh1
+        .cfi_endproc
+
+.subsections_via_symbols

llvm/test/ExecutionEngine/JITLink/AArch64/MachO_weak_link.test

@@ -0,0 +1,8 @@
+# RUN: rm -rf %t && mkdir -p %t
+# RUN: llvm-mc -triple=arm64-apple-darwin19 -filetype=obj -o %t/main.o \
+# RUN:         %S/Inputs/MachO_main_ret_foo.s
+# RUN: llvm-jitlink -noexec %t/main.o -weak_library %S/Inputs/MachO_Foo.tbd
+
+# Check that we can load main.o, which unconditionally uses symbol foo, by
+# using -weak_library on a TBD file to emulate forced weak linking against
+# a library that supplies foo, but is missing at runtime.

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

@@ -29,6 +29,7 @@
 #include "llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h"
 #include "llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/GetTapiInterface.h"
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h"
 #include "llvm/ExecutionEngine/Orc/JITLinkReentryTrampolines.h"
@@ -58,6 +59,7 @@
 #include "llvm/Object/COFF.h"
 #include "llvm/Object/MachO.h"
 #include "llvm/Object/ObjectFile.h"
+#include "llvm/Object/TapiUniversal.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/InitLLVM.h"
@@ -140,6 +142,20 @@ static cl::list<std::string>
                cl::desc("Link against library X with hidden visibility"),
                cl::cat(JITLinkCategory));
 
+static cl::list<std::string>
+    LibrariesWeak("weak-l",
+                  cl::desc("Emulate weak link against library X. Must resolve "
+                           "to a TextAPI file, and all symbols in the "
+                           "interface will resolve to null."),
+                  cl::Prefix, cl::cat(JITLinkCategory));
+
+static cl::list<std::string> WeakLibraries(
+    "weak_library",
+    cl::desc("Emulate weak link against library X. X must point to a "
+             "TextAPI file, and all symbols in the interface will "
+             "resolve to null"),
+    cl::cat(JITLinkCategory));
+
 static cl::opt<bool> SearchSystemLibrary(
     "search-sys-lib",
     cl::desc("Add system library paths to library search paths"),
@@ -2100,6 +2116,27 @@ static SmallVector<StringRef, 5> getSearchPathsFromEnvVar(Session &S) {
   return PathVec;
 }
 
+static Expected<std::unique_ptr<DefinitionGenerator>>
+LoadLibraryWeak(Session &S, StringRef InterfacePath) {
+  auto TapiFileBuffer = getFile(InterfacePath);
+  if (!TapiFileBuffer)
+    return TapiFileBuffer.takeError();
+
+  auto Tapi =
+      object::TapiUniversal::create((*TapiFileBuffer)->getMemBufferRef());
+  if (!Tapi)
+    return Tapi.takeError();
+
+  auto Symbols = getInterfaceFromTapiFile(S.ES, **Tapi);
+  if (!Symbols)
+    return Symbols.takeError();
+
+  return std::make_unique<EPCDynamicLibrarySearchGenerator>(
+      S.ES, [Symbols = std::move(*Symbols)](const SymbolStringPtr &Sym) {
+        return Symbols.count(Sym);
+      });
+}
+
 static Error addLibraries(Session &S,
                           const std::map<unsigned, JITDylib *> &IdxToJD,
                           const DenseSet<unsigned> &LazyLinkIdxs) {
@@ -2142,11 +2179,12 @@ static Error addLibraries(Session &S,
     bool IsPath = false;
     unsigned Position;
     ArrayRef<StringRef> CandidateExtensions;
-    enum { Standard, Hidden } Modifier;
+    enum { Standard, Hidden, Weak } Modifier;
   };
 
   // Queue to load library as in the order as it appears in the argument list.
   std::deque<LibraryLoad> LibraryLoadQueue;
+
   // Add archive files from the inputs to LibraryLoads.
   for (auto InputFileItr = InputFiles.begin(), InputFileEnd = InputFiles.end();
        InputFileItr != InputFileEnd; ++InputFileItr) {
@@ -2173,9 +2211,23 @@ static Error addLibraries(Session &S,
     LL.Modifier = LibraryLoad::Hidden;
     LibraryLoadQueue.push_back(std::move(LL));
   }
+
+  // Add -weak_library arguments to LibraryLoads.
+  for (auto LibItr = WeakLibraries.begin(), LibEnd = WeakLibraries.end();
+       LibItr != LibEnd; ++LibItr) {
+    LibraryLoad LL;
+    LL.LibName = *LibItr;
+    LL.IsPath = true;
+    LL.Position = WeakLibraries.getPosition(LibItr - WeakLibraries.begin());
+    LL.CandidateExtensions = {};
+    LL.Modifier = LibraryLoad::Weak;
+    LibraryLoadQueue.push_back(std::move(LL));
+  }
+
   StringRef StandardExtensions[] = {".so", ".dylib", ".dll", ".a", ".lib"};
   StringRef DynLibExtensionsOnly[] = {".so", ".dylib", ".dll"};
   StringRef ArchiveExtensionsOnly[] = {".a", ".lib"};
+  StringRef InterfaceExtensionsOnly = {".tbd"};
 
   // Add -lx arguments to LibraryLoads.
   for (auto LibItr = Libraries.begin(), LibEnd = Libraries.end();
@@ -2201,6 +2253,18 @@ static Error addLibraries(Session &S,
     LibraryLoadQueue.push_back(std::move(LL));
   }
 
+  // Add -weak-lx arguments to LibraryLoads.
+  for (auto LibWeakItr = LibrariesWeak.begin(),
+            LibWeakEnd = LibrariesWeak.end();
+       LibWeakItr != LibWeakEnd; ++LibWeakItr) {
+    LibraryLoad LL;
+    LL.LibName = *LibWeakItr;
+    LL.Position = LibrariesWeak.getPosition(LibWeakItr - LibrariesWeak.begin());
+    LL.CandidateExtensions = InterfaceExtensionsOnly;
+    LL.Modifier = LibraryLoad::Weak;
+    LibraryLoadQueue.push_back(std::move(LL));
+  }
+
   // Sort library loads by position in the argument list.
   llvm::sort(LibraryLoadQueue,
              [](const LibraryLoad &LHS, const LibraryLoad &RHS) {
@@ -2220,6 +2284,9 @@ static Error addLibraries(Session &S,
       GetObjFileInterface = getObjectFileInterfaceHidden;
       S.HiddenArchives.insert(Path);
       break;
+    case LibraryLoad::Weak:
+      llvm_unreachable("Unsupported");
+      break;
     }
 
     auto &LinkLayer = S.getLinkLayer(LazyLinkIdxs.count(LL.Position));
@@ -2266,11 +2333,26 @@ static Error addLibraries(Session &S,
 
     // If this is the name of a JITDylib then link against that.
     if (auto *LJD = S.ES.getJITDylibByName(LL.LibName)) {
+      if (LL.Modifier == LibraryLoad::Weak)
+        return make_error<StringError>(
+            "Can't use -weak-lx or -weak_library to load JITDylib " +
+                LL.LibName,
+            inconvertibleErrorCode());
       JD.addToLinkOrder(*LJD);
       continue;
     }
 
     if (LL.IsPath) {
+      // Must be -weak_library.
+      if (LL.Modifier == LibraryLoad::Weak) {
+        if (auto G = LoadLibraryWeak(S, LL.LibName)) {
+          JD.addGenerator(std::move(*G));
+          continue;
+        } else
+          return G.takeError();
+      }
+
+      // Otherwise handle archive.
       auto G = AddArchive(JD, LL.LibName.c_str(), LL);
       if (!G)
         return createFileError(LL.LibName, G.takeError());
@@ -2337,6 +2419,14 @@ static Error addLibraries(Session &S,
           });
           break;
         }
+        case file_magic::tapi_file:
+          assert(LL.Modifier == LibraryLoad::Weak &&
+                 "TextAPI file not being loaded as weak?");
+          if (auto G = LoadLibraryWeak(S, LibPath.data()))
+            JD.addGenerator(std::move(*G));
+          else
+            return G.takeError();
+          break;
         default:
           // This file isn't a recognized library kind.
           LLVM_DEBUG({