[llvm-jitlink] Add -harness option to llvm-jitlink.

The -harness option enables new testing use-cases for llvm-jitlink. It takes a
list of objects to treat as a test harness for any regular objects passed to
llvm-jitlink.

If any files are passed using the -harness option then the following
transformations are applied to all other files:

  (1) Symbols definitions that are referenced by the harness files are promoted
      to default scope. (This enables access to statics from test harness).

  (2) Symbols definitions that clash with definitions in the harness files are
      deleted. (This enables interposition by test harness).

  (3) All other definitions in regular files are demoted to local scope.
      (This causes untested code to be dead stripped, reducing memory cost and
      eliminating spurious unresolved symbol errors from untested code).

These transformations allow the harness files to reference and interpose
symbols in the regular object files, which can be used to support execution
tests (including fuzz tests) of functions in relocatable objects produced by a
build.

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

@@ -395,6 +395,10 @@ class Symbol {
     return Name;
   }
 
+  /// Rename this symbol. The client is responsible for updating scope and
+  /// linkage if this name-change requires it.
+  void setName(StringRef Name) { this->Name = Name; }
+
   /// Returns true if this Symbol has content (potentially) defined within this
   /// object file (i.e. is anything but an external or absolute symbol).
   bool isDefined() const {

llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp

@@ -24,7 +24,10 @@ JITLinkerBase::~JITLinkerBase() {}
 
 void JITLinkerBase::linkPhase1(std::unique_ptr<JITLinkerBase> Self) {
 
-  LLVM_DEBUG({ dbgs() << "Building jitlink graph for new input...\n"; });
+  LLVM_DEBUG({
+    dbgs() << "Building jitlink graph for new input "
+           << Ctx->getObjectBuffer().getBufferIdentifier() << "...\n";
+  });
 
   // Build the link graph.
   if (auto GraphOrErr = buildGraph(Ctx->getObjectBuffer()))
@@ -447,16 +450,19 @@ void prune(LinkGraph &G) {
     VisitedBlocks.insert(&B);
 
     for (auto &E : Sym->getBlock().edges()) {
-      if (E.getTarget().isDefined() && !E.getTarget().isLive()) {
-        E.getTarget().setLive(true);
+      // If the edge target is a defined symbol that is being newly marked live
+      // then add it to the worklist.
+      if (E.getTarget().isDefined() && !E.getTarget().isLive())
         Worklist.push_back(&E.getTarget());
-      }
+
+      // Mark the target live.
+      E.getTarget().setLive(true);
     }
   }
 
-  // Collect all the symbols to remove, then remove them.
+  // Collect all defined symbols to remove, then remove them.
   {
-    LLVM_DEBUG(dbgs() << "Dead-stripping symbols:\n");
+    LLVM_DEBUG(dbgs() << "Dead-stripping defined symbols:\n");
     std::vector<Symbol *> SymbolsToRemove;
     for (auto *Sym : G.defined_symbols())
       if (!Sym->isLive())
@@ -479,6 +485,19 @@ void prune(LinkGraph &G) {
       G.removeBlock(*B);
     }
   }
+
+  // Collect all external symbols to remove, then remove them.
+  {
+    LLVM_DEBUG(dbgs() << "Removing unused external symbols:\n");
+    std::vector<Symbol *> SymbolsToRemove;
+    for (auto *Sym : G.external_symbols())
+      if (!Sym->isLive())
+        SymbolsToRemove.push_back(Sym);
+    for (auto *Sym : SymbolsToRemove) {
+      LLVM_DEBUG(dbgs() << "  " << *Sym << "...\n");
+      G.removeExternalSymbol(*Sym);
+    }
+  }
 }
 
 } // end namespace jitlink

llvm/test/ExecutionEngine/JITLink/X86/Inputs/MachO_test_harness_test.s

@@ -0,0 +1,40 @@
+	.section	__TEXT,__text,regular,pure_instructions
+	.globl	_unused_public_function
+	.p2align	4, 0x90
+_unused_public_function:
+	jmp	_unresolvable_external
+
+	.p2align	4, 0x90
+_unused_private_function:
+	jmp	_unresolvable_external
+
+	.globl	_public_func_to_interpose
+	.p2align	4, 0x90
+_public_func_to_interpose:
+	retq
+
+  .p2align	4, 0x90
+_private_func_to_interpose:
+	retq
+
+	.globl	_public_func_to_test
+	.p2align	4, 0x90
+_public_func_to_test:
+	jmp	_public_func_to_interpose
+
+	.p2align	4, 0x90
+_private_func_to_test:
+	jmp	_private_func_to_interpose
+
+	.section	__DATA,__data
+	.globl	_public_func_to_interpose_as_seen_by_test
+	.p2align	3
+_public_func_to_interpose_as_seen_by_test:
+	.quad	_public_func_to_interpose
+
+	.globl	_private_func_to_interpose_as_seen_by_test
+	.p2align	3
+_private_func_to_interpose_as_seen_by_test:
+	.quad	_private_func_to_interpose
+
+.subsections_via_symbols

