Drop linker mangling for in-memory MachO objects in Linux with LLVM JITLink

MachoOnLinuxSearchGenerator.md

Please use this with a grain of salt: linking non-native object formats can cause arbitrary ABI incompatibilities.

#include <llvm/ADT/FunctionExtras.h>
#include <llvm/ExecutionEngine/Orc/TargetProcessControl.h>

class MachoOnLinuxSearchGenerator : public llvm::orc::DefinitionGenerator {
public:
  /// Create a DynamicLibrarySearchGenerator that searches for symbols in the
  /// library with the given handle. It drops one leading underscore where it
  /// exists, before issuing the lookup.
  MachoOnLinuxSearchGenerator(llvm::orc::TargetProcessControl &TPC,
                              llvm::orc::tpctypes::DylibHandle H)
      : TPC(TPC), H(H) {}

  /// Permanently loads the library at the given path and, on success, returns
  /// a DynamicLibrarySearchGenerator that will search it for symbol definitions
  /// in the library. On failure returns the reason the library failed to load.
  static llvm::Expected<std::unique_ptr<MachoOnLinuxSearchGenerator>>
  Load(llvm::orc::TargetProcessControl &TPC, const char *LibraryPath);

  /// Creates a MachoOnLinuxSearchGenerator that searches for symbols in
  /// the target process.
  static llvm::Expected<std::unique_ptr<MachoOnLinuxSearchGenerator>>
  GetForTargetProcess(llvm::orc::TargetProcessControl &TPC) {
    return Load(TPC, nullptr);
  }

  llvm::Error tryToGenerate(llvm::orc::LookupState &LS, llvm::orc::LookupKind K,
                            llvm::orc::JITDylib &JD,
                            llvm::orc::JITDylibLookupFlags JDLookupFlags,
                            const llvm::orc::SymbolLookupSet &Symbols) override;

private:
  llvm::orc::TargetProcessControl &TPC;
  llvm::orc::tpctypes::DylibHandle H;
};

#include "MachoOnLinuxSearchGenerator.h"

using namespace llvm;
using namespace llvm::orc;

static StringRef dropMachOLinkerMangling(StringRef LinkerMangledName) {
  if (*LinkerMangledName.data() == '_')
    return StringRef(LinkerMangledName.data() + 1);
  return LinkerMangledName;
}

Expected<std::unique_ptr<MachoOnLinuxSearchGenerator>>
MachoOnLinuxSearchGenerator::Load(TargetProcessControl &TPC,
                                  const char *LibraryPath) {
  auto Handle = TPC.loadDylib(LibraryPath);
  if (!Handle)
    return Handle.takeError();

  return std::make_unique<MachoOnLinuxSearchGenerator>(TPC, *Handle);
}

Error MachoOnLinuxSearchGenerator::tryToGenerate(
    LookupState &LS, LookupKind K, JITDylib &JD,
    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols) {

  if (Symbols.empty())
    return Error::success();

  SymbolLookupSet LookupSymbols;

  // Lookup all symbols as if they were weak, so that missing symbols don't
  // cause hard failures.
  for (auto &KV : Symbols) {
    StringRef Name = dropMachOLinkerMangling(*KV.first);
    LookupSymbols.add(TPC.intern(Name),
                      SymbolLookupFlags::WeaklyReferencedSymbol);
  }

  SymbolMap NewSymbols;

  tpctypes::LookupRequest Request(H, LookupSymbols);
  using LookupResults = std::vector<tpctypes::LookupResult>;
  Expected<LookupResults> Result = TPC.lookupSymbols(Request);
  if (!Result)
    return Result.takeError();

  assert(Result->size() == 1 && "Results for more than one library returned");
  assert(Result->front().size() == LookupSymbols.size() &&
         "Result has incorrect number of elements");

  auto ResultI = Result->front().begin();
  for (auto &KV : Symbols) {
    if (*ResultI)
      NewSymbols[KV.first] =
          JITEvaluatedSymbol(*ResultI, JITSymbolFlags::Exported);
    ResultI++;
  }

  // If there were no resolved symbols bail out.
  if (NewSymbols.empty())
    return Error::success();

  // Define resolved symbols.
  return JD.define(absoluteSymbols(std::move(NewSymbols)));
}