1283625Sdim//===- LazyEmittingLayer.h - Lazily emit IR to lower JIT layers -*- C++ -*-===//
2283625Sdim//
3353358Sdim// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4353358Sdim// See https://llvm.org/LICENSE.txt for license information.
5353358Sdim// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6283625Sdim//
7283625Sdim//===----------------------------------------------------------------------===//
8283625Sdim//
9283625Sdim// Contains the definition for a lazy-emitting layer for the JIT.
10283625Sdim//
11283625Sdim//===----------------------------------------------------------------------===//
12283625Sdim
13283625Sdim#ifndef LLVM_EXECUTIONENGINE_ORC_LAZYEMITTINGLAYER_H
14283625Sdim#define LLVM_EXECUTIONENGINE_ORC_LAZYEMITTINGLAYER_H
15283625Sdim
16314564Sdim#include "llvm/ADT/STLExtras.h"
17314564Sdim#include "llvm/ADT/StringMap.h"
18314564Sdim#include "llvm/ADT/StringRef.h"
19314564Sdim#include "llvm/ExecutionEngine/JITSymbol.h"
20341825Sdim#include "llvm/ExecutionEngine/Orc/Core.h"
21283625Sdim#include "llvm/IR/GlobalValue.h"
22283625Sdim#include "llvm/IR/Mangler.h"
23283625Sdim#include "llvm/IR/Module.h"
24314564Sdim#include "llvm/Support/ErrorHandling.h"
25314564Sdim#include "llvm/Support/raw_ostream.h"
26314564Sdim#include <algorithm>
27314564Sdim#include <cassert>
28283625Sdim#include <list>
29314564Sdim#include <memory>
30314564Sdim#include <string>
31283625Sdim
32283625Sdimnamespace llvm {
33283625Sdimnamespace orc {
34283625Sdim
35341825Sdim/// Lazy-emitting IR layer.
36283625Sdim///
37353358Sdim///   This layer accepts LLVM IR Modules (via addModule) but does not
38353358Sdim/// immediately emit them the layer below. Instead, emission to the base layer
39321369Sdim/// is deferred until the first time the client requests the address (via
40321369Sdim/// JITSymbol::getAddress) for a symbol contained in this layer.
41283625Sdimtemplate <typename BaseLayerT> class LazyEmittingLayer {
42283625Sdimprivate:
43321369Sdim  class EmissionDeferredModule {
44283625Sdim  public:
45341825Sdim    EmissionDeferredModule(VModuleKey K, std::unique_ptr<Module> M)
46341825Sdim        : K(std::move(K)), M(std::move(M)) {}
47283625Sdim
48283625Sdim    JITSymbol find(StringRef Name, bool ExportedSymbolsOnly, BaseLayerT &B) {
49283625Sdim      switch (EmitState) {
50283625Sdim      case NotEmitted:
51283625Sdim        if (auto GV = searchGVs(Name, ExportedSymbolsOnly)) {
52314564Sdim          JITSymbolFlags Flags = JITSymbolFlags::fromGlobalValue(*GV);
53360784Sdim          auto GetAddress = [this, ExportedSymbolsOnly, Name = Name.str(),
54360784Sdim                             &B]() -> Expected<JITTargetAddress> {
55360784Sdim            if (this->EmitState == Emitting)
56360784Sdim              return 0;
57360784Sdim            else if (this->EmitState == NotEmitted) {
58360784Sdim              this->EmitState = Emitting;
59360784Sdim              if (auto Err = this->emitToBaseLayer(B))
60321369Sdim                return std::move(Err);
61360784Sdim              this->EmitState = Emitted;
62360784Sdim            }
63360784Sdim            if (auto Sym = B.findSymbolIn(K, Name, ExportedSymbolsOnly))
64360784Sdim              return Sym.getAddress();
65360784Sdim            else if (auto Err = Sym.takeError())
66360784Sdim              return std::move(Err);
67360784Sdim            else
68360784Sdim              llvm_unreachable("Successful symbol lookup should return "
69360784Sdim                               "definition address here");
70283625Sdim          };
71283625Sdim          return JITSymbol(std::move(GetAddress), Flags);
72283625Sdim        } else
73283625Sdim          return nullptr;
74283625Sdim      case Emitting:
75296417Sdim        // Calling "emit" can trigger a recursive call to 'find' (e.g. to check
76296417Sdim        // for pre-existing definitions of common-symbol), but any symbol in
77296417Sdim        // this module would already have been found internally (in the
78296417Sdim        // RuntimeDyld that did the lookup), so just return a nullptr here.
79283625Sdim        return nullptr;
80283625Sdim      case Emitted:
81341825Sdim        return B.findSymbolIn(K, Name, ExportedSymbolsOnly);
82283625Sdim      }
83283625Sdim      llvm_unreachable("Invalid emit-state.");
84283625Sdim    }
85283625Sdim
86322740Sdim    Error removeModuleFromBaseLayer(BaseLayerT& BaseLayer) {
87341825Sdim      return EmitState != NotEmitted ? BaseLayer.removeModule(K)
88322740Sdim                                     : Error::success();
89283625Sdim    }
90283625Sdim
91283625Sdim    void emitAndFinalize(BaseLayerT &BaseLayer) {
92283625Sdim      assert(EmitState != Emitting &&
93283625Sdim             "Cannot emitAndFinalize while already emitting");
94283625Sdim      if (EmitState == NotEmitted) {
95283625Sdim        EmitState = Emitting;
96341825Sdim        emitToBaseLayer(BaseLayer);
97283625Sdim        EmitState = Emitted;
98283625Sdim      }
99341825Sdim      BaseLayer.emitAndFinalize(K);
100283625Sdim    }
101283625Sdim
102283625Sdim  private:
103283625Sdim
104283625Sdim    const GlobalValue* searchGVs(StringRef Name,
105321369Sdim                                 bool ExportedSymbolsOnly) const {
106283625Sdim      // FIXME: We could clean all this up if we had a way to reliably demangle
107283625Sdim      //        names: We could just demangle name and search, rather than
108283625Sdim      //        mangling everything else.
109283625Sdim
110283625Sdim      // If we have already built the mangled name set then just search it.
111283625Sdim      if (MangledSymbols) {
112283625Sdim        auto VI = MangledSymbols->find(Name);
113283625Sdim        if (VI == MangledSymbols->end())
114283625Sdim          return nullptr;
115283625Sdim        auto GV = VI->second;
116283625Sdim        if (!ExportedSymbolsOnly || GV->hasDefaultVisibility())
117283625Sdim          return GV;
118283625Sdim        return nullptr;
119283625Sdim      }
120283625Sdim
121283625Sdim      // If we haven't built the mangled name set yet, try to build it. As an
122283625Sdim      // optimization this will leave MangledNames set to nullptr if we find
123283625Sdim      // Name in the process of building the set.
124283625Sdim      return buildMangledSymbols(Name, ExportedSymbolsOnly);
125283625Sdim    }
126283625Sdim
127341825Sdim    Error emitToBaseLayer(BaseLayerT &BaseLayer) {
128283625Sdim      // We don't need the mangled names set any more: Once we've emitted this
129283625Sdim      // to the base layer we'll just look for symbols there.
130283625Sdim      MangledSymbols.reset();
131341825Sdim      return BaseLayer.addModule(std::move(K), std::move(M));
132283625Sdim    }
133283625Sdim
134283625Sdim    // If the mangled name of the given GlobalValue matches the given search
135283625Sdim    // name (and its visibility conforms to the ExportedSymbolsOnly flag) then
136283625Sdim    // return the symbol. Otherwise, add the mangled name to the Names map and
137283625Sdim    // return nullptr.
138283625Sdim    const GlobalValue* addGlobalValue(StringMap<const GlobalValue*> &Names,
139283625Sdim                                      const GlobalValue &GV,
140283625Sdim                                      const Mangler &Mang, StringRef SearchName,
141283625Sdim                                      bool ExportedSymbolsOnly) const {
142283625Sdim      // Modules don't "provide" decls or common symbols.
143283625Sdim      if (GV.isDeclaration() || GV.hasCommonLinkage())
144283625Sdim        return nullptr;
145283625Sdim
146283625Sdim      // Mangle the GV name.
147283625Sdim      std::string MangledName;
148283625Sdim      {
149283625Sdim        raw_string_ostream MangledNameStream(MangledName);
150283625Sdim        Mang.getNameWithPrefix(MangledNameStream, &GV, false);
151283625Sdim      }
152283625Sdim
153283625Sdim      // Check whether this is the name we were searching for, and if it is then
154283625Sdim      // bail out early.
155283625Sdim      if (MangledName == SearchName)
156283625Sdim        if (!ExportedSymbolsOnly || GV.hasDefaultVisibility())
157283625Sdim          return &GV;
158283625Sdim
159283625Sdim      // Otherwise add this to the map for later.
160283625Sdim      Names[MangledName] = &GV;
161283625Sdim      return nullptr;
162283625Sdim    }
163283625Sdim
164283625Sdim    // Build the MangledSymbols map. Bails out early (with MangledSymbols left set
165283625Sdim    // to nullptr) if the given SearchName is found while building the map.
166283625Sdim    const GlobalValue* buildMangledSymbols(StringRef SearchName,
167283625Sdim                                           bool ExportedSymbolsOnly) const {
168283625Sdim      assert(!MangledSymbols && "Mangled symbols map already exists?");
169283625Sdim
170360784Sdim      auto Symbols = std::make_unique<StringMap<const GlobalValue*>>();
171283625Sdim
172321369Sdim      Mangler Mang;
173283625Sdim
174321369Sdim      for (const auto &GO : M->global_objects())
175309124Sdim          if (auto GV = addGlobalValue(*Symbols, GO, Mang, SearchName,
176283625Sdim                                       ExportedSymbolsOnly))
177283625Sdim            return GV;
178283625Sdim
179283625Sdim      MangledSymbols = std::move(Symbols);
180283625Sdim      return nullptr;
181283625Sdim    }
182283625Sdim
183321369Sdim    enum { NotEmitted, Emitting, Emitted } EmitState = NotEmitted;
184341825Sdim    VModuleKey K;
185341825Sdim    std::unique_ptr<Module> M;
186283625Sdim    mutable std::unique_ptr<StringMap<const GlobalValue*>> MangledSymbols;
187283625Sdim  };
188283625Sdim
189283625Sdim  BaseLayerT &BaseLayer;
190341825Sdim  std::map<VModuleKey, std::unique_ptr<EmissionDeferredModule>> ModuleMap;
191283625Sdim
192283625Sdimpublic:
193283625Sdim
194341825Sdim  /// Construct a lazy emitting layer.
195353358Sdim  LLVM_ATTRIBUTE_DEPRECATED(
196353358Sdim      LazyEmittingLayer(BaseLayerT &BaseLayer),
197353358Sdim      "ORCv1 layers (including LazyEmittingLayer) are deprecated. Please use "
198353358Sdim      "ORCv2, where lazy emission is the default");
199283625Sdim
200353358Sdim  /// Construct a lazy emitting layer.
201353358Sdim  LazyEmittingLayer(ORCv1DeprecationAcknowledgement, BaseLayerT &BaseLayer)
202353358Sdim      : BaseLayer(BaseLayer) {}
203353358Sdim
204341825Sdim  /// Add the given module to the lazy emitting layer.
205341825Sdim  Error addModule(VModuleKey K, std::unique_ptr<Module> M) {
206341825Sdim    assert(!ModuleMap.count(K) && "VModuleKey K already in use");
207341825Sdim    ModuleMap[K] =
208360784Sdim        std::make_unique<EmissionDeferredModule>(std::move(K), std::move(M));
209341825Sdim    return Error::success();
210283625Sdim  }
211283625Sdim
212341825Sdim  /// Remove the module represented by the given handle.
213283625Sdim  ///
214321369Sdim  ///   This method will free the memory associated with the given module, both
215321369Sdim  /// in this layer, and the base layer.
216341825Sdim  Error removeModule(VModuleKey K) {
217341825Sdim    auto I = ModuleMap.find(K);
218341825Sdim    assert(I != ModuleMap.end() && "VModuleKey K not valid here");
219341825Sdim    auto EDM = std::move(I.second);
220341825Sdim    ModuleMap.erase(I);
221341825Sdim    return EDM->removeModuleFromBaseLayer(BaseLayer);
222283625Sdim  }
223283625Sdim
224341825Sdim  /// Search for the given named symbol.
225283625Sdim  /// @param Name The name of the symbol to search for.
226283625Sdim  /// @param ExportedSymbolsOnly If true, search only for exported symbols.
227283625Sdim  /// @return A handle for the given named symbol, if it exists.
228283625Sdim  JITSymbol findSymbol(const std::string &Name, bool ExportedSymbolsOnly) {
229283625Sdim    // Look for the symbol among existing definitions.
230283625Sdim    if (auto Symbol = BaseLayer.findSymbol(Name, ExportedSymbolsOnly))
231283625Sdim      return Symbol;
232283625Sdim
233321369Sdim    // If not found then search the deferred modules. If any of these contain a
234283625Sdim    // definition of 'Name' then they will return a JITSymbol that will emit
235283625Sdim    // the corresponding module when the symbol address is requested.
236341825Sdim    for (auto &KV : ModuleMap)
237341825Sdim      if (auto Symbol = KV.second->find(Name, ExportedSymbolsOnly, BaseLayer))
238283625Sdim        return Symbol;
239283625Sdim
240283625Sdim    // If no definition found anywhere return a null symbol.
241283625Sdim    return nullptr;
242283625Sdim  }
243283625Sdim
244341825Sdim  /// Get the address of the given symbol in the context of the of
245341825Sdim  ///        compiled modules represented by the key K.
246341825Sdim  JITSymbol findSymbolIn(VModuleKey K, const std::string &Name,
247283625Sdim                         bool ExportedSymbolsOnly) {
248341825Sdim    assert(ModuleMap.count(K) && "VModuleKey K not valid here");
249341825Sdim    return ModuleMap[K]->find(Name, ExportedSymbolsOnly, BaseLayer);
250283625Sdim  }
251283625Sdim
252341825Sdim  /// Immediately emit and finalize the module represented by the given
253341825Sdim  ///        key.
254341825Sdim  Error emitAndFinalize(VModuleKey K) {
255341825Sdim    assert(ModuleMap.count(K) && "VModuleKey K not valid here");
256341825Sdim    return ModuleMap[K]->emitAndFinalize(BaseLayer);
257283625Sdim  }
258283625Sdim};
259283625Sdim
260353358Sdimtemplate <typename BaseLayerT>
261353358SdimLazyEmittingLayer<BaseLayerT>::LazyEmittingLayer(BaseLayerT &BaseLayer)
262353358Sdim    : BaseLayer(BaseLayer) {}
263353358Sdim
264314564Sdim} // end namespace orc
265314564Sdim} // end namespace llvm
266283625Sdim
267283625Sdim#endif // LLVM_EXECUTIONENGINE_ORC_LAZYEMITTINGLAYER_H
268