GeneratorValueStream.hpp

#pragma once

#include "BackendComponents.hpp"
#include "ValueStream.hpp"

#include <Radiant/BGThreadExecutor.hpp>

#include <folly/MoveWrapper.h>

#include <queue>
#include <vector>

namespace Showcase
{

  template <typename Generator>
  class GeneratorValueStream : public ValueStream<typename Generator::ValueType>
  {
  public:
    static_assert(IsGenerator<Generator>::value,
                  "Need to use Generator as a type parameter "
                  "for GeneratorValueStream.");
    typedef typename Generator::ValueType ValueType;
    using GeneratorPtr = std::shared_ptr<Generator>;

    GeneratorValueStream(std::shared_ptr<folly::Executor> generatorExecutor);

    virtual ~GeneratorValueStream();

    virtual folly::Future<ValueType> next() override;

    void setSynchronizationTest(std::function<bool(const ValueType&)> f);

    void setCallback(std::function<void(const ValueType&)> f);

    void addGenerator(GeneratorPtr generator);

    void addStream(std::shared_ptr<GeneratorValueStream<Generator>> other,
                   const QByteArray & key = QByteArray());

    void removeAddedGenerators(const QByteArray & key);

    bool isEmpty() const;

  private:
    void init();
    bool isPauseValue(const ValueType& v);
    void initFuture(size_t i);
    void futureDone(size_t i, const ValueType& ev);

    std::vector<GeneratorPtr> m_generators;
    std::shared_ptr<folly::Executor> m_executor;

    std::function<bool(const ValueType&)> m_syncTest;
    std::function<void(const ValueType&)> m_onValue;

    // Mutex needs to be recursive, Otherwise setting callback in
    // initFuture can deadlock if future is fulfilled immediately
    mutable std::recursive_mutex m_mutex;
    std::atomic<bool> m_init;

    std::vector<folly::Future<ValueType>> m_futures;
    std::vector<bool> m_isPaused;
    std::vector<bool> m_futuresPending;
    std::vector<bool> m_isDisabled;

    size_t m_promisesWithoutFutures;

    size_t m_generatorsPending;

    size_t m_futuresWaited;

    size_t m_disabledGenerators;

    std::queue<folly::Promise<ValueType>> m_promises;

    std::map<QByteArray, std::vector<size_t> > m_addedGenerators;
  };

  // --------------------------------------------------------------------------

  template <typename Generator>
  GeneratorValueStream<Generator>::GeneratorValueStream(std::shared_ptr<folly::Executor> exec)
    : m_executor(exec),
      m_syncTest(nullptr),
      m_init(false),
      m_promisesWithoutFutures(0),
      m_generatorsPending(0),
      m_futuresWaited(0),
      m_disabledGenerators(0)
  {
  }

  template<typename Generator>
  GeneratorValueStream<Generator>::~GeneratorValueStream()
  {
    if(m_promisesWithoutFutures > 0)
      return;

    for(size_t i = 0; i < m_futures.size(); ++i) {
      if(m_futuresPending[i] && !m_promises.empty()) {
        auto promise = folly::makeMoveWrapper(m_promises.front());
        m_promises.pop();

        std::move(m_futures[i]).thenValue([promise] (const ValueType& ev) mutable {
          promise->setValue(ev);
        });
      }
    }

    size_t i = 0;
    while(!m_promises.empty()) {
      auto promise = folly::makeMoveWrapper(m_promises.front());
      m_promises.pop();

      if (m_disabledGenerators < m_generators.size()) {
        while (m_isDisabled[i])
          i = (i+1) % m_generators.size();
      }

      auto gen = m_generators[i];
      auto future = gen->next(m_executor.get());

      std::move(future).thenValue([promise, gen] (ValueType ev) mutable {
        promise->setValue(ev);
      });

      i = (i+1) % m_generators.size();
    }
  }

  template <typename Generator>
  void GeneratorValueStream<Generator>::addGenerator(GeneratorPtr generator)
  {
    assert(!m_init);
    m_generators.push_back(generator);
  }

  template<typename Generator>
  void GeneratorValueStream<Generator>::futureDone(size_t i, const ValueType &ev)
  {
    {
      std::unique_lock<std::recursive_mutex> lck(m_mutex);
      if (m_isDisabled[i])
        return;
    }

    if(m_onValue)
      m_onValue(ev);

    std::unique_lock<std::recursive_mutex> lck(m_mutex);

    assert(m_futuresWaited > 0);
    --m_futuresWaited;

    bool fulfillPromise = true;

    assert(!m_isPaused[i]);
    assert(m_futuresPending[i]);
    m_futuresPending[i] = false;

    if(isPauseValue(ev)) {
      ++m_generatorsPending;
      m_isPaused[i] = true;
      fulfillPromise = (m_generators.size() - m_disabledGenerators) == m_generatorsPending;
    }

    bool setPromise = false;
    folly::Promise<ValueType> promise;

    if(fulfillPromise) {

      if(m_promisesWithoutFutures > 0 || m_promises.empty()) {
        folly::Promise<ValueType> promise;
        promise.setValue(ev);
        m_promises.emplace(std::move(promise));
        ++m_promisesWithoutFutures;
      } else {
        promise = std::move(m_promises.front());
        m_promises.pop();
        setPromise = true;
      }
    }

    if(isPauseValue(ev) && fulfillPromise) {
      m_generatorsPending = 0;
      std::fill(m_isPaused.begin(), m_isPaused.end(), false);
    }

    size_t promises = m_promises.size();
    if(promises > 0 && m_promisesWithoutFutures == 0) {
      if(!isPauseValue(ev)) {
        initFuture(i);
      } else {
        if(fulfillPromise) {
          for(size_t i = 0; i < m_generators.size(); ++i) {
            if (!m_isDisabled[i])
              initFuture(i);
          }
        } else {
          for(size_t i = 0; i < m_generators.size(); ++i) {
            if(!m_isPaused[i] && !m_futuresPending[i] && !m_isDisabled[i]) {
              initFuture(i);
            }
          }
        }
      }
    }
    lck.unlock();

    if(setPromise)
      promise.setValue(ev);
  }

