[MCJIT] Add a FindGlobalVariableNamed utility

Summary: This adds FindGlobalVariableNamed to ExecutionEngine
(plus implementation in MCJIT), which is an analog of
FindFunctionNamed for GlobalVariables.

Reviewers: lhames

Reviewed By: lhames

Subscribers: llvm-commits

Differential Revision: http://reviews.llvm.org/D10421

llvm-svn: 240202

5 files changed, 63 insertions, 6 deletions

diff --git a/llvm/include/llvm/ExecutionEngine/ExecutionEngine.h b/llvm/include/llvm/ExecutionEngine/ExecutionEngine.h
index f4e5d3850ca..f86490b55bd 100644
--- a/llvm/include/llvm/ExecutionEngine/ExecutionEngine.h
+++ b/llvm/include/llvm/ExecutionEngine/ExecutionEngine.h
@@ -197,11 +197,16 @@ public:
   /// M is found.
   virtual bool removeModule(Module *M);
 
-  /// FindFunctionNamed - Search all of the active modules to find the one that
+  /// FindFunctionNamed - Search all of the active modules to find the function that
   /// defines FnName.  This is very slow operation and shouldn't be used for
   /// general code.
   virtual Function *FindFunctionNamed(const char *FnName);
 
+  /// FindGlobalVariableNamed - Search all of the active modules to find the global variable
+  /// that defines Name.  This is very slow operation and shouldn't be used for
+  /// general code.
+  virtual GlobalVariable *FindGlobalVariableNamed(const char *Name, bool AllowInternal = false);
+
   /// runFunction - Execute the specified function with the specified arguments,
   /// and return the result.
   virtual GenericValue runFunction(Function *F,
diff --git a/llvm/lib/ExecutionEngine/ExecutionEngine.cpp b/llvm/lib/ExecutionEngine/ExecutionEngine.cpp
index b0d41905aeb..94e809061c7 100644
--- a/llvm/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/llvm/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -153,6 +153,14 @@ Function *ExecutionEngine::FindFunctionNamed(const char *FnName) {
   return nullptr;
 }
 
+GlobalVariable *ExecutionEngine::FindGlobalVariableNamed(const char *Name, bool AllowInternal) {
+  for (unsigned i = 0, e = Modules.size(); i != e; ++i) {
+    GlobalVariable *GV = Modules[i]->getGlobalVariable(Name,AllowInternal);
+    if (GV && !GV->isDeclaration())
+      return GV;
+  }
+  return nullptr;
+}
 
 uint64_t ExecutionEngineState::RemoveMapping(StringRef Name) {
   GlobalAddressMapTy::iterator I = GlobalAddressMap.find(Name);
diff --git a/llvm/lib/ExecutionEngine/MCJIT/MCJIT.cpp b/llvm/lib/ExecutionEngine/MCJIT/MCJIT.cpp
index 358d3649d90..87243e4221f 100644
--- a/llvm/lib/ExecutionEngine/MCJIT/MCJIT.cpp
+++ b/llvm/lib/ExecutionEngine/MCJIT/MCJIT.cpp
@@ -429,6 +429,19 @@ Function *MCJIT::FindFunctionNamedInModulePtrSet(const char *FnName,
   return nullptr;
 }
 
+GlobalVariable *MCJIT::FindGlobalVariableNamedInModulePtrSet(const char *Name,
+                                                             bool AllowInternal,
+                                                             ModulePtrSet::iterator I,
+                                                             ModulePtrSet::iterator E) {
+  for (; I != E; ++I) {
+    GlobalVariable *GV = (*I)->getGlobalVariable(Name, AllowInternal);
+    if (GV && !GV->isDeclaration())
+      return GV;
+  }
+  return nullptr;
+}
+
+
 Function *MCJIT::FindFunctionNamed(const char *FnName) {
   Function *F = FindFunctionNamedInModulePtrSet(
       FnName, OwnedModules.begin_added(), OwnedModules.end_added());
@@ -441,6 +454,18 @@ Function *MCJIT::FindFunctionNamed(const char *FnName) {
   return F;
 }
 
+GlobalVariable *MCJIT::FindGlobalVariableNamed(const char *Name, bool AllowInternal) {
+  GlobalVariable *GV = FindGlobalVariableNamedInModulePtrSet(
+      Name, AllowInternal, OwnedModules.begin_added(), OwnedModules.end_added());
+  if (!GV)
+    GV = FindGlobalVariableNamedInModulePtrSet(Name, AllowInternal, OwnedModules.begin_loaded(),
+                                        OwnedModules.end_loaded());
+  if (!GV)
+    GV = FindGlobalVariableNamedInModulePtrSet(Name, AllowInternal, OwnedModules.begin_finalized(),
+                                        OwnedModules.end_finalized());
+  return GV;
+}
+
 GenericValue MCJIT::runFunction(Function *F, ArrayRef<GenericValue> ArgValues) {
   assert(F && "Function *F was null at entry to run()");
 
diff --git a/llvm/lib/ExecutionEngine/MCJIT/MCJIT.h b/llvm/lib/ExecutionEngine/MCJIT/MCJIT.h
index 20924e55806..7fda1e0fed6 100644
--- a/llvm/lib/ExecutionEngine/MCJIT/MCJIT.h
+++ b/llvm/lib/ExecutionEngine/MCJIT/MCJIT.h
@@ -200,6 +200,11 @@ class MCJIT : public ExecutionEngine {
                                             ModulePtrSet::iterator I,
                                             ModulePtrSet::iterator E);
 
+  GlobalVariable *FindGlobalVariableNamedInModulePtrSet(const char *Name,
+                                                        bool AllowInternal,
+                                                        ModulePtrSet::iterator I,
+                                                        ModulePtrSet::iterator E);
+
   void runStaticConstructorsDestructorsInModulePtrSet(bool isDtors,
                                                       ModulePtrSet::iterator I,
                                                       ModulePtrSet::iterator E);
@@ -215,10 +220,15 @@ public:
   void addArchive(object::OwningBinary<object::Archive> O) override;
   bool removeModule(Module *M) override;
 
-  /// FindFunctionNamed - Search all of the active modules to find the one that
+  /// FindFunctionNamed - Search all of the active modules to find the function that
   /// defines FnName.  This is very slow operation and shouldn't be used for
   /// general code.
-  Function *FindFunctionNamed(const char *FnName) override;
+  virtual Function *FindFunctionNamed(const char *FnName) override;
+
+  /// FindGlobalVariableNamed - Search all of the active modules to find the global variable
+  /// that defines Name.  This is very slow operation and shouldn't be used for
+  /// general code.
+  virtual GlobalVariable *FindGlobalVariableNamed(const char *Name, bool AllowInternal = false) override;
 
   /// Sets the object manager that MCJIT should use to avoid compilation.
   void setObjectCache(ObjectCache *manager) override;
diff --git a/llvm/unittests/ExecutionEngine/MCJIT/MCJITMultipleModuleTest.cpp b/llvm/unittests/ExecutionEngine/MCJIT/MCJITMultipleModuleTest.cpp
index da6e25a3d51..7d52a9acca7 100644
--- a/llvm/unittests/ExecutionEngine/MCJIT/MCJITMultipleModuleTest.cpp
+++ b/llvm/unittests/ExecutionEngine/MCJIT/MCJITMultipleModuleTest.cpp
@@ -194,14 +194,15 @@ TEST_F(MCJITMultipleModuleTest, two_module_consecutive_call_case) {
 
 
 // Module A { Global Variable GVA, Function FA loads GVA },
-// Module B { Global Variable GVB, Function FB loads GVB },
-// execute FB then FA
+// Module B { Global Variable GVB, Internal Global GVC, Function FB loads GVB },
+// execute FB then FA, also check that the global variables are properly accesible
+// through the ExecutionEngine APIs
 TEST_F(MCJITMultipleModuleTest, two_module_global_variables_case) {
   SKIP_UNSUPPORTED_PLATFORM;
 
   std::unique_ptr<Module> A, B;
   Function *FA, *FB;
-  GlobalVariable *GVA, *GVB;
+  GlobalVariable *GVA, *GVB, *GVC;
   A.reset(createEmptyModule("A"));
   B.reset(createEmptyModule("B"));
 
@@ -213,9 +214,17 @@ TEST_F(MCJITMultipleModuleTest, two_module_global_variables_case) {
   FB = startFunction<int32_t(void)>(B.get(), "FB");
   endFunctionWithRet(FB, Builder.CreateLoad(GVB));
 
+  GVC = insertGlobalInt32(B.get(), "GVC", initialNum);
+  GVC->setLinkage(GlobalValue::InternalLinkage);
+
   createJIT(std::move(A));
   TheJIT->addModule(std::move(B));
 
+  EXPECT_EQ(GVA, TheJIT->FindGlobalVariableNamed("GVA"));
+  EXPECT_EQ(GVB, TheJIT->FindGlobalVariableNamed("GVB"));
+  EXPECT_EQ(GVC, TheJIT->FindGlobalVariableNamed("GVC",true));
+  EXPECT_EQ(NULL, TheJIT->FindGlobalVariableNamed("GVC"));
+
   uint64_t FBPtr = TheJIT->getFunctionAddress(FB->getName().str());
   TheJIT->finalizeObject();
   EXPECT_TRUE(0 != FBPtr);