xref: /llvm-project/llvm/lib/ExecutionEngine/Orc/COFFPlatform.cpp (revision 4eaff6c58ae2f130ac8d63cf2c87bbb483114876)
1 //===------- COFFPlatform.cpp - Utilities for executing COFF in Orc -------===//
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 #include "llvm/ExecutionEngine/Orc/COFFPlatform.h"
10 #include "llvm/ExecutionEngine/Orc/AbsoluteSymbols.h"
11 #include "llvm/ExecutionEngine/Orc/DebugUtils.h"
12 #include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
13 #include "llvm/ExecutionEngine/Orc/ObjectFileInterface.h"
14 #include "llvm/ExecutionEngine/Orc/Shared/ObjectFormats.h"
15 
16 #include "llvm/Object/COFF.h"
17 
18 #include "llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h"
19 
20 #include "llvm/ExecutionEngine/JITLink/x86_64.h"
21 
22 #define DEBUG_TYPE "orc"
23 
24 using namespace llvm;
25 using namespace llvm::orc;
26 using namespace llvm::orc::shared;
27 
28 namespace llvm {
29 namespace orc {
30 namespace shared {
31 
32 using SPSCOFFJITDylibDepInfo = SPSSequence<SPSExecutorAddr>;
33 using SPSCOFFJITDylibDepInfoMap =
34     SPSSequence<SPSTuple<SPSExecutorAddr, SPSCOFFJITDylibDepInfo>>;
35 using SPSCOFFObjectSectionsMap =
36     SPSSequence<SPSTuple<SPSString, SPSExecutorAddrRange>>;
37 using SPSCOFFRegisterObjectSectionsArgs =
38     SPSArgList<SPSExecutorAddr, SPSCOFFObjectSectionsMap, bool>;
39 using SPSCOFFDeregisterObjectSectionsArgs =
40     SPSArgList<SPSExecutorAddr, SPSCOFFObjectSectionsMap>;
41 
42 } // namespace shared
43 } // namespace orc
44 } // namespace llvm
45 namespace {
46 
47 class COFFHeaderMaterializationUnit : public MaterializationUnit {
48 public:
49   COFFHeaderMaterializationUnit(COFFPlatform &CP,
50                                 const SymbolStringPtr &HeaderStartSymbol)
51       : MaterializationUnit(createHeaderInterface(CP, HeaderStartSymbol)),
52         CP(CP) {}
53 
54   StringRef getName() const override { return "COFFHeaderMU"; }
55 
56   void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
57     auto G = std::make_unique<jitlink::LinkGraph>(
58         "<COFFHeaderMU>", CP.getExecutionSession().getSymbolStringPool(),
59         CP.getExecutionSession().getTargetTriple(), SubtargetFeatures(),
60         jitlink::getGenericEdgeKindName);
61     auto &HeaderSection = G->createSection("__header", MemProt::Read);
62     auto &HeaderBlock = createHeaderBlock(*G, HeaderSection);
63 
64     // Init symbol is __ImageBase symbol.
65     auto &ImageBaseSymbol = G->addDefinedSymbol(
66         HeaderBlock, 0, *R->getInitializerSymbol(), HeaderBlock.getSize(),
67         jitlink::Linkage::Strong, jitlink::Scope::Default, false, true);
68 
69     addImageBaseRelocationEdge(HeaderBlock, ImageBaseSymbol);
70 
71     CP.getObjectLinkingLayer().emit(std::move(R), std::move(G));
72   }
73 
74   void discard(const JITDylib &JD, const SymbolStringPtr &Sym) override {}
75 
76 private:
77   struct HeaderSymbol {
78     const char *Name;
79     uint64_t Offset;
80   };
81 
82   struct NTHeader {
83     support::ulittle32_t PEMagic;
84     object::coff_file_header FileHeader;
85     struct PEHeader {
86       object::pe32plus_header Header;
87       object::data_directory DataDirectory[COFF::NUM_DATA_DIRECTORIES + 1];
88     } OptionalHeader;
89   };
90 
91   struct HeaderBlockContent {
92     object::dos_header DOSHeader;
93     COFFHeaderMaterializationUnit::NTHeader NTHeader;
94   };
95 
96   static jitlink::Block &createHeaderBlock(jitlink::LinkGraph &G,
97                                            jitlink::Section &HeaderSection) {
98     HeaderBlockContent Hdr = {};
99 
100     // Set up magic
101     Hdr.DOSHeader.Magic[0] = 'M';
102     Hdr.DOSHeader.Magic[1] = 'Z';
103     Hdr.DOSHeader.AddressOfNewExeHeader =
104         offsetof(HeaderBlockContent, NTHeader);
105     uint32_t PEMagic = *reinterpret_cast<const uint32_t *>(COFF::PEMagic);
106     Hdr.NTHeader.PEMagic = PEMagic;
107     Hdr.NTHeader.OptionalHeader.Header.Magic = COFF::PE32Header::PE32_PLUS;
108 
109     switch (G.getTargetTriple().getArch()) {
110     case Triple::x86_64:
111       Hdr.NTHeader.FileHeader.Machine = COFF::IMAGE_FILE_MACHINE_AMD64;
112       break;
113     default:
114       llvm_unreachable("Unrecognized architecture");
115     }
116 
117     auto HeaderContent = G.allocateContent(
118         ArrayRef<char>(reinterpret_cast<const char *>(&Hdr), sizeof(Hdr)));
119 
120     return G.createContentBlock(HeaderSection, HeaderContent, ExecutorAddr(), 8,
121                                 0);
122   }
123 
124   static void addImageBaseRelocationEdge(jitlink::Block &B,
125                                          jitlink::Symbol &ImageBase) {
126     auto ImageBaseOffset = offsetof(HeaderBlockContent, NTHeader) +
127                            offsetof(NTHeader, OptionalHeader) +
128                            offsetof(object::pe32plus_header, ImageBase);
129     B.addEdge(jitlink::x86_64::Pointer64, ImageBaseOffset, ImageBase, 0);
130   }
131 
132   static MaterializationUnit::Interface
133   createHeaderInterface(COFFPlatform &MOP,
134                         const SymbolStringPtr &HeaderStartSymbol) {
135     SymbolFlagsMap HeaderSymbolFlags;
136 
137     HeaderSymbolFlags[HeaderStartSymbol] = JITSymbolFlags::Exported;
138 
139     return MaterializationUnit::Interface(std::move(HeaderSymbolFlags),
140                                           HeaderStartSymbol);
141   }
142 
143   COFFPlatform &CP;
144 };
145 
146 } // end anonymous namespace
147 
148 namespace llvm {
149 namespace orc {
150 
151 Expected<std::unique_ptr<COFFPlatform>>
152 COFFPlatform::Create(ObjectLinkingLayer &ObjLinkingLayer, JITDylib &PlatformJD,
153                      std::unique_ptr<MemoryBuffer> OrcRuntimeArchiveBuffer,
154                      LoadDynamicLibrary LoadDynLibrary, bool StaticVCRuntime,
155                      const char *VCRuntimePath,
156                      std::optional<SymbolAliasMap> RuntimeAliases) {
157 
158   auto &ES = ObjLinkingLayer.getExecutionSession();
159 
160   // If the target is not supported then bail out immediately.
161   if (!supportedTarget(ES.getTargetTriple()))
162     return make_error<StringError>("Unsupported COFFPlatform triple: " +
163                                        ES.getTargetTriple().str(),
164                                    inconvertibleErrorCode());
165 
166   auto &EPC = ES.getExecutorProcessControl();
167 
168   auto GeneratorArchive =
169       object::Archive::create(OrcRuntimeArchiveBuffer->getMemBufferRef());
170   if (!GeneratorArchive)
171     return GeneratorArchive.takeError();
172 
173   auto OrcRuntimeArchiveGenerator = StaticLibraryDefinitionGenerator::Create(
174       ObjLinkingLayer, nullptr, std::move(*GeneratorArchive));
175   if (!OrcRuntimeArchiveGenerator)
176     return OrcRuntimeArchiveGenerator.takeError();
177 
178   // We need a second instance of the archive (for now) for the Platform. We
179   // can `cantFail` this call, since if it were going to fail it would have
180   // failed above.
181   auto RuntimeArchive = cantFail(
182       object::Archive::create(OrcRuntimeArchiveBuffer->getMemBufferRef()));
183 
184   // Create default aliases if the caller didn't supply any.
185   if (!RuntimeAliases)
186     RuntimeAliases = standardPlatformAliases(ES);
187 
188   // Define the aliases.
189   if (auto Err = PlatformJD.define(symbolAliases(std::move(*RuntimeAliases))))
190     return std::move(Err);
191 
192   auto &HostFuncJD = ES.createBareJITDylib("$<PlatformRuntimeHostFuncJD>");
193 
194   // Add JIT-dispatch function support symbols.
195   if (auto Err = HostFuncJD.define(
196           absoluteSymbols({{ES.intern("__orc_rt_jit_dispatch"),
197                             {EPC.getJITDispatchInfo().JITDispatchFunction,
198                              JITSymbolFlags::Exported}},
199                            {ES.intern("__orc_rt_jit_dispatch_ctx"),
200                             {EPC.getJITDispatchInfo().JITDispatchContext,
201                              JITSymbolFlags::Exported}}})))
202     return std::move(Err);
203 
204   PlatformJD.addToLinkOrder(HostFuncJD);
205 
206   // Create the instance.
207   Error Err = Error::success();
208   auto P = std::unique_ptr<COFFPlatform>(new COFFPlatform(
209       ObjLinkingLayer, PlatformJD, std::move(*OrcRuntimeArchiveGenerator),
210       std::move(OrcRuntimeArchiveBuffer), std::move(RuntimeArchive),
211       std::move(LoadDynLibrary), StaticVCRuntime, VCRuntimePath, Err));
212   if (Err)
213     return std::move(Err);
214   return std::move(P);
215 }
216 
217 Expected<std::unique_ptr<COFFPlatform>>
218 COFFPlatform::Create(ObjectLinkingLayer &ObjLinkingLayer, JITDylib &PlatformJD,
219                      const char *OrcRuntimePath,
220                      LoadDynamicLibrary LoadDynLibrary, bool StaticVCRuntime,
221                      const char *VCRuntimePath,
222                      std::optional<SymbolAliasMap> RuntimeAliases) {
223 
224   auto ArchiveBuffer = MemoryBuffer::getFile(OrcRuntimePath);
225   if (!ArchiveBuffer)
226     return createFileError(OrcRuntimePath, ArchiveBuffer.getError());
227 
228   return Create(ObjLinkingLayer, PlatformJD, std::move(*ArchiveBuffer),
229                 std::move(LoadDynLibrary), StaticVCRuntime, VCRuntimePath,
230                 std::move(RuntimeAliases));
231 }
232 
233 Expected<MemoryBufferRef> COFFPlatform::getPerJDObjectFile() {
234   auto PerJDObj = OrcRuntimeArchive->findSym("__orc_rt_coff_per_jd_marker");
235   if (!PerJDObj)
236     return PerJDObj.takeError();
237 
238   if (!*PerJDObj)
239     return make_error<StringError>("Could not find per jd object file",
240                                    inconvertibleErrorCode());
241 
242   auto Buffer = (*PerJDObj)->getAsBinary();
243   if (!Buffer)
244     return Buffer.takeError();
245 
246   return (*Buffer)->getMemoryBufferRef();
247 }
248 
249 static void addAliases(ExecutionSession &ES, SymbolAliasMap &Aliases,
250                        ArrayRef<std::pair<const char *, const char *>> AL) {
251   for (auto &KV : AL) {
252     auto AliasName = ES.intern(KV.first);
253     assert(!Aliases.count(AliasName) && "Duplicate symbol name in alias map");
254     Aliases[std::move(AliasName)] = {ES.intern(KV.second),
255                                      JITSymbolFlags::Exported};
256   }
257 }
258 
259 Error COFFPlatform::setupJITDylib(JITDylib &JD) {
260   if (auto Err = JD.define(std::make_unique<COFFHeaderMaterializationUnit>(
261           *this, COFFHeaderStartSymbol)))
262     return Err;
263 
264   if (auto Err = ES.lookup({&JD}, COFFHeaderStartSymbol).takeError())
265     return Err;
266 
267   // Define the CXX aliases.
268   SymbolAliasMap CXXAliases;
269   addAliases(ES, CXXAliases, requiredCXXAliases());
270   if (auto Err = JD.define(symbolAliases(std::move(CXXAliases))))
271     return Err;
272 
273   auto PerJDObj = getPerJDObjectFile();
274   if (!PerJDObj)
275     return PerJDObj.takeError();
276 
277   auto I = getObjectFileInterface(ES, *PerJDObj);
278   if (!I)
279     return I.takeError();
280 
281   if (auto Err = ObjLinkingLayer.add(
282           JD, MemoryBuffer::getMemBuffer(*PerJDObj, false), std::move(*I)))
283     return Err;
284 
285   if (!Bootstrapping) {
286     auto ImportedLibs = StaticVCRuntime
287                             ? VCRuntimeBootstrap->loadStaticVCRuntime(JD)
288                             : VCRuntimeBootstrap->loadDynamicVCRuntime(JD);
289     if (!ImportedLibs)
290       return ImportedLibs.takeError();
291     for (auto &Lib : *ImportedLibs)
292       if (auto Err = LoadDynLibrary(JD, Lib))
293         return Err;
294     if (StaticVCRuntime)
295       if (auto Err = VCRuntimeBootstrap->initializeStaticVCRuntime(JD))
296         return Err;
297   }
298 
299   JD.addGenerator(DLLImportDefinitionGenerator::Create(ES, ObjLinkingLayer));
300   return Error::success();
301 }
302 
303 Error COFFPlatform::teardownJITDylib(JITDylib &JD) {
304   std::lock_guard<std::mutex> Lock(PlatformMutex);
305   auto I = JITDylibToHeaderAddr.find(&JD);
306   if (I != JITDylibToHeaderAddr.end()) {
307     assert(HeaderAddrToJITDylib.count(I->second) &&
308            "HeaderAddrToJITDylib missing entry");
309     HeaderAddrToJITDylib.erase(I->second);
310     JITDylibToHeaderAddr.erase(I);
311   }
312   return Error::success();
313 }
314 
315 Error COFFPlatform::notifyAdding(ResourceTracker &RT,
316                                  const MaterializationUnit &MU) {
317   auto &JD = RT.getJITDylib();
318   const auto &InitSym = MU.getInitializerSymbol();
319   if (!InitSym)
320     return Error::success();
321 
322   RegisteredInitSymbols[&JD].add(InitSym,
323                                  SymbolLookupFlags::WeaklyReferencedSymbol);
324 
325   LLVM_DEBUG({
326     dbgs() << "COFFPlatform: Registered init symbol " << *InitSym << " for MU "
327            << MU.getName() << "\n";
328   });
329   return Error::success();
330 }
331 
332 Error COFFPlatform::notifyRemoving(ResourceTracker &RT) {
333   llvm_unreachable("Not supported yet");
334 }
335 
336 SymbolAliasMap COFFPlatform::standardPlatformAliases(ExecutionSession &ES) {
337   SymbolAliasMap Aliases;
338   addAliases(ES, Aliases, standardRuntimeUtilityAliases());
339   return Aliases;
340 }
341 
342 ArrayRef<std::pair<const char *, const char *>>
343 COFFPlatform::requiredCXXAliases() {
344   static const std::pair<const char *, const char *> RequiredCXXAliases[] = {
345       {"_CxxThrowException", "__orc_rt_coff_cxx_throw_exception"},
346       {"_onexit", "__orc_rt_coff_onexit_per_jd"},
347       {"atexit", "__orc_rt_coff_atexit_per_jd"}};
348 
349   return ArrayRef<std::pair<const char *, const char *>>(RequiredCXXAliases);
350 }
351 
352 ArrayRef<std::pair<const char *, const char *>>
353 COFFPlatform::standardRuntimeUtilityAliases() {
354   static const std::pair<const char *, const char *>
355       StandardRuntimeUtilityAliases[] = {
356           {"__orc_rt_run_program", "__orc_rt_coff_run_program"},
357           {"__orc_rt_jit_dlerror", "__orc_rt_coff_jit_dlerror"},
358           {"__orc_rt_jit_dlopen", "__orc_rt_coff_jit_dlopen"},
359           {"__orc_rt_jit_dlclose", "__orc_rt_coff_jit_dlclose"},
360           {"__orc_rt_jit_dlsym", "__orc_rt_coff_jit_dlsym"},
361           {"__orc_rt_log_error", "__orc_rt_log_error_to_stderr"}};
362 
363   return ArrayRef<std::pair<const char *, const char *>>(
364       StandardRuntimeUtilityAliases);
365 }
366 
367 bool COFFPlatform::supportedTarget(const Triple &TT) {
368   switch (TT.getArch()) {
369   case Triple::x86_64:
370     return true;
371   default:
372     return false;
373   }
374 }
375 
376 COFFPlatform::COFFPlatform(
377     ObjectLinkingLayer &ObjLinkingLayer, JITDylib &PlatformJD,
378     std::unique_ptr<StaticLibraryDefinitionGenerator> OrcRuntimeGenerator,
379     std::unique_ptr<MemoryBuffer> OrcRuntimeArchiveBuffer,
380     std::unique_ptr<object::Archive> OrcRuntimeArchive,
381     LoadDynamicLibrary LoadDynLibrary, bool StaticVCRuntime,
382     const char *VCRuntimePath, Error &Err)
383     : ES(ObjLinkingLayer.getExecutionSession()),
384       ObjLinkingLayer(ObjLinkingLayer),
385       LoadDynLibrary(std::move(LoadDynLibrary)),
386       OrcRuntimeArchiveBuffer(std::move(OrcRuntimeArchiveBuffer)),
387       OrcRuntimeArchive(std::move(OrcRuntimeArchive)),
388       StaticVCRuntime(StaticVCRuntime),
389       COFFHeaderStartSymbol(ES.intern("__ImageBase")) {
390   ErrorAsOutParameter _(Err);
391 
392   Bootstrapping.store(true);
393   ObjLinkingLayer.addPlugin(std::make_unique<COFFPlatformPlugin>(*this));
394 
395   // Load vc runtime
396   auto VCRT =
397       COFFVCRuntimeBootstrapper::Create(ES, ObjLinkingLayer, VCRuntimePath);
398   if (!VCRT) {
399     Err = VCRT.takeError();
400     return;
401   }
402   VCRuntimeBootstrap = std::move(*VCRT);
403 
404   for (auto &Lib : OrcRuntimeGenerator->getImportedDynamicLibraries())
405     DylibsToPreload.insert(Lib);
406 
407   auto ImportedLibs =
408       StaticVCRuntime ? VCRuntimeBootstrap->loadStaticVCRuntime(PlatformJD)
409                       : VCRuntimeBootstrap->loadDynamicVCRuntime(PlatformJD);
410   if (!ImportedLibs) {
411     Err = ImportedLibs.takeError();
412     return;
413   }
414 
415   for (auto &Lib : *ImportedLibs)
416     DylibsToPreload.insert(Lib);
417 
418   PlatformJD.addGenerator(std::move(OrcRuntimeGenerator));
419 
420   // PlatformJD hasn't been set up by the platform yet (since we're creating
421   // the platform now), so set it up.
422   if (auto E2 = setupJITDylib(PlatformJD)) {
423     Err = std::move(E2);
424     return;
425   }
426 
427   for (auto& Lib : DylibsToPreload)
428     if (auto E2 = this->LoadDynLibrary(PlatformJD, Lib)) {
429       Err = std::move(E2);
430       return;
431     }
432 
433   if (StaticVCRuntime)
434       if (auto E2 = VCRuntimeBootstrap->initializeStaticVCRuntime(PlatformJD)) {
435           Err = std::move(E2);
436           return;
437       }
438 
439   // Associate wrapper function tags with JIT-side function implementations.
440   if (auto E2 = associateRuntimeSupportFunctions(PlatformJD)) {
441       Err = std::move(E2);
442       return;
443   }
444 
445   // Lookup addresses of runtime functions callable by the platform,
446   // call the platform bootstrap function to initialize the platform-state
447   // object in the executor.
448   if (auto E2 = bootstrapCOFFRuntime(PlatformJD)) {
449       Err = std::move(E2);
450       return;
451   }
452 
453   Bootstrapping.store(false);
454   JDBootstrapStates.clear();
455 }
456 
457 Expected<COFFPlatform::JITDylibDepMap>
458 COFFPlatform::buildJDDepMap(JITDylib &JD) {
459   return ES.runSessionLocked([&]() -> Expected<JITDylibDepMap> {
460     JITDylibDepMap JDDepMap;
461 
462     SmallVector<JITDylib *, 16> Worklist({&JD});
463     while (!Worklist.empty()) {
464       auto CurJD = Worklist.back();
465       Worklist.pop_back();
466 
467       auto &DM = JDDepMap[CurJD];
468       CurJD->withLinkOrderDo([&](const JITDylibSearchOrder &O) {
469         DM.reserve(O.size());
470         for (auto &KV : O) {
471           if (KV.first == CurJD)
472             continue;
473           {
474             // Bare jitdylibs not known to the platform
475             std::lock_guard<std::mutex> Lock(PlatformMutex);
476             if (!JITDylibToHeaderAddr.count(KV.first)) {
477               LLVM_DEBUG({
478                 dbgs() << "JITDylib unregistered to COFFPlatform detected in "
479                           "LinkOrder: "
480                        << CurJD->getName() << "\n";
481               });
482               continue;
483             }
484           }
485           DM.push_back(KV.first);
486           // Push unvisited entry.
487           if (!JDDepMap.count(KV.first)) {
488             Worklist.push_back(KV.first);
489             JDDepMap[KV.first] = {};
490           }
491         }
492       });
493     }
494     return std::move(JDDepMap);
495   });
496 }
497 
498 void COFFPlatform::pushInitializersLoop(PushInitializersSendResultFn SendResult,
499                                         JITDylibSP JD,
500                                         JITDylibDepMap &JDDepMap) {
501   SmallVector<JITDylib *, 16> Worklist({JD.get()});
502   DenseSet<JITDylib *> Visited({JD.get()});
503   DenseMap<JITDylib *, SymbolLookupSet> NewInitSymbols;
504   ES.runSessionLocked([&]() {
505     while (!Worklist.empty()) {
506       auto CurJD = Worklist.back();
507       Worklist.pop_back();
508 
509       auto RISItr = RegisteredInitSymbols.find(CurJD);
510       if (RISItr != RegisteredInitSymbols.end()) {
511         NewInitSymbols[CurJD] = std::move(RISItr->second);
512         RegisteredInitSymbols.erase(RISItr);
513       }
514 
515       for (auto *DepJD : JDDepMap[CurJD])
516         if (Visited.insert(DepJD).second)
517           Worklist.push_back(DepJD);
518     }
519   });
520 
521   // If there are no further init symbols to look up then send the link order
522   // (as a list of header addresses) to the caller.
523   if (NewInitSymbols.empty()) {
524     // Build the dep info map to return.
525     COFFJITDylibDepInfoMap DIM;
526     DIM.reserve(JDDepMap.size());
527     for (auto &KV : JDDepMap) {
528       std::lock_guard<std::mutex> Lock(PlatformMutex);
529       COFFJITDylibDepInfo DepInfo;
530       DepInfo.reserve(KV.second.size());
531       for (auto &Dep : KV.second) {
532         DepInfo.push_back(JITDylibToHeaderAddr[Dep]);
533       }
534       auto H = JITDylibToHeaderAddr[KV.first];
535       DIM.push_back(std::make_pair(H, std::move(DepInfo)));
536     }
537     SendResult(DIM);
538     return;
539   }
540 
541   // Otherwise issue a lookup and re-run this phase when it completes.
542   lookupInitSymbolsAsync(
543       [this, SendResult = std::move(SendResult), &JD,
544        JDDepMap = std::move(JDDepMap)](Error Err) mutable {
545         if (Err)
546           SendResult(std::move(Err));
547         else
548           pushInitializersLoop(std::move(SendResult), JD, JDDepMap);
549       },
550       ES, std::move(NewInitSymbols));
551 }
552 
553 void COFFPlatform::rt_pushInitializers(PushInitializersSendResultFn SendResult,
554                                        ExecutorAddr JDHeaderAddr) {
555   JITDylibSP JD;
556   {
557     std::lock_guard<std::mutex> Lock(PlatformMutex);
558     auto I = HeaderAddrToJITDylib.find(JDHeaderAddr);
559     if (I != HeaderAddrToJITDylib.end())
560       JD = I->second;
561   }
562 
563   LLVM_DEBUG({
564     dbgs() << "COFFPlatform::rt_pushInitializers(" << JDHeaderAddr << ") ";
565     if (JD)
566       dbgs() << "pushing initializers for " << JD->getName() << "\n";
567     else
568       dbgs() << "No JITDylib for header address.\n";
569   });
570 
571   if (!JD) {
572     SendResult(make_error<StringError>("No JITDylib with header addr " +
573                                            formatv("{0:x}", JDHeaderAddr),
574                                        inconvertibleErrorCode()));
575     return;
576   }
577 
578   auto JDDepMap = buildJDDepMap(*JD);
579   if (!JDDepMap) {
580     SendResult(JDDepMap.takeError());
581     return;
582   }
583 
584   pushInitializersLoop(std::move(SendResult), JD, *JDDepMap);
585 }
586 
587 void COFFPlatform::rt_lookupSymbol(SendSymbolAddressFn SendResult,
588                                    ExecutorAddr Handle, StringRef SymbolName) {
589   LLVM_DEBUG(dbgs() << "COFFPlatform::rt_lookupSymbol(\"" << Handle << "\")\n");
590 
591   JITDylib *JD = nullptr;
592 
593   {
594     std::lock_guard<std::mutex> Lock(PlatformMutex);
595     auto I = HeaderAddrToJITDylib.find(Handle);
596     if (I != HeaderAddrToJITDylib.end())
597       JD = I->second;
598   }
599 
600   if (!JD) {
601     LLVM_DEBUG(dbgs() << "  No JITDylib for handle " << Handle << "\n");
602     SendResult(make_error<StringError>("No JITDylib associated with handle " +
603                                            formatv("{0:x}", Handle),
604                                        inconvertibleErrorCode()));
605     return;
606   }
607 
608   // Use functor class to work around XL build compiler issue on AIX.
609   class RtLookupNotifyComplete {
610   public:
611     RtLookupNotifyComplete(SendSymbolAddressFn &&SendResult)
612         : SendResult(std::move(SendResult)) {}
613     void operator()(Expected<SymbolMap> Result) {
614       if (Result) {
615         assert(Result->size() == 1 && "Unexpected result map count");
616         SendResult(Result->begin()->second.getAddress());
617       } else {
618         SendResult(Result.takeError());
619       }
620     }
621 
622   private:
623     SendSymbolAddressFn SendResult;
624   };
625 
626   ES.lookup(
627       LookupKind::DLSym, {{JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
628       SymbolLookupSet(ES.intern(SymbolName)), SymbolState::Ready,
629       RtLookupNotifyComplete(std::move(SendResult)), NoDependenciesToRegister);
630 }
631 
632 Error COFFPlatform::associateRuntimeSupportFunctions(JITDylib &PlatformJD) {
633   ExecutionSession::JITDispatchHandlerAssociationMap WFs;
634 
635   using LookupSymbolSPSSig =
636       SPSExpected<SPSExecutorAddr>(SPSExecutorAddr, SPSString);
637   WFs[ES.intern("__orc_rt_coff_symbol_lookup_tag")] =
638       ES.wrapAsyncWithSPS<LookupSymbolSPSSig>(this,
639                                               &COFFPlatform::rt_lookupSymbol);
640   using PushInitializersSPSSig =
641       SPSExpected<SPSCOFFJITDylibDepInfoMap>(SPSExecutorAddr);
642   WFs[ES.intern("__orc_rt_coff_push_initializers_tag")] =
643       ES.wrapAsyncWithSPS<PushInitializersSPSSig>(
644           this, &COFFPlatform::rt_pushInitializers);
645 
646   return ES.registerJITDispatchHandlers(PlatformJD, std::move(WFs));
647 }
648 
649 Error COFFPlatform::runBootstrapInitializers(JDBootstrapState &BState) {
650   llvm::sort(BState.Initializers);
651   if (auto Err =
652           runBootstrapSubsectionInitializers(BState, ".CRT$XIA", ".CRT$XIZ"))
653     return Err;
654 
655   if (auto Err = runSymbolIfExists(*BState.JD, "__run_after_c_init"))
656     return Err;
657 
658   if (auto Err =
659           runBootstrapSubsectionInitializers(BState, ".CRT$XCA", ".CRT$XCZ"))
660     return Err;
661   return Error::success();
662 }
663 
664 Error COFFPlatform::runBootstrapSubsectionInitializers(JDBootstrapState &BState,
665                                                        StringRef Start,
666                                                        StringRef End) {
667   for (auto &Initializer : BState.Initializers)
668     if (Initializer.first >= Start && Initializer.first <= End &&
669         Initializer.second) {
670       auto Res =
671           ES.getExecutorProcessControl().runAsVoidFunction(Initializer.second);
672       if (!Res)
673         return Res.takeError();
674     }
675   return Error::success();
676 }
677 
678 Error COFFPlatform::bootstrapCOFFRuntime(JITDylib &PlatformJD) {
679   // Lookup of runtime symbols causes the collection of initializers if
680   // it's static linking setting.
681   if (auto Err = lookupAndRecordAddrs(
682           ES, LookupKind::Static, makeJITDylibSearchOrder(&PlatformJD),
683           {
684               {ES.intern("__orc_rt_coff_platform_bootstrap"),
685                &orc_rt_coff_platform_bootstrap},
686               {ES.intern("__orc_rt_coff_platform_shutdown"),
687                &orc_rt_coff_platform_shutdown},
688               {ES.intern("__orc_rt_coff_register_jitdylib"),
689                &orc_rt_coff_register_jitdylib},
690               {ES.intern("__orc_rt_coff_deregister_jitdylib"),
691                &orc_rt_coff_deregister_jitdylib},
692               {ES.intern("__orc_rt_coff_register_object_sections"),
693                &orc_rt_coff_register_object_sections},
694               {ES.intern("__orc_rt_coff_deregister_object_sections"),
695                &orc_rt_coff_deregister_object_sections},
696           }))
697     return Err;
698 
699   // Call bootstrap functions
700   if (auto Err = ES.callSPSWrapper<void()>(orc_rt_coff_platform_bootstrap))
701     return Err;
702 
703   // Do the pending jitdylib registration actions that we couldn't do
704   // because orc runtime was not linked fully.
705   for (auto KV : JDBootstrapStates) {
706     auto &JDBState = KV.second;
707     if (auto Err = ES.callSPSWrapper<void(SPSString, SPSExecutorAddr)>(
708             orc_rt_coff_register_jitdylib, JDBState.JDName,
709             JDBState.HeaderAddr))
710       return Err;
711 
712     for (auto &ObjSectionMap : JDBState.ObjectSectionsMaps)
713       if (auto Err = ES.callSPSWrapper<void(SPSExecutorAddr,
714                                             SPSCOFFObjectSectionsMap, bool)>(
715               orc_rt_coff_register_object_sections, JDBState.HeaderAddr,
716               ObjSectionMap, false))
717         return Err;
718   }
719 
720   // Run static initializers collected in bootstrap stage.
721   for (auto KV : JDBootstrapStates) {
722     auto &JDBState = KV.second;
723     if (auto Err = runBootstrapInitializers(JDBState))
724       return Err;
725   }
726 
727   return Error::success();
728 }
729 
730 Error COFFPlatform::runSymbolIfExists(JITDylib &PlatformJD,
731                                       StringRef SymbolName) {
732   ExecutorAddr jit_function;
733   auto AfterCLookupErr = lookupAndRecordAddrs(
734       ES, LookupKind::Static, makeJITDylibSearchOrder(&PlatformJD),
735       {{ES.intern(SymbolName), &jit_function}});
736   if (!AfterCLookupErr) {
737     auto Res = ES.getExecutorProcessControl().runAsVoidFunction(jit_function);
738     if (!Res)
739       return Res.takeError();
740     return Error::success();
741   }
742   if (!AfterCLookupErr.isA<SymbolsNotFound>())
743     return AfterCLookupErr;
744   consumeError(std::move(AfterCLookupErr));
745   return Error::success();
746 }
747 
748 void COFFPlatform::COFFPlatformPlugin::modifyPassConfig(
749     MaterializationResponsibility &MR, jitlink::LinkGraph &LG,
750     jitlink::PassConfiguration &Config) {
751 
752   bool IsBootstrapping = CP.Bootstrapping.load();
753 
754   if (auto InitSymbol = MR.getInitializerSymbol()) {
755     if (InitSymbol == CP.COFFHeaderStartSymbol) {
756       Config.PostAllocationPasses.push_back(
757           [this, &MR, IsBootstrapping](jitlink::LinkGraph &G) {
758             return associateJITDylibHeaderSymbol(G, MR, IsBootstrapping);
759           });
760       return;
761     }
762     Config.PrePrunePasses.push_back([this, &MR](jitlink::LinkGraph &G) {
763       return preserveInitializerSections(G, MR);
764     });
765   }
766 
767   if (!IsBootstrapping)
768     Config.PostFixupPasses.push_back(
769         [this, &JD = MR.getTargetJITDylib()](jitlink::LinkGraph &G) {
770           return registerObjectPlatformSections(G, JD);
771         });
772   else
773     Config.PostFixupPasses.push_back(
774         [this, &JD = MR.getTargetJITDylib()](jitlink::LinkGraph &G) {
775           return registerObjectPlatformSectionsInBootstrap(G, JD);
776         });
777 }
778 
779 Error COFFPlatform::COFFPlatformPlugin::associateJITDylibHeaderSymbol(
780     jitlink::LinkGraph &G, MaterializationResponsibility &MR,
781     bool IsBootstraping) {
782   auto I = llvm::find_if(G.defined_symbols(), [this](jitlink::Symbol *Sym) {
783     return *Sym->getName() == *CP.COFFHeaderStartSymbol;
784   });
785   assert(I != G.defined_symbols().end() && "Missing COFF header start symbol");
786 
787   auto &JD = MR.getTargetJITDylib();
788   std::lock_guard<std::mutex> Lock(CP.PlatformMutex);
789   auto HeaderAddr = (*I)->getAddress();
790   CP.JITDylibToHeaderAddr[&JD] = HeaderAddr;
791   CP.HeaderAddrToJITDylib[HeaderAddr] = &JD;
792   if (!IsBootstraping) {
793     G.allocActions().push_back(
794         {cantFail(WrapperFunctionCall::Create<
795                   SPSArgList<SPSString, SPSExecutorAddr>>(
796              CP.orc_rt_coff_register_jitdylib, JD.getName(), HeaderAddr)),
797          cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
798              CP.orc_rt_coff_deregister_jitdylib, HeaderAddr))});
799   } else {
800     G.allocActions().push_back(
801         {{},
802          cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
803              CP.orc_rt_coff_deregister_jitdylib, HeaderAddr))});
804     JDBootstrapState BState;
805     BState.JD = &JD;
806     BState.JDName = JD.getName();
807     BState.HeaderAddr = HeaderAddr;
808     CP.JDBootstrapStates.emplace(&JD, BState);
809   }
810 
811   return Error::success();
812 }
813 
814 Error COFFPlatform::COFFPlatformPlugin::registerObjectPlatformSections(
815     jitlink::LinkGraph &G, JITDylib &JD) {
816   COFFObjectSectionsMap ObjSecs;
817   auto HeaderAddr = CP.JITDylibToHeaderAddr[&JD];
818   assert(HeaderAddr && "Must be registered jitdylib");
819   for (auto &S : G.sections()) {
820     jitlink::SectionRange Range(S);
821     if (Range.getSize())
822       ObjSecs.push_back(std::make_pair(S.getName().str(), Range.getRange()));
823   }
824 
825   G.allocActions().push_back(
826       {cantFail(WrapperFunctionCall::Create<SPSCOFFRegisterObjectSectionsArgs>(
827            CP.orc_rt_coff_register_object_sections, HeaderAddr, ObjSecs, true)),
828        cantFail(
829            WrapperFunctionCall::Create<SPSCOFFDeregisterObjectSectionsArgs>(
830                CP.orc_rt_coff_deregister_object_sections, HeaderAddr,
831                ObjSecs))});
832 
833   return Error::success();
834 }
835 
836 Error COFFPlatform::COFFPlatformPlugin::preserveInitializerSections(
837     jitlink::LinkGraph &G, MaterializationResponsibility &MR) {
838 
839   if (const auto &InitSymName = MR.getInitializerSymbol()) {
840 
841     jitlink::Symbol *InitSym = nullptr;
842 
843     for (auto &InitSection : G.sections()) {
844       // Skip non-init sections.
845       if (!isCOFFInitializerSection(InitSection.getName()) ||
846           InitSection.empty())
847         continue;
848 
849       // Create the init symbol if it has not been created already and attach it
850       // to the first block.
851       if (!InitSym) {
852         auto &B = **InitSection.blocks().begin();
853         InitSym = &G.addDefinedSymbol(
854             B, 0, *InitSymName, B.getSize(), jitlink::Linkage::Strong,
855             jitlink::Scope::SideEffectsOnly, false, true);
856       }
857 
858       // Add keep-alive edges to anonymous symbols in all other init blocks.
859       for (auto *B : InitSection.blocks()) {
860         if (B == &InitSym->getBlock())
861           continue;
862 
863         auto &S = G.addAnonymousSymbol(*B, 0, B->getSize(), false, true);
864         InitSym->getBlock().addEdge(jitlink::Edge::KeepAlive, 0, S, 0);
865       }
866     }
867   }
868 
869   return Error::success();
870 }
871 
872 Error COFFPlatform::COFFPlatformPlugin::
873     registerObjectPlatformSectionsInBootstrap(jitlink::LinkGraph &G,
874                                               JITDylib &JD) {
875   std::lock_guard<std::mutex> Lock(CP.PlatformMutex);
876   auto HeaderAddr = CP.JITDylibToHeaderAddr[&JD];
877   COFFObjectSectionsMap ObjSecs;
878   for (auto &S : G.sections()) {
879     jitlink::SectionRange Range(S);
880     if (Range.getSize())
881       ObjSecs.push_back(std::make_pair(S.getName().str(), Range.getRange()));
882   }
883 
884   G.allocActions().push_back(
885       {{},
886        cantFail(
887            WrapperFunctionCall::Create<SPSCOFFDeregisterObjectSectionsArgs>(
888                CP.orc_rt_coff_deregister_object_sections, HeaderAddr,
889                ObjSecs))});
890 
891   auto &BState = CP.JDBootstrapStates[&JD];
892   BState.ObjectSectionsMaps.push_back(std::move(ObjSecs));
893 
894   // Collect static initializers
895   for (auto &S : G.sections())
896     if (isCOFFInitializerSection(S.getName()))
897       for (auto *B : S.blocks()) {
898         if (B->edges_empty())
899           continue;
900         for (auto &E : B->edges())
901           BState.Initializers.push_back(std::make_pair(
902               S.getName().str(), E.getTarget().getAddress() + E.getAddend()));
903       }
904 
905   return Error::success();
906 }
907 
908 } // End namespace orc.
909 } // End namespace llvm.
910