CallbackManager.h

Go to the documentation of this file.

 
#pragma once
 
#include <tvm/internal/PairElementToken.h>
 
#include <algorithm>
#include <functional>
#include <vector>
 
namespace tvm::internal
{
class CallbackManager
{
public:
  using Callback = std::function<void()>;
  using TCPair = std::pair<PairElementToken, Callback>;
 
  CallbackManager() = default;
  CallbackManager(const CallbackManager &) = delete;
  CallbackManager & operator=(const CallbackManager &) = delete;
 
  virtual ~CallbackManager() = default;
 
  PairElementTokenHandle registerCallback(std::function<void()> c);
 
  void unregisterCallback(const PairElementToken & t);
 
  void run();
 
private:
  std::vector<TCPair> callbacks_;
};
 
inline PairElementTokenHandle CallbackManager::registerCallback(std::function<void()> c)
{
  callbacks_.emplace_back(PairElementToken(), c);
  return PairElementTokenHandle(callbacks_.back().first);
}
 
inline void CallbackManager::unregisterCallback(const PairElementToken & t)
{
  auto it =
      std::find_if(callbacks_.begin(), callbacks_.end(), [&](const TCPair & p) { return p.first.isPairedWith(t); });
  if(it != callbacks_.end())
    callbacks_.erase(it);
  else
    throw std::runtime_error("Element to unregister not found.");
}
 
inline void CallbackManager::run()
{
  for(size_t i = 0; i < callbacks_.size();)
  {
    if(callbacks_[i].first.isPaired())
    {
      std::invoke(callbacks_[i].second);
      ++i;
    }
    else
    {
      callbacks_.erase(callbacks_.begin() + i);
    }
  }
}
 
} // namespace tvm::internal