1//===----- JITLinkMocks.h - Mock APIs for JITLink unit tests ----*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// Mock APIs for JITLink unit tests.
10//
11//===----------------------------------------------------------------------===//
12
13#ifndef LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKMOCKS_H
14#define LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKMOCKS_H
15
16#include "llvm/ExecutionEngine/JITLink/JITLink.h"
17
18class MockJITLinkMemoryManager : public llvm::jitlink::JITLinkMemoryManager {
19public:
20 class Alloc {
21 public:
22 virtual ~Alloc() {}
23 };
24
25 class SimpleAlloc : public Alloc {
26 public:
27 SimpleAlloc(const llvm::jitlink::JITLinkDylib *JD,
28 llvm::jitlink::LinkGraph &G) {
29 for (auto *B : G.blocks())
30 (void)B->getMutableContent(G);
31 }
32 };
33
34 class MockInFlightAlloc : public InFlightAlloc {
35 public:
36 MockInFlightAlloc(MockJITLinkMemoryManager &MJMM, std::unique_ptr<Alloc> A)
37 : MJMM(MJMM), A(std::move(A)) {}
38
39 void abandon(OnAbandonedFunction OnAbandoned) override {
40 OnAbandoned(MJMM.Abandon(std::move(A)));
41 }
42
43 void finalize(OnFinalizedFunction OnFinalized) override {
44 OnFinalized(MJMM.Finalize(std::move(A)));
45 }
46
47 private:
48 MockJITLinkMemoryManager &MJMM;
49 std::unique_ptr<Alloc> A;
50 };
51
52 MockJITLinkMemoryManager() {
53 Allocate = [this](const llvm::jitlink::JITLinkDylib *JD,
54 llvm::jitlink::LinkGraph &G) {
55 return defaultAllocate(JD, G);
56 };
57
58 Deallocate = [this](std::vector<FinalizedAlloc> Allocs) {
59 return defaultDeallocate(Allocs: std::move(Allocs));
60 };
61
62 Abandon = [this](std::unique_ptr<Alloc> A) {
63 return defaultAbandon(A: std::move(A));
64 };
65
66 Finalize = [this](std::unique_ptr<Alloc> A) {
67 return defaultFinalize(A: std::move(A));
68 };
69 }
70
71 void allocate(const llvm::jitlink::JITLinkDylib *JD,
72 llvm::jitlink::LinkGraph &G,
73 OnAllocatedFunction OnAllocated) override {
74 auto A = Allocate(JD, G);
75 if (!A)
76 OnAllocated(A.takeError());
77 else
78 OnAllocated(std::make_unique<MockInFlightAlloc>(args&: *this, args: std::move(*A)));
79 }
80
81 void deallocate(std::vector<FinalizedAlloc> Allocs,
82 OnDeallocatedFunction OnDeallocated) override {
83 OnDeallocated(Deallocate(std::move(Allocs)));
84 }
85
86 using JITLinkMemoryManager::allocate;
87 using JITLinkMemoryManager::deallocate;
88
89 llvm::Expected<std::unique_ptr<Alloc>>
90 defaultAllocate(const llvm::jitlink::JITLinkDylib *JD,
91 llvm::jitlink::LinkGraph &G) {
92 return std::make_unique<SimpleAlloc>(args&: JD, args&: G);
93 }
94
95 llvm::Error defaultDeallocate(std::vector<FinalizedAlloc> Allocs) {
96 for (auto &A : Allocs)
97 delete A.release().toPtr<Alloc *>();
98 return llvm::Error::success();
99 }
100
101 llvm::Error defaultAbandon(std::unique_ptr<Alloc> A) {
102 return llvm::Error::success();
103 }
104
105 llvm::Expected<FinalizedAlloc> defaultFinalize(std::unique_ptr<Alloc> A) {
106 return FinalizedAlloc(llvm::orc::ExecutorAddr::fromPtr(Ptr: A.release()));
107 }
108
109 llvm::unique_function<llvm::Expected<std::unique_ptr<Alloc>>(
110 const llvm::jitlink::JITLinkDylib *, llvm::jitlink::LinkGraph &)>
111 Allocate;
112 llvm::unique_function<llvm::Error(std::vector<FinalizedAlloc>)> Deallocate;
113 llvm::unique_function<llvm::Error(std::unique_ptr<Alloc>)> Abandon;
114 llvm::unique_function<llvm::Expected<FinalizedAlloc>(std::unique_ptr<Alloc>)>
115 Finalize;
116};
117
118void lookupResolveEverythingToNull(
119 const llvm::jitlink::JITLinkContext::LookupMap &Symbols,
120 std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation> LC);
121
122void lookupErrorOut(
123 const llvm::jitlink::JITLinkContext::LookupMap &Symbols,
124 std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation> LC);
125
126class MockJITLinkContext : public llvm::jitlink::JITLinkContext {
127public:
128 using HandleFailedFn = llvm::unique_function<void(llvm::Error)>;
129
130 MockJITLinkContext(std::unique_ptr<MockJITLinkMemoryManager> MJMM,
131 HandleFailedFn HandleFailed)
132 : JITLinkContext(&JD), MJMM(std::move(MJMM)),
133 HandleFailed(std::move(HandleFailed)) {}
134
135 ~MockJITLinkContext() {
136 if (auto Err = MJMM->deallocate(Allocs: std::move(FinalizedAllocs)))
137 notifyFailed(Err: std::move(Err));
138 }
139
140 llvm::jitlink::JITLinkMemoryManager &getMemoryManager() override {
141 return *MJMM;
142 }
143
144 void notifyFailed(llvm::Error Err) override { HandleFailed(std::move(Err)); }
145
146 void lookup(const LookupMap &Symbols,
147 std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation> LC)
148 override {
149 Lookup(Symbols, std::move(LC));
150 }
151
152 llvm::Error notifyResolved(llvm::jitlink::LinkGraph &G) override {
153 return NotifyResolved ? NotifyResolved(G) : llvm::Error::success();
154 }
155
156 void notifyFinalized(
157 llvm::jitlink::JITLinkMemoryManager::FinalizedAlloc Alloc) override {
158 if (NotifyFinalized)
159 NotifyFinalized(std::move(Alloc));
160 else
161 FinalizedAllocs.push_back(x: std::move(Alloc));
162 }
163
164 bool shouldAddDefaultTargetPasses(const llvm::Triple &TT) const override {
165 return true;
166 }
167
168 llvm::jitlink::LinkGraphPassFunction
169 getMarkLivePass(const llvm::Triple &TT) const override {
170 return MarkLivePass ? llvm::jitlink::LinkGraphPassFunction(
171 [this](llvm::jitlink::LinkGraph &G) {
172 return MarkLivePass(G);
173 })
174 : llvm::jitlink::LinkGraphPassFunction(
175 [](llvm::jitlink::LinkGraph &G) {
176 return markAllSymbolsLive(G);
177 });
178 }
179
180 llvm::Error
181 modifyPassConfig(llvm::jitlink::LinkGraph &G,
182 llvm::jitlink::PassConfiguration &Config) override {
183 if (ModifyPassConfig)
184 return ModifyPassConfig(G, Config);
185 return llvm::Error::success();
186 }
187
188 llvm::jitlink::JITLinkDylib JD{"JD"};
189 std::unique_ptr<MockJITLinkMemoryManager> MJMM;
190 HandleFailedFn HandleFailed;
191 llvm::unique_function<void(
192 const LookupMap &,
193 std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation>)>
194 Lookup;
195 llvm::unique_function<llvm::Error(llvm::jitlink::LinkGraph &)> NotifyResolved;
196 llvm::unique_function<void(
197 llvm::jitlink::JITLinkMemoryManager::FinalizedAlloc)>
198 NotifyFinalized;
199 mutable llvm::unique_function<llvm::Error(llvm::jitlink::LinkGraph &)>
200 MarkLivePass;
201 llvm::unique_function<llvm::Error(llvm::jitlink::LinkGraph &,
202 llvm::jitlink::PassConfiguration &)>
203 ModifyPassConfig;
204
205 std::vector<llvm::jitlink::JITLinkMemoryManager::FinalizedAlloc>
206 FinalizedAllocs;
207};
208
209std::unique_ptr<MockJITLinkContext> makeMockContext(
210 llvm::unique_function<void(llvm::Error)> HandleFailed,
211 llvm::unique_function<void(MockJITLinkMemoryManager &)> SetupMemMgr,
212 llvm::unique_function<void(MockJITLinkContext &)> SetupContext);
213
214void defaultMemMgrSetup(MockJITLinkMemoryManager &);
215void defaultCtxSetup(MockJITLinkContext &);
216
217class JoinErrorsInto {
218public:
219 JoinErrorsInto(llvm::Error &Err) : Err(Err) {}
220 void operator()(llvm::Error E2) {
221 Err = llvm::joinErrors(E1: std::move(Err), E2: std::move(E2));
222 }
223
224private:
225 llvm::Error &Err;
226};
227
228#endif // LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKMOCKS_H
229

source code of llvm/unittests/ExecutionEngine/JITLink/JITLinkMocks.h