llvm/test/ExecutionEngine/JITLink/X86/MachO_test_harness_harness.s

@@ -0,0 +1,65 @@
+# RUN: rm -rf %t && mkdir -p %t
+# RUN: llvm-mc -triple=x86_64-apple-macosx10.9 -filetype=obj \
+# RUN:   -o %t/file_to_test.o %S/Inputs/MachO_test_harness_test.s
+# RUN: llvm-mc -triple=x86_64-apple-macosx10.9 -filetype=obj \
+# RUN:   -o %t/test_harness.o %s
+# RUN: llvm-jitlink -noexec -check %s %t/file_to_test.o \
+# RUN:    -harness %t/test_harness.o
+#
+# Check that we
+#   (1) Can call global symbols in the test object.
+#   (2) Can call private symbols in the test object.
+#   (3) Can interpose global symbols in the test object.
+#   (4) Can interpose private symbols in the test object.
+#   (5) Don't need to resolve unused externals in the test object.
+
+.section	__TEXT,__text,regular,pure_instructions
+
+  .globl	_public_func_to_interpose
+	.p2align	4, 0x90
+_public_func_to_interpose:
+	retq
+
+	.globl	_private_func_to_interpose
+	.p2align	4, 0x90
+_private_func_to_interpose:
+	retq
+
+	.globl	_main
+	.p2align	4, 0x90
+_main:
+	callq	_public_func_to_test
+	callq	_private_func_to_test
+  xorl  %eax, %eax
+	retq
+
+	.section	__DATA,__data
+
+# Check that the harness and test file agree on the address of the addresses
+# of the interposes:
+
+# jitlink-check: *{8}_public_func_to_interpose_as_seen_by_harness = \
+# jitlink-check:   *{8}_public_func_to_interpose_as_seen_by_test
+
+# jitlink-check: *{8}_private_func_to_interpose_as_seen_by_harness = \
+# jitlink-check:   *{8}_private_func_to_interpose_as_seen_by_test
+
+  .globl	_public_func_to_interpose_as_seen_by_harness
+	.p2align	3
+_public_func_to_interpose_as_seen_by_harness:
+	.quad	_public_func_to_interpose
+
+	.globl	_private_func_to_interpose_as_seen_by_harness
+	.p2align	3
+_private_func_to_interpose_as_seen_by_harness:
+	.quad	_private_func_to_interpose
+
+# We need to reference the *_as_seen_by_test pointers used above to ensure
+# that they're not dead-stripped as unused.
+  .globl  _anchor_test_case_pointers
+  .p2align  3
+_anchor_test_case_pointers:
+  .quad _public_func_to_interpose_as_seen_by_test
+  .quad _private_func_to_interpose_as_seen_by_test
+
+.subsections_via_symbols

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

@@ -86,6 +86,11 @@ static cl::list<std::string> AbsoluteDefs(
     cl::desc("Inject absolute symbol definitions (syntax: <name>=<addr>)"),
     cl::ZeroOrMore);
 
