[llvm] e12a028 - [llvm-jitlink] Support promotion of ODR weak symbols in -harness mode.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2020-08-01T18:33:44-07:00
New Revision: e12a028ed3ed79fba1e1aa5212c7e22d83949192

URL: https://github.com/llvm/llvm-project/commit/e12a028ed3ed79fba1e1aa5212c7e22d83949192
DIFF: https://github.com/llvm/llvm-project/commit/e12a028ed3ed79fba1e1aa5212c7e22d83949192.diff

LOG: [llvm-jitlink] Support promotion of ODR weak symbols in -harness mode.

This prevents weak symbols from being immediately dead-stripped when not
directly referenced from the test harneess, enabling use of weak symbols
from the code under test.

Added: 
    

Modified: 
    llvm/test/ExecutionEngine/JITLink/X86/Inputs/MachO_test_harness_test.s
    llvm/tools/llvm-jitlink/llvm-jitlink.cpp
    llvm/tools/llvm-jitlink/llvm-jitlink.h

Removed: 
    


################################################################################
diff  --git a/llvm/test/ExecutionEngine/JITLink/X86/Inputs/MachO_test_harness_test.s b/llvm/test/ExecutionEngine/JITLink/X86/Inputs/MachO_test_harness_test.s
index e0764a93fa43..3ca616a4a729 100644
--- a/llvm/test/ExecutionEngine/JITLink/X86/Inputs/MachO_test_harness_test.s
+++ b/llvm/test/ExecutionEngine/JITLink/X86/Inputs/MachO_test_harness_test.s
@@ -17,9 +17,16 @@ _public_func_to_interpose:
 _private_func_to_interpose:
 	retq
 
+	.globl	_used_weak
+	.weak_definition	_used_weak
+	.p2align	4, 0x90
+_used_weak:
+	retq
+
 	.globl	_public_func_to_test
 	.p2align	4, 0x90
 _public_func_to_test:
+  callq _used_weak
 	jmp	_public_func_to_interpose
 
 	.p2align	4, 0x90

diff  --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index 6828944ced23..22d29c12e192 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -188,10 +188,28 @@ static Error applyHarnessPromotions(Session &S, LinkGraph &G) {
     if (!Sym->hasName())
       continue;
 
-    if (S.HarnessExternals.count(Sym->getName())) {
+    if (Sym->getLinkage() == Linkage::Weak) {
+      if (!S.CanonicalWeakDefs.count(Sym->getName()) ||
+          S.CanonicalWeakDefs[Sym->getName()] != G.getName()) {
+        LLVM_DEBUG({
+          dbgs() << "  Externalizing weak symbol " << Sym->getName() << "\n";
+        });
+        DefinitionsToRemove.push_back(Sym);
+      } else {
+        LLVM_DEBUG({
+          dbgs() << "  Making weak symbol " << Sym->getName() << " strong\n";
+        });
+        if (S.HarnessExternals.count(Sym->getName()))
+          Sym->setScope(Scope::Default);
+        else
+          Sym->setScope(Scope::Hidden);
+        Sym->setLinkage(Linkage::Strong);
+      }
+    } else if (S.HarnessExternals.count(Sym->getName())) {
       LLVM_DEBUG(dbgs() << "  Promoting " << Sym->getName() << "\n");
       Sym->setScope(Scope::Default);
       Sym->setLive(true);
+      continue;
     } else if (S.HarnessDefinitions.count(Sym->getName())) {
       LLVM_DEBUG(dbgs() << "  Externalizing " << Sym->getName() << "\n");
       DefinitionsToRemove.push_back(Sym);
@@ -504,10 +522,6 @@ Error LLVMJITLinkObjectLinkingLayer::add(JITDylib &JD,
     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)
@@ -515,13 +529,28 @@ Error LLVMJITLinkObjectLinkingLayer::add(JITDylib &JD,
     } else
       return SymType.takeError();
 
-    auto InternedName = S.ES.intern(*Name);
     auto SymFlags = JITSymbolFlags::fromObjectSymbol(Sym);
     if (!SymFlags)
       return SymFlags.takeError();
 
-    *SymFlags |= JITSymbolFlags::Exported;
+    if (SymFlags->isWeak()) {
+      // If this is a weak symbol that's not defined in the harness then we
+      // need to either mark it as strong (if this is the first definition
+      // that we've seen) or discard it.
+      if (S.HarnessDefinitions.count(*Name) || S.CanonicalWeakDefs.count(*Name))
+        continue;
+      S.CanonicalWeakDefs[*Name] = O->getBufferIdentifier();
+      *SymFlags &= ~JITSymbolFlags::Weak;
+      if (!S.HarnessExternals.count(*Name))
+        *SymFlags &= ~JITSymbolFlags::Exported;
+    } else if (S.HarnessExternals.count(*Name)) {
+      *SymFlags |= JITSymbolFlags::Exported;
+    } else {
+      // Skip symbols that aren't in the HarnessExternals set.
+      continue;
+    }
 
+    auto InternedName = S.ES.intern(*Name);
     SymbolFlags[InternedName] = std::move(*SymFlags);
   }
 

diff  --git a/llvm/tools/llvm-jitlink/llvm-jitlink.h b/llvm/tools/llvm-jitlink/llvm-jitlink.h
index c16aed9f2b50..227d3a9225c8 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.h
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.h
@@ -86,6 +86,7 @@ struct Session {
   StringSet<> HarnessFiles;
   StringSet<> HarnessExternals;
   StringSet<> HarnessDefinitions;
+  DenseMap<StringRef, StringRef> CanonicalWeakDefs;
 
 private:
   Session(Triple TT, Error &Err);