1 files changed, 40 insertions, 50 deletions

diff --git a/unittests/ExecutionEngine/Orc/OrcTestCommon.h b/unittests/ExecutionEngine/Orc/OrcTestCommon.h
index 6c6b4918c205..28a5f08ee439 100644
--- a/unittests/ExecutionEngine/Orc/OrcTestCommon.h
+++ b/unittests/ExecutionEngine/Orc/OrcTestCommon.h
@@ -105,73 +105,63 @@ public:
   }
 };
 
-template <typename HandleT,
-          typename AddModuleFtor,
-          typename RemoveModuleFtor,
-          typename FindSymbolFtor,
-          typename FindSymbolInFtor>
+template <typename HandleT, typename ModuleT>
 class MockBaseLayer {
 public:
 
-  typedef HandleT ModuleHandleT;
-
-  MockBaseLayer(AddModuleFtor &&AddModule,
-                RemoveModuleFtor &&RemoveModule,
-                FindSymbolFtor &&FindSymbol,
-                FindSymbolInFtor &&FindSymbolIn)
-      : AddModule(std::move(AddModule)),
-        RemoveModule(std::move(RemoveModule)),
-        FindSymbol(std::move(FindSymbol)),
-        FindSymbolIn(std::move(FindSymbolIn))
-  {}
-
-  template <typename ModuleT, typename MemoryManagerPtrT,
-            typename SymbolResolverPtrT>
-  Expected<ModuleHandleT> addModule(ModuleT Ms, MemoryManagerPtrT MemMgr,
-                                    SymbolResolverPtrT Resolver) {
-    return AddModule(std::move(Ms), std::move(MemMgr), std::move(Resolver));
+  using ModuleHandleT = HandleT;
+
+  using AddModuleSignature =
+    Expected<ModuleHandleT>(ModuleT M,
+                            std::shared_ptr<JITSymbolResolver> R);
+
+  using RemoveModuleSignature = Error(ModuleHandleT H);
+  using FindSymbolSignature = JITSymbol(const std::string &Name,
+                                        bool ExportedSymbolsOnly);
+  using FindSymbolInSignature = JITSymbol(ModuleHandleT H,
+                                          const std::string &Name,
+                                          bool ExportedSymbolsONly);
+  using EmitAndFinalizeSignature = Error(ModuleHandleT H);
+
+  std::function<AddModuleSignature> addModuleImpl;
+  std::function<RemoveModuleSignature> removeModuleImpl;
+  std::function<FindSymbolSignature> findSymbolImpl;
+  std::function<FindSymbolInSignature> findSymbolInImpl;
+  std::function<EmitAndFinalizeSignature> emitAndFinalizeImpl;
+
+  Expected<ModuleHandleT> addModule(ModuleT M,
+                                    std::shared_ptr<JITSymbolResolver> R) {
+    assert(addModuleImpl &&
+           "addModule called, but no mock implementation was provided");
+    return addModuleImpl(std::move(M), std::move(R));
   }
 
   Error removeModule(ModuleHandleT H) {
-    return RemoveModule(H);
+    assert(removeModuleImpl &&
+           "removeModule called, but no mock implementation was provided");
+    return removeModuleImpl(H);
   }
 
   JITSymbol findSymbol(const std::string &Name, bool ExportedSymbolsOnly) {
-    return FindSymbol(Name, ExportedSymbolsOnly);
+    assert(findSymbolImpl &&
+           "findSymbol called, but no mock implementation was provided");
+    return findSymbolImpl(Name, ExportedSymbolsOnly);
   }
 
   JITSymbol findSymbolIn(ModuleHandleT H, const std::string &Name,
                          bool ExportedSymbolsOnly) {
-    return FindSymbolIn(H, Name, ExportedSymbolsOnly);
+    assert(findSymbolInImpl &&
+           "findSymbolIn called, but no mock implementation was provided");
+    return findSymbolInImpl(H, Name, ExportedSymbolsOnly);
   }
 
-private:
-  AddModuleFtor AddModule;
-  RemoveModuleFtor RemoveModule;
-  FindSymbolFtor FindSymbol;
-  FindSymbolInFtor FindSymbolIn;
+  Error emitAndFinaliez(ModuleHandleT H) {
+    assert(emitAndFinalizeImpl &&
+           "emitAndFinalize called, but no mock implementation was provided");
+    return emitAndFinalizeImpl(H);
+  }
 };
 
-template <typename ModuleHandleT,
-          typename AddModuleFtor,
-          typename RemoveModuleFtor,
-          typename FindSymbolFtor,
-          typename FindSymbolInFtor>
-MockBaseLayer<ModuleHandleT, AddModuleFtor, RemoveModuleFtor,
-              FindSymbolFtor, FindSymbolInFtor>
-createMockBaseLayer(AddModuleFtor &&AddModule,
-                    RemoveModuleFtor &&RemoveModule,
-                    FindSymbolFtor &&FindSymbol,
-                    FindSymbolInFtor &&FindSymbolIn) {
-  return MockBaseLayer<ModuleHandleT, AddModuleFtor, RemoveModuleFtor,
-                       FindSymbolFtor, FindSymbolInFtor>(
-                         std::forward<AddModuleFtor>(AddModule),
-                         std::forward<RemoveModuleFtor>(RemoveModule),
-                         std::forward<FindSymbolFtor>(FindSymbol),
-                         std::forward<FindSymbolInFtor>(FindSymbolIn));
-}
-
-
 class ReturnNullJITSymbol {
 public:
   template <typename... Args>