231 lines · c
1//===--- JITLinkTestUtils.h - Utilities 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-exception6//7//===----------------------------------------------------------------------===//8//9// Utilities for JITLink unit tests.10//11//===----------------------------------------------------------------------===//12 13#ifndef LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKTESTUTILS_H14#define LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKTESTUTILS_H15 16#include "llvm/ExecutionEngine/JITLink/JITLink.h"17 18class MockJITLinkMemoryManager : public llvm::jitlink::JITLinkMemoryManager {19public:20 class Alloc {21 public:22 virtual ~Alloc() = default;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(std::move(Allocs));60 };61 62 Abandon = [this](std::unique_ptr<Alloc> A) {63 return defaultAbandon(std::move(A));64 };65 66 Finalize = [this](std::unique_ptr<Alloc> A) {67 return defaultFinalize(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 else78 OnAllocated(std::make_unique<MockInFlightAlloc>(*this, 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>(JD, 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(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() override {136 if (auto Err = MJMM->deallocate(std::move(FinalizedAllocs)))137 notifyFailed(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 else161 FinalizedAllocs.push_back(std::move(Alloc));162 }163 164 bool shouldAddDefaultTargetPasses(const llvm::Triple &TT) const override {165 return true;166 }167 168 llvm::jitlink::LinkGraphPassFunction169 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::Error181 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(std::move(Err), std::move(E2));222 }223 224private:225 llvm::Error &Err;226};227 228extern llvm::ArrayRef<char> BlockContent;229 230#endif // LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKTESTUTILS_H231