text: //===------- ObjectLinkingLayer.cpp - JITLink backed ORC ObjectLayer ------===//
  //
  // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
  // See https://llvm.org/LICENSE.txt for license information.
  // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
  //
  //===----------------------------------------------------------------------===//
  
  #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
  
  #include "llvm/ADT/Optional.h"
  #include "llvm/ExecutionEngine/JITLink/EHFrameSupport.h"
  
  #include <vector>
  
  #define DEBUG_TYPE "orc"
  
  using namespace llvm;
  using namespace llvm::jitlink;
  using namespace llvm::orc;
  
  namespace llvm {
  namespace orc {
  
  class ObjectLinkingLayerJITLinkContext final : public JITLinkContext {
  public:
    ObjectLinkingLayerJITLinkContext(ObjectLinkingLayer &Layer,
                                     MaterializationResponsibility MR,
                                     std::unique_ptr<MemoryBuffer> ObjBuffer)
        : Layer(Layer), MR(std::move(MR)), ObjBuffer(std::move(ObjBuffer)) {}
  
    ~ObjectLinkingLayerJITLinkContext() {
      // If there is an object buffer return function then use it to
      // return ownership of the buffer.
      if (Layer.ReturnObjectBuffer)
        Layer.ReturnObjectBuffer(std::move(ObjBuffer));
    }
  
    JITLinkMemoryManager &getMemoryManager() override { return Layer.MemMgr; }
  
    MemoryBufferRef getObjectBuffer() const override {
      return ObjBuffer->getMemBufferRef();
    }
  
    void notifyFailed(Error Err) override {
      Layer.getExecutionSession().reportError(std::move(Err));
      MR.failMaterialization();
    }
  
    void lookup(const LookupMap &Symbols,
                std::unique_ptr<JITLinkAsyncLookupContinuation> LC) override {
  
      JITDylibSearchOrder SearchOrder;
      MR.getTargetJITDylib().withSearchOrderDo(
          [&](const JITDylibSearchOrder &O) { SearchOrder = O; });
  
      auto &ES = Layer.getExecutionSession();
  
      SymbolLookupSet LookupSet;
      for (auto &KV : Symbols) {
        orc::SymbolLookupFlags LookupFlags;
        switch (KV.second) {
        case jitlink::SymbolLookupFlags::RequiredSymbol:
          LookupFlags = orc::SymbolLookupFlags::RequiredSymbol;
          break;
        case jitlink::SymbolLookupFlags::WeaklyReferencedSymbol:
          LookupFlags = orc::SymbolLookupFlags::WeaklyReferencedSymbol;
          break;
        }
        LookupSet.add(ES.intern(KV.first), LookupFlags);
      }
  
      // OnResolve -- De-intern the symbols and pass the result to the linker.
      auto OnResolve = [this, LookupContinuation = std::move(LC)](
                           Expected<SymbolMap> Result) mutable {
        auto Main = Layer.getExecutionSession().intern("_main");
        if (!Result)
          LookupContinuation->run(Result.takeError());
        else {
          AsyncLookupResult LR;
          for (auto &KV : *Result)
            LR[*KV.first] = KV.second;
          LookupContinuation->run(std::move(LR));
        }
      };
  
      ES.lookup(LookupKind::Static, SearchOrder, std::move(LookupSet),
                SymbolState::Resolved, std::move(OnResolve),
                [this](const SymbolDependenceMap &Deps) {
                  registerDependencies(Deps);
                });
    }
  
    void notifyResolved(LinkGraph &G) override {
      auto &ES = Layer.getExecutionSession();
  
      SymbolFlagsMap ExtraSymbolsToClaim;
      bool AutoClaim = Layer.AutoClaimObjectSymbols;
  
      SymbolMap InternedResult;
      for (auto *Sym : G.defined_symbols())
        if (Sym->hasName() && Sym->getScope() != Scope::Local) {
          auto InternedName = ES.intern(Sym->getName());
          JITSymbolFlags Flags;
  
          if (Sym->isCallable())
            Flags |= JITSymbolFlags::Callable;
          if (Sym->getScope() == Scope::Default)
            Flags |= JITSymbolFlags::Exported;
  
          InternedResult[InternedName] =
              JITEvaluatedSymbol(Sym->getAddress(), Flags);
          if (AutoClaim && !MR.getSymbols().count(InternedName)) {
            assert(!ExtraSymbolsToClaim.count(InternedName) &&
                   "Duplicate symbol to claim?");
            ExtraSymbolsToClaim[InternedName] = Flags;
          }
        }
  
      for (auto *Sym : G.absolute_symbols())
        if (Sym->hasName()) {
          auto InternedName = ES.intern(Sym->getName());
          JITSymbolFlags Flags;
          Flags |= JITSymbolFlags::Absolute;
          if (Sym->isCallable())
            Flags |= JITSymbolFlags::Callable;
          if (Sym->getLinkage() == Linkage::Weak)
            Flags |= JITSymbolFlags::Weak;
          InternedResult[InternedName] =
              JITEvaluatedSymbol(Sym->getAddress(), Flags);
          if (AutoClaim && !MR.getSymbols().count(InternedName)) {
            assert(!ExtraSymbolsToClaim.count(InternedName) &&
                   "Duplicate symbol to claim?");
            ExtraSymbolsToClaim[InternedName] = Flags;
          }
        }
  
      if (!ExtraSymbolsToClaim.empty())
        if (auto Err = MR.defineMaterializing(ExtraSymbolsToClaim))
          return notifyFailed(std::move(Err));
      if (auto Err = MR.notifyResolved(InternedResult)) {
        Layer.getExecutionSession().reportError(std::move(Err));
        MR.failMaterialization();
        return;
      }
      Layer.notifyLoaded(MR);
    }
  
    void notifyFinalized(
        std::unique_ptr<JITLinkMemoryManager::Allocation> A) override {
      if (auto Err = Layer.notifyEmitted(MR, std::move(A))) {
        Layer.getExecutionSession().reportError(std::move(Err));
        MR.failMaterialization();
        return;
      }
      if (auto Err = MR.notifyEmitted()) {
        Layer.getExecutionSession().reportError(std::move(Err));
        MR.failMaterialization();
      }
    }
  
    LinkGraphPassFunction getMarkLivePass(const Triple &TT) const override {
      return [this](LinkGraph &G) { return markResponsibilitySymbolsLive(G); };
    }
  
    Error modifyPassConfig(const Triple &TT, PassConfiguration &Config) override {
      // Add passes to mark duplicate defs as should-discard, and to walk the
      // link graph to build the symbol dependence graph.
      Config.PrePrunePasses.push_back(
          [this](LinkGraph &G) { return externalizeWeakAndCommonSymbols(G); });
      Config.PostPrunePasses.push_back(
          [this](LinkGraph &G) { return computeNamedSymbolDependencies(G); });
  
      Layer.modifyPassConfig(MR, TT, Config);
  
      return Error::success();
    }
  
  private:
    using AnonToNamedDependenciesMap = DenseMap<const Symbol *, SymbolNameSet>;
  
    Error externalizeWeakAndCommonSymbols(LinkGraph &G) {
      auto &ES = Layer.getExecutionSession();
      for (auto *Sym : G.defined_symbols())
        if (Sym->hasName() && Sym->getLinkage() == Linkage::Weak) {
          if (!MR.getSymbols().count(ES.intern(Sym->getName())))
            G.makeExternal(*Sym);
        }
  
      for (auto *Sym : G.absolute_symbols())
        if (Sym->hasName() && Sym->getLinkage() == Linkage::Weak) {
          if (!MR.getSymbols().count(ES.intern(Sym->getName())))
            G.makeExternal(*Sym);
        }
  
      return Error::success();
    }
  
    Error markResponsibilitySymbolsLive(LinkGraph &G) const {
      auto &ES = Layer.getExecutionSession();
      for (auto *Sym : G.defined_symbols())
        if (Sym->hasName() && MR.getSymbols().count(ES.intern(Sym->getName())))
          Sym->setLive(true);
      return Error::success();
    }
  
    Error computeNamedSymbolDependencies(LinkGraph &G) {
      auto &ES = MR.getTargetJITDylib().getExecutionSession();
      auto AnonDeps = computeAnonDeps(G);
  
      for (auto *Sym : G.defined_symbols()) {
  
        // Skip anonymous and non-global atoms: we do not need dependencies for
        // these.
        if (Sym->getScope() == Scope::Local)
          continue;
  
        auto SymName = ES.intern(Sym->getName());
        SymbolNameSet &SymDeps = NamedSymbolDeps[SymName];
  
        for (auto &E : Sym->getBlock().edges()) {
          auto &TargetSym = E.getTarget();
  
          if (TargetSym.getScope() != Scope::Local)
            SymDeps.insert(ES.intern(TargetSym.getName()));
          else {
            assert(TargetSym.isDefined() &&
                   "Anonymous/local symbols must be defined");
            auto I = AnonDeps.find(&TargetSym);
            if (I != AnonDeps.end())
              for (auto &S : I->second)
                SymDeps.insert(S);
          }
        }
      }
  
      return Error::success();
    }
  
    AnonToNamedDependenciesMap computeAnonDeps(LinkGraph &G) {
  
      auto &ES = MR.getTargetJITDylib().getExecutionSession();
      AnonToNamedDependenciesMap DepMap;
  
      // For all anonymous symbols:
      // (1) Add their named dependencies.
      // (2) Add them to the worklist for further iteration if they have any
      //     depend on any other anonymous symbols.
      struct WorklistEntry {
        WorklistEntry(Symbol *Sym, DenseSet<Symbol *> SymAnonDeps)
            : Sym(Sym), SymAnonDeps(std::move(SymAnonDeps)) {}
  
        Symbol *Sym = nullptr;
        DenseSet<Symbol *> SymAnonDeps;
      };
      std::vector<WorklistEntry> Worklist;
      for (auto *Sym : G.defined_symbols())
        if (!Sym->hasName()) {
          auto &SymNamedDeps = DepMap[Sym];
          DenseSet<Symbol *> SymAnonDeps;
  
          for (auto &E : Sym->getBlock().edges()) {
            auto &TargetSym = E.getTarget();
            if (TargetSym.hasName())
              SymNamedDeps.insert(ES.intern(TargetSym.getName()));
            else {
              assert(TargetSym.isDefined() &&
                     "Anonymous symbols must be defined");
              SymAnonDeps.insert(&TargetSym);
            }
          }
  
          if (!SymAnonDeps.empty())
            Worklist.push_back(WorklistEntry(Sym, std::move(SymAnonDeps)));
        }
  
      // Loop over all anonymous symbols with anonymous dependencies, propagating
      // their respective *named* dependencies. Iterate until we hit a stable
      // state.
      bool Changed;
      do {
        Changed = false;
        for (auto &WLEntry : Worklist) {
          auto *Sym = WLEntry.Sym;
          auto &SymNamedDeps = DepMap[Sym];
          auto &SymAnonDeps = WLEntry.SymAnonDeps;
  
          for (auto *TargetSym : SymAnonDeps) {
            auto I = DepMap.find(TargetSym);
            if (I != DepMap.end())
              for (const auto &S : I->second)
                Changed |= SymNamedDeps.insert(S).second;
          }
        }
      } while (Changed);
  
      return DepMap;
    }
  
    void registerDependencies(const SymbolDependenceMap &QueryDeps) {
      for (auto &NamedDepsEntry : NamedSymbolDeps) {
        auto &Name = NamedDepsEntry.first;
        auto &NameDeps = NamedDepsEntry.second;
        SymbolDependenceMap SymbolDeps;
  
        for (const auto &QueryDepsEntry : QueryDeps) {
          JITDylib &SourceJD = *QueryDepsEntry.first;
          const SymbolNameSet &Symbols = QueryDepsEntry.second;
          auto &DepsForJD = SymbolDeps[&SourceJD];
  
          for (const auto &S : Symbols)
            if (NameDeps.count(S))
              DepsForJD.insert(S);
  
          if (DepsForJD.empty())
            SymbolDeps.erase(&SourceJD);
        }
  
        MR.addDependencies(Name, SymbolDeps);
      }
    }
  
    ObjectLinkingLayer &Layer;
    MaterializationResponsibility MR;
    std::unique_ptr<MemoryBuffer> ObjBuffer;
    DenseMap<SymbolStringPtr, SymbolNameSet> NamedSymbolDeps;
  };
  
  ObjectLinkingLayer::Plugin::~Plugin() {}
  
  ObjectLinkingLayer::ObjectLinkingLayer(ExecutionSession &ES,
                                         JITLinkMemoryManager &MemMgr)
      : ObjectLayer(ES), MemMgr(MemMgr) {}
  
  ObjectLinkingLayer::~ObjectLinkingLayer() {
    if (auto Err = removeAllModules())
      getExecutionSession().reportError(std::move(Err));
  }
  
  void ObjectLinkingLayer::emit(MaterializationResponsibility R,
                                std::unique_ptr<MemoryBuffer> O) {
    assert(O && "Object must not be null");
    jitLink(std::make_unique<ObjectLinkingLayerJITLinkContext>(
        *this, std::move(R), std::move(O)));
  }
  
  void ObjectLinkingLayer::modifyPassConfig(MaterializationResponsibility &MR,
                                            const Triple &TT,
                                            PassConfiguration &PassConfig) {
    for (auto &P : Plugins)
      P->modifyPassConfig(MR, TT, PassConfig);
  }
  
  void ObjectLinkingLayer::notifyLoaded(MaterializationResponsibility &MR) {
    for (auto &P : Plugins)
      P->notifyLoaded(MR);
  }
  
  Error ObjectLinkingLayer::notifyEmitted(MaterializationResponsibility &MR,
                                          AllocPtr Alloc) {
    Error Err = Error::success();
    for (auto &P : Plugins)
      Err = joinErrors(std::move(Err), P->notifyEmitted(MR));
  
    if (Err)
      return Err;
  
    {
      std::lock_guard<std::mutex> Lock(LayerMutex);
      UntrackedAllocs.push_back(std::move(Alloc));
    }
  
    return Error::success();
  }
  
  Error ObjectLinkingLayer::removeModule(VModuleKey K) {
    Error Err = Error::success();
  
    for (auto &P : Plugins)
      Err = joinErrors(std::move(Err), P->notifyRemovingModule(K));
  
    AllocPtr Alloc;
  
    {
      std::lock_guard<std::mutex> Lock(LayerMutex);
      auto AllocItr = TrackedAllocs.find(K);
      Alloc = std::move(AllocItr->second);
      TrackedAllocs.erase(AllocItr);
    }
  
    assert(Alloc && "No allocation for key K");
  
    return joinErrors(std::move(Err), Alloc->deallocate());
  }
  
  Error ObjectLinkingLayer::removeAllModules() {
  
    Error Err = Error::success();
  
    for (auto &P : Plugins)
      Err = joinErrors(std::move(Err), P->notifyRemovingAllModules());
  
    std::vector<AllocPtr> Allocs;
    {
      std::lock_guard<std::mutex> Lock(LayerMutex);
      Allocs = std::move(UntrackedAllocs);
  
      for (auto &KV : TrackedAllocs)
        Allocs.push_back(std::move(KV.second));
  
      TrackedAllocs.clear();
    }
  
    while (!Allocs.empty()) {
      Err = joinErrors(std::move(Err), Allocs.back()->deallocate());
      Allocs.pop_back();
    }
  
    return Err;
  }
  
  EHFrameRegistrationPlugin::EHFrameRegistrationPlugin(
      EHFrameRegistrar &Registrar)
      : Registrar(Registrar) {}
  
  void EHFrameRegistrationPlugin::modifyPassConfig(
      MaterializationResponsibility &MR, const Triple &TT,
      PassConfiguration &PassConfig) {
    assert(!InProcessLinks.count(&MR) && "Link for MR already being tracked?");
  
    PassConfig.PostFixupPasses.push_back(
        createEHFrameRecorderPass(TT, [this, &MR](JITTargetAddress Addr,
                                                  size_t Size) {
          if (Addr)
            InProcessLinks[&MR] = { Addr, Size };
        }));
  }
  
  Error EHFrameRegistrationPlugin::notifyEmitted(
      MaterializationResponsibility &MR) {
  
    auto EHFrameRangeItr = InProcessLinks.find(&MR);
    if (EHFrameRangeItr == InProcessLinks.end())
      return Error::success();
  
    auto EHFrameRange = EHFrameRangeItr->second;
    assert(EHFrameRange.Addr &&
           "eh-frame addr to register can not be null");
  
    InProcessLinks.erase(EHFrameRangeItr);
    if (auto Key = MR.getVModuleKey())
      TrackedEHFrameRanges[Key] = EHFrameRange;
    else
      UntrackedEHFrameRanges.push_back(EHFrameRange);
  
    return Registrar.registerEHFrames(EHFrameRange.Addr, EHFrameRange.Size);
  }
  
  Error EHFrameRegistrationPlugin::notifyRemovingModule(VModuleKey K) {
    auto EHFrameRangeItr = TrackedEHFrameRanges.find(K);
    if (EHFrameRangeItr == TrackedEHFrameRanges.end())
      return Error::success();
  
    auto EHFrameRange = EHFrameRangeItr->second;
    assert(EHFrameRange.Addr && "Tracked eh-frame range must not be null");
  
    TrackedEHFrameRanges.erase(EHFrameRangeItr);
  
    return Registrar.deregisterEHFrames(EHFrameRange.Addr, EHFrameRange.Size);
  }
  
  Error EHFrameRegistrationPlugin::notifyRemovingAllModules() {
  
    std::vector<EHFrameRange> EHFrameRanges =
      std::move(UntrackedEHFrameRanges);
    EHFrameRanges.reserve(EHFrameRanges.size() + TrackedEHFrameRanges.size());
  
    for (auto &KV : TrackedEHFrameRanges)
      EHFrameRanges.push_back(KV.second);
  
    TrackedEHFrameRanges.clear();
  
    Error Err = Error::success();
  
    while (!EHFrameRanges.empty()) {
      auto EHFrameRange = EHFrameRanges.back();
      assert(EHFrameRange.Addr && "Untracked eh-frame range must not be null");
      EHFrameRanges.pop_back();
      Err = joinErrors(std::move(Err),
                       Registrar.deregisterEHFrames(EHFrameRange.Addr,
                                                    EHFrameRange.Size));
    }
  
    return Err;
  }
  
  } // End namespace orc.
  } // End namespace llvm.
  