text: //===------ ResourceTrackerTest.cpp - Unit tests ResourceTracker API
  //-------===//
  //
  // 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 "OrcTestCommon.h"
  #include "llvm/ADT/FunctionExtras.h"
  #include "llvm/Config/llvm-config.h"
  #include "llvm/ExecutionEngine/Orc/Core.h"
  #include "llvm/ExecutionEngine/Orc/OrcError.h"
  #include "llvm/Testing/Support/Error.h"
  
  using namespace llvm;
  using namespace llvm::orc;
  
  class ResourceTrackerStandardTest : public CoreAPIsBasedStandardTest {};
  
  namespace {
  
  template <typename ResourceT = unsigned>
  class SimpleResourceManager : public ResourceManager {
  public:
    using HandleRemoveFunction = unique_function<Error(ResourceKey)>;
  
    using HandleTransferFunction =
        unique_function<void(ResourceKey, ResourceKey)>;
  
    using RecordedResourcesMap = DenseMap<ResourceKey, ResourceT>;
  
    SimpleResourceManager(ExecutionSession &ES) : ES(ES) {
      HandleRemove = [&](ResourceKey K) -> Error {
        ES.runSessionLocked([&] { removeResource(K); });
        return Error::success();
      };
  
      HandleTransfer = [this](ResourceKey DstKey, ResourceKey SrcKey) {
        transferResources(DstKey, SrcKey);
      };
  
      ES.registerResourceManager(*this);
    }
  
    SimpleResourceManager(const SimpleResourceManager &) = delete;
    SimpleResourceManager &operator=(const SimpleResourceManager &) = delete;
    SimpleResourceManager(SimpleResourceManager &&) = delete;
    SimpleResourceManager &operator=(SimpleResourceManager &&) = delete;
  
    ~SimpleResourceManager() { ES.deregisterResourceManager(*this); }
  
    /// Set the HandleRemove function object.
    void setHandleRemove(HandleRemoveFunction HandleRemove) {
      this->HandleRemove = std::move(HandleRemove);
    }
  
    /// Set the HandleTransfer function object.
    void setHandleTransfer(HandleTransferFunction HandleTransfer) {
      this->HandleTransfer = std::move(HandleTransfer);
    }
  
    /// Create an association between the given key and resource.
    template <typename MergeOp = std::plus<ResourceT>>
    void recordResource(ResourceKey K, ResourceT Val = ResourceT(),
                        MergeOp Merge = MergeOp()) {
      auto Tmp = std::move(Resources[K]);
      Resources[K] = Merge(std::move(Tmp), std::move(Val));
    }
  
    /// Remove the resource associated with K from the map if present.
    void removeResource(ResourceKey K) { Resources.erase(K); }
  
    /// Transfer resources from DstKey to SrcKey.
    template <typename MergeOp = std::plus<ResourceT>>
    void transferResources(ResourceKey DstKey, ResourceKey SrcKey,
                           MergeOp Merge = MergeOp()) {
      auto &DstResourceRef = Resources[DstKey];
      ResourceT DstResources;
      std::swap(DstResourceRef, DstResources);
  
      auto SI = Resources.find(SrcKey);
      assert(SI != Resources.end() && "No resource associated with SrcKey");
  
      DstResourceRef = Merge(std::move(DstResources), std::move(SI->second));
      Resources.erase(SI);
    }
  
    /// Return a reference to the Resources map.
    RecordedResourcesMap &getRecordedResources() { return Resources; }
    const RecordedResourcesMap &getRecordedResources() const { return Resources; }
  
    Error handleRemoveResources(ResourceKey K) override {
      return HandleRemove(K);
    }
  
    void handleTransferResources(ResourceKey DstKey,
                                 ResourceKey SrcKey) override {
      HandleTransfer(DstKey, SrcKey);
    }
  
    static void transferNotAllowed(ResourceKey DstKey, ResourceKey SrcKey) {
      llvm_unreachable("Resource transfer not allowed");
    }
  
  private:
    ExecutionSession &ES;
    HandleRemoveFunction HandleRemove;
    HandleTransferFunction HandleTransfer;
    RecordedResourcesMap Resources;
  };
  
  TEST_F(ResourceTrackerStandardTest,
         BasicDefineAndRemoveAllBeforeMaterializing) {
  
    bool ResourceManagerGotRemove = false;
    SimpleResourceManager<> SRM(ES);
    SRM.setHandleRemove([&](ResourceKey K) -> Error {
      ResourceManagerGotRemove = true;
      EXPECT_EQ(SRM.getRecordedResources().size(), 0U)
          << "Unexpected resources recorded";
      SRM.removeResource(K);
      return Error::success();
    });
  
    bool MaterializationUnitDestroyed = false;
    auto MU = std::make_unique<SimpleMaterializationUnit>(
        SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
        [&](std::unique_ptr<MaterializationResponsibility> R) {
          llvm_unreachable("Never called");
        },
        nullptr, SimpleMaterializationUnit::DiscardFunction(),
        [&]() { MaterializationUnitDestroyed = true; });
  
    auto RT = JD.createResourceTracker();
    cantFail(JD.define(std::move(MU), RT));
    cantFail(RT->remove());
    auto SymFlags = cantFail(ES.lookupFlags(
        LookupKind::Static,
        {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
        SymbolLookupSet(Foo, SymbolLookupFlags::WeaklyReferencedSymbol)));
  
    EXPECT_EQ(SymFlags.size(), 0U)
        << "Symbols should have been removed from the symbol table";
    EXPECT_TRUE(ResourceManagerGotRemove)
        << "ResourceManager did not receive handleRemoveResources";
    EXPECT_TRUE(MaterializationUnitDestroyed)
        << "MaterializationUnit not destroyed in response to removal";
  }
  
  TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllAfterMaterializing) {
  
    bool ResourceManagerGotRemove = false;
    SimpleResourceManager<> SRM(ES);
    SRM.setHandleRemove([&](ResourceKey K) -> Error {
      ResourceManagerGotRemove = true;
      EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
          << "Unexpected number of resources recorded";
      EXPECT_EQ(SRM.getRecordedResources().count(K), 1U)
          << "Unexpected recorded resource";
      SRM.removeResource(K);
      return Error::success();
    });
  
    auto MU = std::make_unique<SimpleMaterializationUnit>(
        SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
        [&](std::unique_ptr<MaterializationResponsibility> R) {
          cantFail(R->withResourceKeyDo(
              [&](ResourceKey K) { SRM.recordResource(K); }));
          cantFail(R->notifyResolved({{Foo, FooSym}}));
          cantFail(R->notifyEmitted());
        });
  
    auto RT = JD.createResourceTracker();
    cantFail(JD.define(std::move(MU), RT));
    cantFail(ES.lookup({&JD}, Foo));
    cantFail(RT->remove());
    auto SymFlags = cantFail(ES.lookupFlags(
        LookupKind::Static,
        {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
        SymbolLookupSet(Foo, SymbolLookupFlags::WeaklyReferencedSymbol)));
  
    EXPECT_EQ(SymFlags.size(), 0U)
        << "Symbols should have been removed from the symbol table";
    EXPECT_TRUE(ResourceManagerGotRemove)
        << "ResourceManager did not receive handleRemoveResources";
  }
  
  TEST_F(ResourceTrackerStandardTest, BasicDefineAndRemoveAllWhileMaterializing) {
  
    bool ResourceManagerGotRemove = false;
    SimpleResourceManager<> SRM(ES);
    SRM.setHandleRemove([&](ResourceKey K) -> Error {
      ResourceManagerGotRemove = true;
      EXPECT_EQ(SRM.getRecordedResources().size(), 0U)
          << "Unexpected resources recorded";
      SRM.removeResource(K);
      return Error::success();
    });
  
    std::unique_ptr<MaterializationResponsibility> MR;
    auto MU = std::make_unique<SimpleMaterializationUnit>(
        SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
        [&](std::unique_ptr<MaterializationResponsibility> R) {
          MR = std::move(R);
        });
  
    auto RT = JD.createResourceTracker();
    cantFail(JD.define(std::move(MU), RT));
  
    ES.lookup(
        LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Foo),
        SymbolState::Ready,
        [](Expected<SymbolMap> Result) {
          EXPECT_THAT_EXPECTED(Result, Failed<FailedToMaterialize>())
              << "Lookup failed unexpectedly";
        },
        NoDependenciesToRegister);
  
    cantFail(RT->remove());
    auto SymFlags = cantFail(ES.lookupFlags(
        LookupKind::Static,
        {{&JD, JITDylibLookupFlags::MatchExportedSymbolsOnly}},
        SymbolLookupSet(Foo, SymbolLookupFlags::WeaklyReferencedSymbol)));
  
    EXPECT_EQ(SymFlags.size(), 0U)
        << "Symbols should have been removed from the symbol table";
    EXPECT_TRUE(ResourceManagerGotRemove)
        << "ResourceManager did not receive handleRemoveResources";
  
    EXPECT_THAT_ERROR(MR->withResourceKeyDo([](ResourceKey K) {
      ADD_FAILURE() << "Should not reach withResourceKeyDo body for removed key";
    }),
                      Failed<ResourceTrackerDefunct>())
        << "withResourceKeyDo on MR with removed tracker should have failed";
    EXPECT_THAT_ERROR(MR->notifyResolved({{Foo, FooSym}}),
                      Failed<ResourceTrackerDefunct>())
        << "notifyResolved on MR with removed tracker should have failed";
  
    MR->failMaterialization();
  }
  
  TEST_F(ResourceTrackerStandardTest, JITDylibClear) {
    SimpleResourceManager<> SRM(ES);
  
    // Add materializer for Foo.
    cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(
        SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
        [&](std::unique_ptr<MaterializationResponsibility> R) {
          cantFail(R->withResourceKeyDo(
              [&](ResourceKey K) { ++SRM.getRecordedResources()[K]; }));
          cantFail(R->notifyResolved({{Foo, FooSym}}));
          cantFail(R->notifyEmitted());
        })));
  
    // Add materializer for Bar.
    cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(
        SymbolFlagsMap({{Bar, BarSym.getFlags()}}),
        [&](std::unique_ptr<MaterializationResponsibility> R) {
          cantFail(R->withResourceKeyDo(
              [&](ResourceKey K) { ++SRM.getRecordedResources()[K]; }));
          cantFail(R->notifyResolved({{Bar, BarSym}}));
          cantFail(R->notifyEmitted());
        })));
  
    EXPECT_TRUE(SRM.getRecordedResources().empty())
        << "Expected no resources recorded yet.";
  
    cantFail(
        ES.lookup(makeJITDylibSearchOrder(&JD), SymbolLookupSet({Foo, Bar})));
  
    auto JDResourceKey = JD.getDefaultResourceTracker()->getKeyUnsafe();
    EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
        << "Expected exactly one entry (for JD's ResourceKey)";
    EXPECT_EQ(SRM.getRecordedResources().count(JDResourceKey), 1U)
        << "Expected an entry for JD's ResourceKey";
    EXPECT_EQ(SRM.getRecordedResources()[JDResourceKey], 2U)
        << "Expected value of 2 for JD's ResourceKey "
           "(+1 for each of Foo and Bar)";
  
    cantFail(JD.clear());
  
    EXPECT_TRUE(SRM.getRecordedResources().empty())
        << "Expected no resources recorded after clear";
  }
  
  TEST_F(ResourceTrackerStandardTest,
         BasicDefineAndExplicitTransferBeforeMaterializing) {
  
    bool ResourceManagerGotTransfer = false;
    SimpleResourceManager<> SRM(ES);
    SRM.setHandleTransfer([&](ResourceKey DstKey, ResourceKey SrcKey) {
      ResourceManagerGotTransfer = true;
      auto &RR = SRM.getRecordedResources();
      EXPECT_EQ(RR.size(), 0U) << "Expected no resources recorded yet";
    });
  
    auto MakeMU = [&](SymbolStringPtr Name, JITEvaluatedSymbol Sym) {
      return std::make_unique<SimpleMaterializationUnit>(
          SymbolFlagsMap({{Name, Sym.getFlags()}}),
          [=, &SRM](std::unique_ptr<MaterializationResponsibility> R) {
            cantFail(R->withResourceKeyDo(
                [&](ResourceKey K) { SRM.recordResource(K); }));
            cantFail(R->notifyResolved({{Name, Sym}}));
            cantFail(R->notifyEmitted());
          });
    };
  
    auto FooRT = JD.createResourceTracker();
    cantFail(JD.define(MakeMU(Foo, FooSym), FooRT));
  
    auto BarRT = JD.createResourceTracker();
    cantFail(JD.define(MakeMU(Bar, BarSym), BarRT));
  
    BarRT->transferTo(*FooRT);
  
    EXPECT_TRUE(ResourceManagerGotTransfer)
        << "ResourceManager did not receive transfer";
    EXPECT_TRUE(BarRT->isDefunct()) << "BarRT should now be defunct";
  
    cantFail(
        ES.lookup(makeJITDylibSearchOrder({&JD}), SymbolLookupSet({Foo, Bar})));
  
    EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
        << "Expected exactly one entry (for FooRT's Key)";
    EXPECT_EQ(SRM.getRecordedResources().count(FooRT->getKeyUnsafe()), 1U)
        << "Expected an entry for FooRT's ResourceKey";
    EXPECT_EQ(SRM.getRecordedResources().count(BarRT->getKeyUnsafe()), 0U)
        << "Expected no entry for BarRT's ResourceKey";
  
    // We need to explicitly destroy FooRT or its resources will be implicitly
    // transferred to the default tracker triggering a second call to our
    // transfer function above (which expects only one call).
    cantFail(FooRT->remove());
  }
  
  TEST_F(ResourceTrackerStandardTest,
         BasicDefineAndExplicitTransferAfterMaterializing) {
  
    bool ResourceManagerGotTransfer = false;
    SimpleResourceManager<> SRM(ES);
    SRM.setHandleTransfer([&](ResourceKey DstKey, ResourceKey SrcKey) {
      ResourceManagerGotTransfer = true;
      SRM.transferResources(DstKey, SrcKey);
    });
  
    auto MakeMU = [&](SymbolStringPtr Name, JITEvaluatedSymbol Sym) {
      return std::make_unique<SimpleMaterializationUnit>(
          SymbolFlagsMap({{Name, Sym.getFlags()}}),
          [=, &SRM](std::unique_ptr<MaterializationResponsibility> R) {
            cantFail(R->withResourceKeyDo(
                [&](ResourceKey K) { SRM.recordResource(K, 1); }));
            cantFail(R->notifyResolved({{Name, Sym}}));
            cantFail(R->notifyEmitted());
          });
    };
  
    auto FooRT = JD.createResourceTracker();
    cantFail(JD.define(MakeMU(Foo, FooSym), FooRT));
  
    auto BarRT = JD.createResourceTracker();
    cantFail(JD.define(MakeMU(Bar, BarSym), BarRT));
  
    EXPECT_EQ(SRM.getRecordedResources().size(), 0U)
        << "Expected no recorded resources yet";
  
    cantFail(
        ES.lookup(makeJITDylibSearchOrder({&JD}), SymbolLookupSet({Foo, Bar})));
  
    EXPECT_EQ(SRM.getRecordedResources().size(), 2U)
        << "Expected recorded resources for both Foo and Bar";
  
    BarRT->transferTo(*FooRT);
  
    EXPECT_TRUE(ResourceManagerGotTransfer)
        << "ResourceManager did not receive transfer";
    EXPECT_TRUE(BarRT->isDefunct()) << "BarRT should now be defunct";
  
    EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
        << "Expected recorded resources for Foo only";
    EXPECT_EQ(SRM.getRecordedResources().count(FooRT->getKeyUnsafe()), 1U)
        << "Expected recorded resources for Foo";
    EXPECT_EQ(SRM.getRecordedResources()[FooRT->getKeyUnsafe()], 2U)
        << "Expected resources value for for Foo to be '2'";
  }
  
  TEST_F(ResourceTrackerStandardTest,
         BasicDefineAndExplicitTransferWhileMaterializing) {
  
    bool ResourceManagerGotTransfer = false;
    SimpleResourceManager<> SRM(ES);
    SRM.setHandleTransfer([&](ResourceKey DstKey, ResourceKey SrcKey) {
      ResourceManagerGotTransfer = true;
      SRM.transferResources(DstKey, SrcKey);
    });
  
    auto FooRT = JD.createResourceTracker();
    std::unique_ptr<MaterializationResponsibility> FooMR;
    cantFail(JD.define(std::make_unique<SimpleMaterializationUnit>(
                           SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
                           [&](std::unique_ptr<MaterializationResponsibility> R) {
                             FooMR = std::move(R);
                           }),
                       FooRT));
  
    auto BarRT = JD.createResourceTracker();
  
    ES.lookup(
        LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Foo),
        SymbolState::Ready,
        [](Expected<SymbolMap> Result) { cantFail(Result.takeError()); },
        NoDependenciesToRegister);
  
    cantFail(FooMR->withResourceKeyDo([&](ResourceKey K) {
      EXPECT_EQ(FooRT->getKeyUnsafe(), K)
          << "Expected FooRT's ResourceKey for Foo here";
      SRM.recordResource(K, 1);
    }));
  
    EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
        << "Expected one recorded resource here";
    EXPECT_EQ(SRM.getRecordedResources()[FooRT->getKeyUnsafe()], 1U)
        << "Expected Resource value for FooRT to be '1' here";
  
    FooRT->transferTo(*BarRT);
  
    EXPECT_TRUE(ResourceManagerGotTransfer)
        << "Expected resource manager to receive handleTransferResources call";
  
    cantFail(FooMR->withResourceKeyDo([&](ResourceKey K) {
      EXPECT_EQ(BarRT->getKeyUnsafe(), K)
          << "Expected BarRT's ResourceKey for Foo here";
      SRM.recordResource(K, 1);
    }));
  
    EXPECT_EQ(SRM.getRecordedResources().size(), 1U)
        << "Expected one recorded resource here";
    EXPECT_EQ(SRM.getRecordedResources().count(BarRT->getKeyUnsafe()), 1U)
        << "Expected RecordedResources to contain an entry for BarRT";
    EXPECT_EQ(SRM.getRecordedResources()[BarRT->getKeyUnsafe()], 2U)
        << "Expected Resource value for BarRT to be '2' here";
  
    cantFail(FooMR->notifyResolved({{Foo, FooSym}}));
    cantFail(FooMR->notifyEmitted());
  }
  
  } // namespace
  