generator.cpp

#pragma once

template<typename T>
struct Maybe {
  Maybe()
    : have_value_(false) {}
  Maybe(T value)
    : have_value_(true)
    , value_(std::move(value)) {}
  Maybe(const Maybe&) = delete;
  Maybe(Maybe&& other) {
    *this = std::move(other);
  }

  Maybe& operator=(Maybe&& other) {
    if (have_value_ && other.have_value_) {
      value_ = std::move(other.value_);
    } else if (have_value_ && !other.have_value_) {
      value_.~T();
      have_value_ = false;
    } else if (!have_value_ && other.have_value_) {
      new (&value_) T(std::move(other.value_));
      have_value_ = true;
    } else if (!have_value_ && !other.have_value_) {
      /* nothing to do */
    }
    return *this;
  }

  ~Maybe() {
    if (!have_value_) return;
    value_.~T();
  }

  T& value() {
    TORCH_ASSERT(have_value_);
    return value_;
  }

  bool empty() const {
    return !have_value_;
  }

private:
  bool have_value_;
  union {
    T value_;
  };
};


template<typename T>
struct generator {
  using next_type = std::function<Maybe<T> ()>;

  generator(next_type next)
    : next_(std::move(next)) {
    next_value_ = next_();
  }

  struct generator_iter {
    generator_iter(generator& gen)
      : gen_(gen) {}

    bool operator!=(const generator_iter& other) const {
      return !gen_.next_value_.empty();
    }

    generator_iter& operator++() {
      gen_.next_value_ = gen_.next_();
      return *this;
    }

    T&& operator*() {
      return std::move(gen_.next_value_.value());
    }

    generator& gen_;
  };

  generator_iter begin() {
    TORCH_ASSERT(!started_);
    started_ = true;
    return generator_iter(*this);
  }

  generator_iter end() {
    return generator_iter(*this);
  }

  Maybe<T> next_value_;
  next_type next_;
  bool started_ = false;
};