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 | |
18 | class MockJITLinkMemoryManager : public llvm::jitlink::JITLinkMemoryManager { |
19 | public: |
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 | |
118 | void lookupResolveEverythingToNull( |
119 | const llvm::jitlink::JITLinkContext::LookupMap &Symbols, |
120 | std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation> LC); |
121 | |
122 | void lookupErrorOut( |
123 | const llvm::jitlink::JITLinkContext::LookupMap &Symbols, |
124 | std::unique_ptr<llvm::jitlink::JITLinkAsyncLookupContinuation> LC); |
125 | |
126 | class MockJITLinkContext : public llvm::jitlink::JITLinkContext { |
127 | public: |
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 | |
209 | std::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 | |
214 | void defaultMemMgrSetup(MockJITLinkMemoryManager &); |
215 | void defaultCtxSetup(MockJITLinkContext &); |
216 | |
217 | class JoinErrorsInto { |
218 | public: |
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 | |
224 | private: |
225 | llvm::Error &Err; |
226 | }; |
227 | |
228 | #endif // LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKMOCKS_H |
229 | |