//===------ 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/Shared/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(JITDylib &JD, ResourceKey)>;

  using HandleTransferFunction =
      unique_function<void(JITDylib &JD, ResourceKey, ResourceKey)>;

  using RecordedResourcesMap = DenseMap<ResourceKey, ResourceT>;

  SimpleResourceManager(ExecutionSession &ES) : ES(ES) {
    HandleRemove = [&](JITDylib &JD, ResourceKey K) -> Error {
      ES.runSessionLocked([&] { removeResource(JD, K); });
      return Error::success();
    };

    HandleTransfer = [this](JITDylib &JD, ResourceKey DstKey,
                            ResourceKey SrcKey) {
      transferResources(JD, 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(JITDylib &JD, ResourceKey K) { Resources.erase(K); }

  /// Transfer resources from DstKey to SrcKey.
  template <typename MergeOp = std::plus<ResourceT>>
  void transferResources(JITDylib &JD, 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(JITDylib &JD, ResourceKey K) override {
    return HandleRemove(JD, K);
  }

  void handleTransferResources(JITDylib &JD, ResourceKey DstKey,
                               ResourceKey SrcKey) override {
    HandleTransfer(JD, 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([&](JITDylib &JD, ResourceKey K) -> Error {
    ResourceManagerGotRemove = true;
    EXPECT_EQ(SRM.getRecordedResources().size(), 0U)
        << "Unexpected resources recorded";
    SRM.removeResource(JD, 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([&](JITDylib &JD, 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(JD, 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([&](JITDylib &JD, ResourceKey K) -> Error {
    ResourceManagerGotRemove = true;
    EXPECT_EQ(SRM.getRecordedResources().size(), 0U)
        << "Unexpected resources recorded";
    SRM.removeResource(JD, 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(
      [&](JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {
        ResourceManagerGotTransfer = true;
        auto &RR = SRM.getRecordedResources();
        EXPECT_EQ(RR.size(), 0U) << "Expected no resources recorded yet";
      });

  auto MakeMU = [&](SymbolStringPtr Name, ExecutorSymbolDef 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(
      [&](JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {
        ResourceManagerGotTransfer = true;
        SRM.transferResources(JD, DstKey, SrcKey);
      });

  auto MakeMU = [&](SymbolStringPtr Name, ExecutorSymbolDef 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(
      [&](JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {
        ResourceManagerGotTransfer = true;
        SRM.transferResources(JD, 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