+static cl::list<std::string> TestHarnesses("harness", cl::Positional,
+                                           cl::desc("Test harness files"),
+                                           cl::ZeroOrMore,
+                                           cl::PositionalEatsArgs);
+
 static cl::opt<bool> ShowInitialExecutionSessionState(
     "show-init-es",
     cl::desc("Print ExecutionSession state before resolving entry point"),
@@ -166,6 +171,43 @@ operator<<(raw_ostream &OS, const Session::FileInfoMap &FIM) {
   return OS;
 }
 
+static Error applyHarnessPromotions(Session &S, LinkGraph &G) {
+
+  // If this graph is part of the test harness there's nothing to do.
+  if (S.HarnessFiles.empty() || S.HarnessFiles.count(G.getName()))
+    return Error::success();
+
+  LLVM_DEBUG(dbgs() << "Appling promotions to graph " << G.getName() << "\n");
+
+  // If it isn't then promote any symbols referenced by the harness to default
+  // scope, remove all symbols that clash with harness definitions, and demote
+  // all others.
+  std::vector<Symbol *> DefinitionsToRemove;
+  for (auto *Sym : G.defined_symbols()) {
+
+    if (!Sym->hasName())
+      continue;
+
+    if (S.HarnessExternals.count(Sym->getName())) {
+      LLVM_DEBUG(dbgs() << "  Promoting " << Sym->getName() << "\n");
+      Sym->setScope(Scope::Default);
+      Sym->setLive(true);
+    } else if (S.HarnessDefinitions.count(Sym->getName())) {
+      LLVM_DEBUG(dbgs() << "  Externalizing " << Sym->getName() << "\n");
+      DefinitionsToRemove.push_back(Sym);
+    } else {
+      LLVM_DEBUG(dbgs() << "  Demoting " << Sym->getName() << "\n");
+      Sym->setScope(Scope::Local);
+      Sym->setLive(false);
+    }
+  }
+
+  for (auto *Sym : DefinitionsToRemove)
+    G.makeExternal(*Sym);
+
+  return Error::success();
+}
+
 static uint64_t computeTotalBlockSizes(LinkGraph &G) {
   uint64_t TotalSize = 0;
   for (auto *B : G.blocks())
@@ -408,12 +450,85 @@ Expected<uint64_t> getSlabAllocSize(StringRef SizeString) {
   return SlabSize * Units;
 }
 
-static std::unique_ptr<jitlink::JITLinkMemoryManager> createMemoryManager() {
+static std::unique_ptr<JITLinkMemoryManager> createMemoryManager() {
   if (!SlabAllocateSizeString.empty()) {
     auto SlabSize = ExitOnErr(getSlabAllocSize(SlabAllocateSizeString));
     return ExitOnErr(JITLinkSlabAllocator::Create(SlabSize));
   }
-  return std::make_unique<jitlink::InProcessMemoryManager>();
+  return std::make_unique<InProcessMemoryManager>();
+}
+
+LLVMJITLinkObjectLinkingLayer::LLVMJITLinkObjectLinkingLayer(
+    Session &S, std::unique_ptr<JITLinkMemoryManager> MemMgr)
+    : ObjectLinkingLayer(S.ES, std::move(MemMgr)), S(S) {}
+
+Error LLVMJITLinkObjectLinkingLayer::add(JITDylib &JD,
+                                         std::unique_ptr<MemoryBuffer> O,
+                                         VModuleKey K) {
+
+  if (S.HarnessFiles.empty() || S.HarnessFiles.count(O->getBufferIdentifier()))
+    return ObjectLinkingLayer::add(JD, std::move(O), std::move(K));
+
+  // Use getObjectSymbolInfo to compute the init symbol, but ignore
+  // the symbols field. We'll handle that manually to include promotion.
+  auto ObjSymInfo =
+      getObjectSymbolInfo(getExecutionSession(), O->getMemBufferRef());
+
+  if (!ObjSymInfo)
+    return ObjSymInfo.takeError();
+
+  auto &InitSymbol = ObjSymInfo->second;
+
+  // If creating an object file was going to fail it would have happened above,
+  // so we can 'cantFail' this.
+  auto Obj =
+      cantFail(object::ObjectFile::createObjectFile(O->getMemBufferRef()));
+
+  SymbolFlagsMap SymbolFlags;
+
+  // The init symbol must be included in the SymbolFlags map if present.
+  if (InitSymbol)
+    SymbolFlags[InitSymbol] = JITSymbolFlags::MaterializationSideEffectsOnly;
+
+  for (auto &Sym : Obj->symbols()) {
+    Expected<uint32_t> SymFlagsOrErr = Sym.getFlags();
+    if (!SymFlagsOrErr)
+      // TODO: Test this error.
+      return SymFlagsOrErr.takeError();
+
+    // Skip symbols not defined in this object file.
+    if (*SymFlagsOrErr & object::BasicSymbolRef::SF_Undefined)
+      continue;
+
+    auto Name = Sym.getName();
+    if (!Name)
+      return Name.takeError();
+
+    // Skip symbols that aren't in the HarnessExternals set.
+    if (!S.HarnessExternals.count(*Name))
+      continue;
+
+    // Skip symbols that have type SF_File.
+    if (auto SymType = Sym.getType()) {
+      if (*SymType == object::SymbolRef::ST_File)
+        continue;
+    } else
+      return SymType.takeError();
+
+    auto InternedName = S.ES.intern(*Name);
+    auto SymFlags = JITSymbolFlags::fromObjectSymbol(Sym);
+    if (!SymFlags)
+      return SymFlags.takeError();
+
+    *SymFlags |= JITSymbolFlags::Exported;
+
+    SymbolFlags[InternedName] = std::move(*SymFlags);
+  }
+
+  auto MU = std::make_unique<BasicObjectLayerMaterializationUnit>(
+      *this, K, std::move(O), std::move(SymbolFlags), std::move(InitSymbol));
+
+  return JD.define(std::move(MU));
 }
 
 Expected<std::unique_ptr<Session>> Session::Create(Triple TT) {
@@ -427,7 +542,7 @@ Expected<std::unique_ptr<Session>> Session::Create(Triple TT) {
 // FIXME: Move to createJITDylib if/when we start using Platform support in
 // llvm-jitlink.
 Session::Session(Triple TT, Error &Err)
-    : ObjLayer(ES, createMemoryManager()), TT(std::move(TT)) {
+    : ObjLayer(*this, createMemoryManager()), TT(std::move(TT)) {
 
   /// Local ObjectLinkingLayer::Plugin class to forward modifyPassConfig to the
   /// Session.
@@ -457,6 +572,39 @@ Session::Session(Triple TT, Error &Err)
         InProcessEHFrameRegistrar::getInstance()));
 
   ObjLayer.addPlugin(std::make_unique<JITLinkSessionPlugin>(*this));
+
+  // Process any harness files.
+  for (auto &HarnessFile : TestHarnesses) {
+    HarnessFiles.insert(HarnessFile);
+
+    auto ObjBuffer =
+        ExitOnErr(errorOrToExpected(MemoryBuffer::getFile(HarnessFile)));
+
+    auto ObjSymbolInfo =
+        ExitOnErr(getObjectSymbolInfo(ES, ObjBuffer->getMemBufferRef()));
+
+    for (auto &KV : ObjSymbolInfo.first)
+      HarnessDefinitions.insert(*KV.first);
+
+    auto Obj = ExitOnErr(
+        object::ObjectFile::createObjectFile(ObjBuffer->getMemBufferRef()));
+
+    for (auto &Sym : Obj->symbols()) {
+      uint32_t SymFlags = ExitOnErr(Sym.getFlags());
+      auto Name = ExitOnErr(Sym.getName());
+
+      if (Name.empty())
+        continue;
+
+      if (SymFlags & object::BasicSymbolRef::SF_Undefined)
+        HarnessExternals.insert(Name);
+    }
+  }
+
+  // If a name is defined by some harness file then it's a definition, not an
+  // external.
+  for (auto &DefName : HarnessDefinitions)
+    HarnessExternals.erase(DefName.getKey());
 }
 
 void Session::dumpSessionInfo(raw_ostream &OS) {
@@ -481,11 +629,14 @@ void Session::modifyPassConfig(const Triple &FTT,
 
   if (ShowLinkGraph)
     PassConfig.PostFixupPasses.push_back([](LinkGraph &G) -> Error {
-      outs() << "Link graph post-fixup:\n";
+      outs() << "Link graph \"" << G.getName() << "\" post-fixup:\n";
       G.dump(outs());
       return Error::success();
     });
 
+  PassConfig.PrePrunePasses.push_back(
+      [this](LinkGraph &G) { return applyHarnessPromotions(*this, G); });
+
   if (ShowSizes) {
     PassConfig.PrePrunePasses.push_back([this](LinkGraph &G) -> Error {
       SizeBeforePruning += computeTotalBlockSizes(G);
@@ -672,6 +823,14 @@ Error loadObjects(Session &S) {
     }
   }
 
+  LLVM_DEBUG(dbgs() << "Adding test harness objects...\n");
+  for (auto HarnessFile : TestHarnesses) {
+    LLVM_DEBUG(dbgs() << "  " << HarnessFile << "\n");
+    auto ObjBuffer =
+        ExitOnErr(errorOrToExpected(MemoryBuffer::getFile(HarnessFile)));
+    ExitOnErr(S.ObjLayer.add(*S.MainJD, std::move(ObjBuffer)));
+  }
+
   // Load each object into the corresponding JITDylib..
   LLVM_DEBUG(dbgs() << "Adding objects...\n");
   for (auto InputFileItr = InputFiles.begin(), InputFileEnd = InputFiles.end();