1//===- MCJITObjectCacheTest.cpp - Unit tests for MCJIT object caching -----===//
2//
3//                     The LLVM Compiler Infrastructure
4//
5// This file is distributed under the University of Illinois Open Source
6// License. See LICENSE.TXT for details.
7//
8//===----------------------------------------------------------------------===//
9
10#include "MCJITTestBase.h"
11#include "llvm/ADT/SmallVector.h"
12#include "llvm/ADT/StringMap.h"
13#include "llvm/ADT/StringSet.h"
14#include "llvm/ExecutionEngine/MCJIT.h"
15#include "llvm/ExecutionEngine/ObjectCache.h"
16#include "llvm/ExecutionEngine/SectionMemoryManager.h"
17#include "gtest/gtest.h"
18
19using namespace llvm;
20
21namespace {
22
23class TestObjectCache : public ObjectCache {
24public:
25  TestObjectCache() : DuplicateInserted(false) { }
26
27  void notifyObjectCompiled(const Module *M, MemoryBufferRef Obj) override {
28    // If we've seen this module before, note that.
29    const std::string ModuleID = M->getModuleIdentifier();
30    if (ObjMap.find(ModuleID) != ObjMap.end())
31      DuplicateInserted = true;
32    // Store a copy of the buffer in our map.
33    ObjMap[ModuleID] = copyBuffer(Obj);
34  }
35
36  std::unique_ptr<MemoryBuffer> getObject(const Module *M) override {
37    const MemoryBuffer* BufferFound = getObjectInternal(M);
38    ModulesLookedUp.insert(M->getModuleIdentifier());
39    if (!BufferFound)
40      return nullptr;
41    // Our test cache wants to maintain ownership of its object buffers
42    // so we make a copy here for the execution engine.
43    return MemoryBuffer::getMemBufferCopy(BufferFound->getBuffer());
44  }
45
46  // Test-harness-specific functions
47  bool wereDuplicatesInserted() { return DuplicateInserted; }
48
49  bool wasModuleLookedUp(const Module *M) {
50    return ModulesLookedUp.find(M->getModuleIdentifier())
51                                      != ModulesLookedUp.end();
52  }
53
54  const MemoryBuffer* getObjectInternal(const Module* M) {
55    // Look for the module in our map.
56    const std::string ModuleID = M->getModuleIdentifier();
57    StringMap<const MemoryBuffer *>::iterator it = ObjMap.find(ModuleID);
58    if (it == ObjMap.end())
59      return nullptr;
60    return it->second;
61  }
62
63private:
64  MemoryBuffer *copyBuffer(MemoryBufferRef Buf) {
65    // Create a local copy of the buffer.
66    std::unique_ptr<MemoryBuffer> NewBuffer =
67        MemoryBuffer::getMemBufferCopy(Buf.getBuffer());
68    MemoryBuffer *Ret = NewBuffer.get();
69    AllocatedBuffers.push_back(std::move(NewBuffer));
70    return Ret;
71  }
72
73  StringMap<const MemoryBuffer *> ObjMap;
74  StringSet<>                     ModulesLookedUp;
75  SmallVector<std::unique_ptr<MemoryBuffer>, 2> AllocatedBuffers;
76  bool                            DuplicateInserted;
77};
78
79class MCJITObjectCacheTest : public testing::Test, public MCJITTestBase {
80protected:
81
82  enum {
83    OriginalRC = 6,
84    ReplacementRC = 7
85  };
86
87  void SetUp() override {
88    M.reset(createEmptyModule("<main>"));
89    Main = insertMainFunction(M.get(), OriginalRC);
90  }
91
92  void compileAndRun(int ExpectedRC = OriginalRC) {
93    // This function shouldn't be called until after SetUp.
94    ASSERT_TRUE(bool(TheJIT));
95    ASSERT_TRUE(nullptr != Main);
96
97    // We may be using a null cache, so ensure compilation is valid.
98    TheJIT->finalizeObject();
99    void *vPtr = TheJIT->getPointerToFunction(Main);
100
101    EXPECT_TRUE(nullptr != vPtr)
102      << "Unable to get pointer to main() from JIT";
103
104    int (*FuncPtr)(void) = (int(*)(void))(intptr_t)vPtr;
105    int returnCode = FuncPtr();
106    EXPECT_EQ(returnCode, ExpectedRC);
107  }
108
109  Function *Main;
110};
111
112TEST_F(MCJITObjectCacheTest, SetNullObjectCache) {
113  SKIP_UNSUPPORTED_PLATFORM;
114
115  createJIT(std::move(M));
116
117  TheJIT->setObjectCache(nullptr);
118
119  compileAndRun();
120}
121
122
123TEST_F(MCJITObjectCacheTest, VerifyBasicObjectCaching) {
124  SKIP_UNSUPPORTED_PLATFORM;
125
126  std::unique_ptr<TestObjectCache> Cache(new TestObjectCache);
127
128  // Save a copy of the module pointer before handing it off to MCJIT.
129  const Module * SavedModulePointer = M.get();
130
131  createJIT(std::move(M));
132
133  TheJIT->setObjectCache(Cache.get());
134
135  // Verify that our object cache does not contain the module yet.
136  const MemoryBuffer *ObjBuffer = Cache->getObjectInternal(SavedModulePointer);
137  EXPECT_EQ(nullptr, ObjBuffer);
138
139  compileAndRun();
140
141  // Verify that MCJIT tried to look-up this module in the cache.
142  EXPECT_TRUE(Cache->wasModuleLookedUp(SavedModulePointer));
143
144  // Verify that our object cache now contains the module.
145  ObjBuffer = Cache->getObjectInternal(SavedModulePointer);
146  EXPECT_TRUE(nullptr != ObjBuffer);
147
148  // Verify that the cache was only notified once.
149  EXPECT_FALSE(Cache->wereDuplicatesInserted());
150}
151
152TEST_F(MCJITObjectCacheTest, VerifyLoadFromCache) {
153  SKIP_UNSUPPORTED_PLATFORM;
154
155  std::unique_ptr<TestObjectCache> Cache(new TestObjectCache);
156
157  // Compile this module with an MCJIT engine
158  createJIT(std::move(M));
159  TheJIT->setObjectCache(Cache.get());
160  TheJIT->finalizeObject();
161
162  // Destroy the MCJIT engine we just used
163  TheJIT.reset();
164
165  // Create a new memory manager.
166  MM.reset(new SectionMemoryManager());
167
168  // Create a new module and save it. Use a different return code so we can
169  // tell if MCJIT compiled this module or used the cache.
170  M.reset(createEmptyModule("<main>"));
171  Main = insertMainFunction(M.get(), ReplacementRC);
172  const Module * SecondModulePointer = M.get();
173
174  // Create a new MCJIT instance to load this module then execute it.
175  createJIT(std::move(M));
176  TheJIT->setObjectCache(Cache.get());
177  compileAndRun();
178
179  // Verify that MCJIT tried to look-up this module in the cache.
180  EXPECT_TRUE(Cache->wasModuleLookedUp(SecondModulePointer));
181
182  // Verify that MCJIT didn't try to cache this again.
183  EXPECT_FALSE(Cache->wereDuplicatesInserted());
184}
185
186TEST_F(MCJITObjectCacheTest, VerifyNonLoadFromCache) {
187  SKIP_UNSUPPORTED_PLATFORM;
188
189  std::unique_ptr<TestObjectCache> Cache(new TestObjectCache);
190
191  // Compile this module with an MCJIT engine
192  createJIT(std::move(M));
193  TheJIT->setObjectCache(Cache.get());
194  TheJIT->finalizeObject();
195
196  // Destroy the MCJIT engine we just used
197  TheJIT.reset();
198
199  // Create a new memory manager.
200  MM.reset(new SectionMemoryManager());
201
202  // Create a new module and save it. Use a different return code so we can
203  // tell if MCJIT compiled this module or used the cache. Note that we use
204  // a new module name here so the module shouldn't be found in the cache.
205  M.reset(createEmptyModule("<not-main>"));
206  Main = insertMainFunction(M.get(), ReplacementRC);
207  const Module * SecondModulePointer = M.get();
208
209  // Create a new MCJIT instance to load this module then execute it.
210  createJIT(std::move(M));
211  TheJIT->setObjectCache(Cache.get());
212
213  // Verify that our object cache does not contain the module yet.
214  const MemoryBuffer *ObjBuffer = Cache->getObjectInternal(SecondModulePointer);
215  EXPECT_EQ(nullptr, ObjBuffer);
216
217  // Run the function and look for the replacement return code.
218  compileAndRun(ReplacementRC);
219
220  // Verify that MCJIT tried to look-up this module in the cache.
221  EXPECT_TRUE(Cache->wasModuleLookedUp(SecondModulePointer));
222
223  // Verify that our object cache now contains the module.
224  ObjBuffer = Cache->getObjectInternal(SecondModulePointer);
225  EXPECT_TRUE(nullptr != ObjBuffer);
226
227  // Verify that MCJIT didn't try to cache this again.
228  EXPECT_FALSE(Cache->wereDuplicatesInserted());
229}
230
231} // Namespace
232
233