wrappers.cc

#include <utility>
#include <iostream>
#include <atomic>
#include <set>

class OnceWatcher {
 public:
  virtual void Called() = 0;
};

template <typename T>
class Once {
 public:
  explicit Once(T func) : impl_(std::make_shared<Impl>(func, nullptr)) {}
  explicit Once(T func, OnceWatcher* watcher) : impl_(std::make_shared<Impl>(func, watcher)) {}

  template <typename... ArgType>
  auto operator()(ArgType&&... args) {
    return impl_->operator()(std::forward<ArgType>(args)...);
  }

  bool HasBeenCalled() { return impl_->called_; }

 private:
  class Impl {
   public:
    explicit Impl(T func, OnceWatcher* watcher) : func_(func), watcher_(watcher) {}

    template <typename... ArgType>
    auto operator()(ArgType&&... args) {
      if (!flag_.test_and_set()) {
        called_ = true;
        if (watcher_ != nullptr) {
          watcher_->Called();
        }
        return func_(std::forward<ArgType>(args)...);
      }
      return;
    }

    T func_;
    OnceWatcher* watcher_;
    bool called_ = false;
    std::atomic_flag flag_ = ATOMIC_FLAG_INIT;
  };
  std::shared_ptr<Impl> impl_;
};

class CallbackWrapperBase {
 public:
  virtual void Abort() = 0;
};

class CallbackRegistry {
 public:
  virtual void Register(CallbackWrapperBase* callback) = 0;
  virtual void Unregister(CallbackWrapperBase* callback) = 0;
};

template <typename T>
class CallbackWrapper {
 public:
  friend class CallbackHolder;
  using ErrorHandler = std::function<void(T&)>;
  explicit CallbackWrapper(T func, ErrorHandler error_handler, CallbackRegistry* registry)
      : impl_(std::make_shared<Impl>(func, error_handler, registry)) {}
  template <typename... ArgType>
  auto operator()(ArgType&&... args) const {
    return (*impl_)(std::forward<ArgType>(args)...);
  }

 private:
  class Impl : public CallbackWrapperBase, OnceWatcher {
   public:
    Impl(T func, ErrorHandler error_handler, CallbackRegistry* registry)
        : once_(func), error_handler_(error_handler), registry_(registry) {
          registry_->Register(this);
        }
    ~Impl() { Abort(); }
    void Abort() override {
      if (!once_.HasBeenCalled()) {
        T f = once_;
        error_handler_(f);
      }
    }

    void Called() override {
      registry_->Unregister(this);
    }

    template <typename... ArgType>
    void operator()(ArgType&&... args) {
      once_(std::forward<ArgType>(args)...);
    }
    Once<T> once_;
    ErrorHandler error_handler_;
    CallbackRegistry* registry_;
  };

  std::shared_ptr<Impl> impl_;
};

class CallbackHolder : public CallbackRegistry {
 public:
  template <typename T>
  CallbackWrapper<T> WrapCallback(T func, typename CallbackWrapper<T>::ErrorHandler error_handler) {
    CallbackWrapper<T> wrapper(func, error_handler, this);
    return std::move(wrapper);
  }

  void Register(CallbackWrapperBase* callback) override {
    wrappers_.insert(callback);
  }

  void Unregister(CallbackWrapperBase* callback) override {
    wrappers_.erase(callback);
  }

  void AbortAll() {
    std::set<CallbackWrapperBase*> wrappers;
    wrappers_.swap(wrappers);
    for (auto wrapper : wrappers) {
      wrapper->Abort();
    }
  }

 private:
  std::set<CallbackWrapperBase*> wrappers_;
};

typedef std::function<void(const char*)> Greeter;

int main(int argc, char** argv) {
  CallbackHolder holder;

  Greeter greeter = [](const char* x) { std::cout << "Hello, " << x << std::endl; };
  auto x = holder.WrapCallback(greeter, [](Greeter& func) { func("FAIL"); });
  // x("World");

  Greeter informal = [](const char* x) { std::cout << "Hi, " << x << std::endl; };
  auto y = holder.WrapCallback(informal, [](Greeter& func) { func("SUCK"); });
  y("Globe");

  std::cout << "Aborting all" << std::endl;
  holder.AbortAll();

  std::cout << "Done" << std::endl;

}