generator.hpp

#pragma once
 
#include <functional>
#include <map>
#include <optional>
 
#include "array.hpp"
#include "templates.hpp"
 
namespace z {
namespace core {

template <typename T, typename S>
class generatorIter {
    std::function<const std::optional<T>(S &)> lambda;
    S state;
    std::optional<T> current_yield;
 
public:
    explicit generatorIter(std::function<const std::optional<T>(S &)> lambda, const S &state, bool dummy = false) : lambda(lambda), state(state), current_yield(T()) {
        if (!dummy) {
            ++(*this); // Load the first value
        }
    }

    const T &operator*() const {
        return current_yield.value();
    }

    generatorIter &operator++() {
        current_yield = lambda(state);
        return *this;
    }

    bool operator!=(const generatorIter &other) const {
        (void)other;
        return current_yield.has_value();
    }
};

template <typename T, typename S>
class generator : public iterable<generatorIter<T, S>> {
protected:
    S state;

    std::function<const std::optional<T>(S &)> lambda;

    struct countedState {
        long count;
        S state;
    };
 
public:
    generator(const S &initial, std::function<const std::optional<T>(S &)> lambda) : state(initial), lambda(lambda) {}

    generatorIter<T, S> begin() const noexcept override {
        return generatorIter<T, S>(lambda, state);
    }

    generatorIter<T, S> end() const noexcept override {
        return generatorIter<T, S>(lambda, state, true);
    }

    inline virtual std::optional<T> next() {
        return lambda(state);
    }

    long count() {
        long count = 0;
        for (auto _ : *this) {
            count++;
        }
        return count;
    }

    array<T> collect() {
        array<T> result;
        for (auto i : *this) {
            result.push(i);
        }
        return result;
    }

    inline long consume() {
        return count();
    }

    array<T> take(int count) {
        array<T> result;
        result.increase(count);
 
        for (int i = 0; i < count; i++) {
            auto item = next();
            if (!item.has_value()) {
                break;
            }
            result.push(item.value());
        }
 
        return result;
    }

    template <typename U>
    generator<U, S> map(std::function<U(const T &)> mapLambda) noexcept {
        auto lambda = this->lambda;
 
        return generator<U, S>(state, [lambda, mapLambda](S &state) -> std::optional<U> {
            auto item = lambda(state);
            if (!item.has_value()) {
                return {};
            } else {
                return mapLambda(item.value());
            }
        });
    }

    generator filter(std::function<bool(const T &)> filterLambda) noexcept {
        auto lambda = this->lambda;
 
        return generator(state, [lambda, filterLambda](S &state) {
            auto val = lambda(state);
            while (val.has_value()) {
                if (filterLambda(val.value())) {
                    return val;
                }
                val = lambda(state);
            }
 
            return val;
        });
    }

    T reduce(std::function<T(const T &, const T &)> reduceLambda, const T &defaultValue = {}) {
        auto result = lambda(state);
 
        if (!result.has_value()) {
            return defaultValue;
        }
 
        auto value = result.value();
 
        while (true) {
            result = lambda(state);
            if (!result.has_value()) {
                break;
            }
            value = reduceLambda(value, result.value());
        }
 
        return value;
    }

    generator &forEach(std::function<void(const T &)> newLambda) noexcept {
        auto lambda = this->lambda;
        this->lambda = [lambda, newLambda](S &state) {
            auto item = lambda(state);
            if (item.has_value()) {
                newLambda(item.value());
            }
            return item;
        };
 
        return *this;
    }

    generator<T, countedState> skip(long count) noexcept {
        auto lambda = this->lambda;
 
        return generator<T, countedState>({count, state}, [lambda](countedState &state) {
            for (long i = 0; i < state.count; i++) {
                auto item = lambda(state.state);
                if (!item.has_value()) {
                    return item;
                }
            }
            state.count = 0;
 
            return lambda(state.state);
        });
    }

    generator<T, countedState> limit(long count) noexcept {
        auto lambda = this->lambda;
 
        return generator<T, countedState>({count, state}, [lambda](countedState &state) {
            if (state.count <= 0) {
                return std::optional<T>();
            }
 
            auto item = lambda(state.state);
            if (!item.has_value()) {
                return item;
            }
 
            state.count--;
 
            return item;
        });
    }

    template <typename U, typename S2>
    generator<std::pair<T, U>, generator<U, S2>> pair(generator<U, S2> &other) noexcept {
        typedef std::pair<T, U> pair_type;
 
        return generator<pair_type, generator<U, S2>>(other, [this](generator<U, S2> &otherGen) {
            auto item1 = next();
            if (!item1.has_value()) {
                return std::optional<pair_type>();
            }
 
            auto item2 = otherGen.next();
            if (!item2.has_value()) {
                return std::optional<pair_type>();
            }
 
            return std::optional<pair_type>({item1.value(), item2.value()});
        });
    }

