llvm/unittests/ExecutionEngine/JITLink/JITLinkTestUtils.h

   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-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // Utilities for JITLink unit tests.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #ifndef LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKTESTUTILS_H
  14 #define LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKTESTUTILS_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(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     else
  78       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
 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(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     else
 161       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::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(std::move(Err), std::move(E2));
 222   }
 223
 224 private:
 225   llvm::Error &Err;
 226 };
 227
 228 extern llvm::ArrayRef<char> BlockContent;
 229
 230 #endif // LLVM_UNITTESTS_EXECUTIONENGINE_JITLINK_JITLINKTESTUTILS_H