  template<typename Generator>
  void GeneratorValueStream<Generator>::initFuture(size_t i)
  {
    assert(i < m_futures.size());

    auto ptr = this->shared_from_this();
    std::weak_ptr<GeneratorValueStream<Generator>> weak =
        std::static_pointer_cast<GeneratorValueStream<Generator>>(ptr);

    assert(i >= m_isPaused.size() || !m_isPaused[i]);

    assert(!m_futuresPending[i]);
    m_futuresPending[i] = true;

    auto cb = [weak, i] (const ValueType& ev) -> ValueType {
      auto me = weak.lock();
      if(me) {
        me->futureDone(i, ev);
      }
      return ev;
    };
    ++m_futuresWaited;

    m_futures[i] = m_generators[i]->next(m_executor.get()).thenValue(cb);
  }

  template<typename Generator>
  void GeneratorValueStream<Generator>::init()
  {
    // Is guaranteed to be executed when mutex is locked
    for(size_t i = 0; i < m_generators.size(); ++i) {
      m_futures.emplace_back(folly::makeFuture<ValueType>(
                folly::exception_wrapper(std::runtime_error("empty"))));
      m_isPaused.push_back(false);
      m_futuresPending.push_back(false);
      m_isDisabled.push_back(false);
      initFuture(i);
    }
  }

  template <typename Generator>
  folly::Future<typename GeneratorValueStream<Generator>::ValueType>
  GeneratorValueStream<Generator>::next()
  {
    if(!m_init) {
      std::lock_guard<std::recursive_mutex> g(m_mutex);

      if(!m_init) {
        init();
        m_init = true;
      }
    }

    assert(!m_generators.empty());
    if(m_generators.empty() || m_disabledGenerators == m_generators.size())
      return folly::makeFuture<ValueType>(
            folly::exception_wrapper(std::runtime_error("No generators")));

    {
      std::lock_guard<std::recursive_mutex> g(m_mutex);

      if(m_promisesWithoutFutures > 0) {
        --m_promisesWithoutFutures;

        auto promise = std::move(m_promises.front());
        m_promises.pop();

        return promise.getFuture();
      } else {

        folly::Promise<ValueType> promise;
        auto result = promise.getFuture();
        m_promises.emplace(std::move(promise));

        for(size_t i = 0; i < m_generators.size(); ++i) {
          if(!m_isPaused[i] && !m_futuresPending[i] && !m_isDisabled[i]) {
            initFuture(i);
          }
        }

        return result;
      }
    }
  }

  template <typename Generator>
  bool GeneratorValueStream<Generator>::isEmpty() const
  {
    std::lock_guard<std::recursive_mutex> g(m_mutex);

    return m_generators.empty();
  }

  template <typename Generator>
  bool GeneratorValueStream<Generator>::isPauseValue(const ValueType& v)
  {
    return m_syncTest && m_syncTest(v);
  }

  template <typename Generator>
  void GeneratorValueStream<Generator>::
  setSynchronizationTest(std::function<bool(const ValueType&)> f)
  {
    m_syncTest = f;
  }

  template <typename Generator>
  void GeneratorValueStream<Generator>::
  setCallback(std::function<void(const ValueType&)> f)
  {
    m_onValue = f;
  }

  template <typename Generator>
  void GeneratorValueStream<Generator>::
  addStream(std::shared_ptr<GeneratorValueStream<Generator>> other,
            const QByteArray & key)
  {
    std::lock_guard<std::recursive_mutex> g(m_mutex);

    assert(!other->m_init);
    size_t offset = m_generators.size();

    // The stream can have a key associated with it, which can be used later to remove
    // generators that were added from that stream
    bool storeIds = !key.isEmpty();
    std::vector<size_t> generators;

    for(size_t i = 0; i < other->m_generators.size(); ++i) {
      m_generators.emplace_back(other->m_generators[i]);
      m_futures.emplace_back(folly::makeFuture<ValueType>(
                folly::exception_wrapper(std::runtime_error("empty"))));
      m_isPaused.push_back(false);
      m_futuresPending.push_back(false);
      m_isDisabled.push_back(false);
      initFuture(offset + i);

      if (storeIds)
        generators.push_back(offset + i);
    }

    if (storeIds)
      m_addedGenerators[key] = generators;
  }

  template <typename Generator>
  void GeneratorValueStream<Generator>::
  removeAddedGenerators(const QByteArray & key)
  {
    std::lock_guard<std::recursive_mutex> g(m_mutex);

    auto mapIt = m_addedGenerators.find(key);
    if (mapIt == m_addedGenerators.end())
      return;

    // Since everything in this class relies on generator positions in vector, don't
    // actually remove the generators, just disable them
    // @todo could do a clean up and remap when we're not waiting for any values
    for(auto it = mapIt->second.begin(); it != mapIt->second.end(); it++) {
      auto i = (*it);
      assert(!m_isDisabled[i]);
      m_isDisabled[i] = true;
      m_disabledGenerators++;

      if (m_isPaused[i])
        m_generatorsPending--;
      if (m_futuresPending[i])
        m_futuresWaited--;
    }

    m_addedGenerators.erase(mapIt);
  }


}