    generator<T, std::pair<bool, generator &>> zip(generator &other) noexcept {
        return generator<T, std::pair<bool, generator &>>({false, other}, [this](std::pair<bool, generator &> &state) {
            state.first = !state.first;
 
            // Draw from first generator
            if (state.first) {
                auto item = next();
                if (item.has_value()) {
                    return item;
                }
                return state.second.next();
            }
 
            // Draw from second generator
            auto item = state.second.next();
            if (item.has_value()) {
                return item;
            }
            return next();
        });
    }

    generator<std::pair<long, T>, std::pair<long, generator<T, S>>> enumerate() noexcept {
        return generator<std::pair<long, T>, std::pair<long, generator<T, S>>>({0, *this}, [](std::pair<long, generator<T, S>> &state) -> std::optional<std::pair<long, T>> {
            auto item = state.second.next();
            if (!item.has_value()) {
                return {};
            }
            return std::pair<long, T>{state.first++, item.value()};
        });
    }

    generator<T, std::pair<generator, std::optional<T>>> diff(generator &other) noexcept {
        return generator<T, std::pair<generator, std::optional<T>>>({other, other.next()}, [this](std::pair<generator, std::optional<T>> &state) -> std::optional<T> {
            while (true) {
                auto item1 = next();
                if (!item1.has_value()) {
                    return {};
                }
 
                if (!state.second.has_value()) {
                    return item1.value(); // Yield the item from this generator, as the other generator is done
                }
 
                if (item1.value() != state.second.value()) {
                    return item1.value(); // Yield the item from this generator, as it is different
                }
 
                // Move to the next item in the other generator
                state.second = state.first.next();
            }
        });
    }

    generator<array<T>, generator> chunk(long chunkSize) noexcept {
        return generator<array<T>, generator>(*this, [chunkSize](generator &state) -> std::optional<array<T>> {
            array<T> chunk;
            for (long i = 0; i < chunkSize; i++) {
                auto item = state.next();
                if (!item.has_value()) {
                    if (chunk.length() == 0) {
                        return {}; // No more items, end the generator
                    }
                    break; // The chunk has data, yield it
                }
                chunk.push(item.value());
            }
            return chunk; // Return the current chunk
        });
    }

    template <typename U = T, typename std::enable_if<is_iterator<U>::value, int>::type = 0>
    auto flatten() noexcept {
        struct unchunkData {
            std::optional<T> data;
            std::optional<iterator_value<T>> iter;
            generator<T, S> gen;
        };
 
        auto unchunkFn = [](unchunkData &state) -> std::optional<dereference<T>> {
            // If the generator has not been initialized,
            // or there is no more data in the chunk,
            // then try to generate another chunk.
            if (!state.data || (state.iter.value() == state.data.value().end())) {
                // If there's nothing left, we're done.
                auto chunkData = state.gen.next();
                if (!chunkData) {
                    return {};
                }
 
                auto &val = chunkData.value();
                // Make sure we didn't just get an empty iterator.
                if (val.begin() == val.end()) {
                    return {};
                }
 
                state.data = chunkData;
                state.iter = state.data.value().begin();
            }
 
            // At this point the generator has been initialized, and current chunk has data in it,
            // so spit out the next value on the chunk.
 
            return *state.iter.value()++;
        };
 
        return generator<dereference<T>, unchunkData>({{}, {}, *this}, unchunkFn);
    }

    generator<std::pair<T, std::optional<T>>, std::pair<std::optional<T>, generator>> peek() noexcept {
        return generator<std::pair<T, std::optional<T>>, std::pair<std::optional<T>, generator>>({next(), *this}, [](std::pair<std::optional<T>, generator> &state) -> std::optional<std::pair<T, std::optional<T>>> {
            const auto prevValue = state.first;
            auto &gen = state.second;
 
            if (!prevValue.has_value()) {
                // No more items, end the generator
                return {};
            }
            auto nextValue = gen.next();
            state.first = nextValue;
 
            return std::pair<T, std::optional<T>>(prevValue.value(), nextValue);
        });
    }

    template <typename U>
    generator<T, std::pair<generator, bool>> chain(generator<T, U> &other) noexcept {
        return generator<T, std::pair<generator, bool>>({*this, false}, [&other](std::pair<generator, bool> &state) {
            if (!state.second) {
                auto item = state.first.next();
                if (!item.has_value()) {
                    state.second = true;
                } else {
                    return item;
                }
            }
 
            return other.next();
        });
    }

