1 files changed, 43 insertions, 5 deletions

diff --git a/unittests/ExecutionEngine/ExecutionEngineTest.cpp b/unittests/ExecutionEngine/ExecutionEngineTest.cpp
index f23745c..19917a4 100644
--- a/unittests/ExecutionEngine/ExecutionEngineTest.cpp
+++ b/unittests/ExecutionEngine/ExecutionEngineTest.cpp
@@ -8,10 +8,13 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Interpreter.h"
+#include "llvm/ExecutionEngine/RTDyldMemoryManager.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
+#include "llvm/Support/DynamicLibrary.h"
+#include "llvm/Support/ManagedStatic.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
@@ -19,10 +22,14 @@ using namespace llvm;
 namespace {
 
 class ExecutionEngineTest : public testing::Test {
+private:
+  llvm_shutdown_obj Y; // Call llvm_shutdown() on exit.
+
 protected:
-  ExecutionEngineTest()
-    : M(new Module("<main>", getGlobalContext())), Error(""),
-      Engine(EngineBuilder(M).setErrorStr(&Error).create()) {
+  ExecutionEngineTest() {
+    auto Owner = make_unique<Module>("<main>", getGlobalContext());
+    M = Owner.get();
+    Engine.reset(EngineBuilder(std::move(Owner)).setErrorStr(&Error).create());
   }
 
   virtual void SetUp() {
@@ -35,9 +42,9 @@ protected:
                               GlobalValue::ExternalLinkage, nullptr, Name);
   }
 
-  Module *const M;
   std::string Error;
-  const std::unique_ptr<ExecutionEngine> Engine;
+  Module *M;  // Owned by ExecutionEngine.
+  std::unique_ptr<ExecutionEngine> Engine;
 };
 
 TEST_F(ExecutionEngineTest, ForwardGlobalMapping) {
@@ -127,4 +134,35 @@ TEST_F(ExecutionEngineTest, DestructionRemovesGlobalMapping) {
   EXPECT_EQ(nullptr, Engine->getGlobalValueAtAddress(&Mem1));
 }
 
+TEST_F(ExecutionEngineTest, LookupWithMangledName) {
+  int x;
+  llvm::sys::DynamicLibrary::AddSymbol("x", &x);
+
+  // Demonstrate that getSymbolAddress accepts mangled names and always strips
+  // the leading underscore.
+  EXPECT_EQ(reinterpret_cast<uint64_t>(&x),
+            RTDyldMemoryManager::getSymbolAddressInProcess("_x"));
+}
+
+TEST_F(ExecutionEngineTest, LookupWithMangledAndDemangledSymbol) {
+  int x;
+  int _x;
+  llvm::sys::DynamicLibrary::AddSymbol("x", &x);
+  llvm::sys::DynamicLibrary::AddSymbol("_x", &_x);
+
+  // Lookup the demangled name first, even if there's a demangled symbol that
+  // matches the input already.
+  EXPECT_EQ(reinterpret_cast<uint64_t>(&x),
+            RTDyldMemoryManager::getSymbolAddressInProcess("_x"));
+}
+
+TEST_F(ExecutionEngineTest, LookupwithDemangledName) {
+  int _x;
+  llvm::sys::DynamicLibrary::AddSymbol("_x", &_x);
+
+  // But do fallback to looking up a demangled name if there's no ambiguity
+  EXPECT_EQ(reinterpret_cast<uint64_t>(&_x),
+            RTDyldMemoryManager::getSymbolAddressInProcess("_x"));
+}
+
 }