1 files changed, 37 insertions, 29 deletions

diff --git a/llvm/unittests/tools/llvm-exegesis/X86/InMemoryAssemblerTest.cpp b/llvm/unittests/tools/llvm-exegesis/X86/InMemoryAssemblerTest.cpp
index d00b223393f..24b55c8f109 100644
--- a/llvm/unittests/tools/llvm-exegesis/X86/InMemoryAssemblerTest.cpp
+++ b/llvm/unittests/tools/llvm-exegesis/X86/InMemoryAssemblerTest.cpp
@@ -7,7 +7,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "InMemoryAssembler.h"
+#include "Assembler.h"
 #include "X86InstrInfo.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
@@ -51,8 +51,8 @@ protected:
         llvm::TargetRegistry::lookupTarget(TT, Error);
     EXPECT_TRUE(TheTarget) << Error << " " << TT;
     const llvm::TargetOptions Options;
-    llvm::TargetMachine* TM = TheTarget->createTargetMachine(
-            TT, CpuName, "", Options, llvm::Reloc::Model::Static);
+    llvm::TargetMachine *TM = TheTarget->createTargetMachine(
+        TT, CpuName, "", Options, llvm::Reloc::Model::Static);
     EXPECT_TRUE(TM) << TT << " " << CpuName;
     return std::unique_ptr<llvm::LLVMTargetMachine>(
         static_cast<llvm::LLVMTargetMachine *>(TM));
@@ -62,58 +62,66 @@ protected:
     return llvm::StringRef(TT).startswith_lower("x86_64");
   }
 
+  ExecutableFunction assembleFunction(llvm::MCInst MCInst) {
+    return assembleToFunction({MCInst});
+  }
+
+  ExecutableFunction assembleEmptyFunction() { return assembleToFunction({}); }
+
 private:
+  ExecutableFunction
+  assembleToFunction(llvm::ArrayRef<llvm::MCInst> Instructions) {
+    llvm::SmallString<256> Buffer;
+    llvm::raw_svector_ostream AsmStream(Buffer);
+    assembleToStream(createTargetMachine(), Instructions, AsmStream);
+    return ExecutableFunction(createTargetMachine(),
+                              getObjectFromBuffer(AsmStream.str()));
+  }
   const std::string TT;
   const std::string CpuName;
 };
 
 // Used to skip tests on unsupported architectures and operating systems.
 // To skip a test, add this macro at the top of a test-case.
-#define SKIP_UNSUPPORTED_PLATFORM \
-  do \
-    if (!IsSupportedTarget()) \
-      return; \
-  while(0)
-
+#define SKIP_UNSUPPORTED_PLATFORM                                              \
+  do                                                                           \
+    if (!IsSupportedTarget())                                                  \
+      return;                                                                  \
+  while (0)
 
-TEST_F(MachineFunctionGeneratorTest, DISABLED_JitFunction) {
+TEST_F(MachineFunctionGeneratorTest, JitFunction) {
   SKIP_UNSUPPORTED_PLATFORM;
-  JitFunctionContext Context(createTargetMachine());
-  JitFunction Function(std::move(Context), {});
+  const ExecutableFunction Function = assembleEmptyFunction();
   ASSERT_THAT(Function.getFunctionBytes().str(), ElementsAre(0xc3));
   // FIXME: Check that the function runs without errors. Right now this is
   // disabled because it fails on some bots.
-  // Function();
+  Function();
 }
 
-TEST_F(MachineFunctionGeneratorTest, DISABLED_JitFunctionXOR32rr) {
+TEST_F(MachineFunctionGeneratorTest, JitFunctionXOR32rr) {
   SKIP_UNSUPPORTED_PLATFORM;
-  JitFunctionContext Context(createTargetMachine());
-  JitFunction Function(
-      std::move(Context),
-      {MCInstBuilder(XOR32rr).addReg(EAX).addReg(EAX).addReg(EAX)});
+  const ExecutableFunction Function = assembleFunction(
+      MCInstBuilder(XOR32rr).addReg(EAX).addReg(EAX).addReg(EAX));
   ASSERT_THAT(Function.getFunctionBytes().str(), ElementsAre(0x31, 0xc0, 0xc3));
-  // Function();
+  Function();
 }
 
-TEST_F(MachineFunctionGeneratorTest, DISABLED_JitFunctionMOV64ri) {
+TEST_F(MachineFunctionGeneratorTest, JitFunctionMOV64ri) {
   SKIP_UNSUPPORTED_PLATFORM;
-  JitFunctionContext Context(createTargetMachine());
-  JitFunction Function(std::move(Context),
-                       {MCInstBuilder(MOV64ri32).addReg(RAX).addImm(42)});
+  const ExecutableFunction Function =
+      assembleFunction(MCInstBuilder(MOV64ri32).addReg(RAX).addImm(42));
   ASSERT_THAT(Function.getFunctionBytes().str(),
               ElementsAre(0x48, 0xc7, 0xc0, 0x2a, 0x00, 0x00, 0x00, 0xc3));
-  // Function();
+  Function();
 }
 
-TEST_F(MachineFunctionGeneratorTest, DISABLED_JitFunctionMOV32ri) {
+TEST_F(MachineFunctionGeneratorTest, JitFunctionMOV32ri) {
   SKIP_UNSUPPORTED_PLATFORM;
-  JitFunctionContext Context(createTargetMachine());
-  JitFunction Function(std::move(Context),
-                       {MCInstBuilder(MOV32ri).addReg(EAX).addImm(42)});
+  const ExecutableFunction Function =
+      assembleFunction(MCInstBuilder(MOV32ri).addReg(EAX).addImm(42));
   ASSERT_THAT(Function.getFunctionBytes().str(),
               ElementsAre(0xb8, 0x2a, 0x00, 0x00, 0x00, 0xc3));
-  // Function();
+  Function();
 }
 
 } // namespace