    template <typename U>
    generator<T, std::pair<std::pair<generator, bool>, generator<T, U>>> chain(generator<T, U> &&other) noexcept {
        return generator<T, std::pair<std::pair<generator, bool>, generator<T, U>>>({{*this, false}, other}, [&other](std::pair<std::pair<generator, bool>, generator<T, U>> &state) {
            if (!state.first.second) {
                auto item = state.first.first.next();
                if (!item.has_value()) {
                    state.first.second = true;
                } else {
                    return item;
                }
            }
 
            return state.second.next();
        });
    }

    generator until(std::function<bool(T)> predicate) noexcept {
        auto lambda = this->lambda;
 
        return generator(state, [lambda, predicate](S &state) {
            auto val = lambda(state);
            if (val.has_value() && predicate(val.value())) {
                return std::optional<T>{};
            }
            return val;
        });
    }

    generator until(const T &sentinel) noexcept {
        auto lambda = this->lambda;
 
        return generator(state, [lambda, sentinel](S &state) {
            auto val = lambda(state);
            if (val.has_value() && val.value() == sentinel) {
                return std::optional<T>{};
            }
            return val;
        });
    }

    template <typename U>
    inline generator<T, std::pair<generator, bool>> operator+(generator<T, U> &other) noexcept {
        return chain(other);
    }

    template <typename U>
    inline generator<T, std::pair<std::pair<generator, bool>, generator<T, U>>> operator+(generator<T, U> &&other) noexcept {
        return generator<T, std::pair<std::pair<generator, bool>, generator<T, U>>>({{*this, false}, other}, [&other](std::pair<std::pair<generator, bool>, generator<T, U>> &state) {
            if (!state.first.second) {
                auto item = state.first.first.next();
                if (!item.has_value()) {
                    state.first.second = true;
                } else {
                    return item;
                }
            }
 
            return state.second.next();
        });
    }

    template <typename U>
    inline generator<U, S> operator|(std::function<U(T)> mapLambda) noexcept {
        return map<U>(mapLambda);
    }

    template <typename U>
    inline generator<U, S> operator|(U (*mapLambda)(T)) noexcept {
        return map<U>(mapLambda);
    }

    inline generator operator&&(std::function<T(const T &)> filterLambda) noexcept {
        return filter(filterLambda);
    }

    inline T operator>>(std::function<T(const T &, const T &)> reduceLambda) {
        return reduce(reduceLambda);
    }

    inline generator operator!=(std::function<bool(T)> predicate) {
        return until(predicate);
    }

    inline generator operator!=(const T &sentinel) {
        return until(sentinel);
    }
};

template <typename T>
generator<dereference<T>, const_iterator_value<T>> generatorFrom(const T &list) {
    return generator<dereference<T>, const_iterator_value<T>>(list.begin(), [&list](const_iterator_value<T> &iter) -> std::optional<dereference<T>> {
        if (iter != list.end()) {
            auto ret = *iter;
            ++iter; // Move to the next item
            return ret;
        }
 
        return {};
    });
}

template <typename T>
generator<dereference<T>, std::pair<T, iterator_value<T>>> generatorFrom(const T &&list) {
    return generator<dereference<T>, std::pair<T, iterator_value<T>>>({list, list.begin()}, [](std::pair<T, iterator_value<T>> &state) -> std::optional<dereference<T>> {
        if (state.second != state.first.end()) {
            auto ret = *state.second;
            ++state.second; // Move to the next item
            return ret;
        }
 
        return {};
    });
}

template <typename T>
generator<T, std::pair<array<T>, long>> generatorFrom(std::initializer_list<T> list) {
    return generator<T, std::pair<array<T>, long>>({list, 0}, [](std::pair<array<T>, long> &state) -> std::optional<T> {
        if (state.second < state.first.length()) {
            return state.first[state.second++];
        }
 
        return {};
    });
}

template <typename K, typename V>
generator<std::pair<K, V>, typename std::map<K, V>::const_iterator> generatorFrom(const std::map<K, V> &map) {
    return generator<std::pair<K, V>, typename std::map<K, V>::const_iterator>(map.begin(), [&map](auto &iter) -> std::optional<std::pair<K, V>> {
        if (iter != map.end()) {
            auto ret = *iter;
            ++iter; // Move to the next item
            return ret;
        }
 
        return {};
    });
}

template <typename K, typename V>
generator<std::pair<K, V>, std::pair<typename std::map<K, V>::const_iterator, std::map<K, V>>> generatorFrom(std::map<K, V> &&map) {
    return generator<std::pair<K, V>, std::pair<typename std::map<K, V>::const_iterator, std::map<K, V>>>({map.begin(), map}, [](auto &state) -> std::optional<std::pair<K, V>> {
        if (state.first != state.second.end()) {
            auto ret = *state.first;
            ++state.first; // Move to the next item
            return ret;
        }
 
        return {};
    });
}
 
} // namespace core
} // namespace z