generator.hpp

#pragma once
 
#include "array.hpp"
#include <functional>
#include <map>
 
namespace z {
namespace core {
 
template <typename T>
struct yield {
    bool done;
 
    T value;
 
    operator bool() const {
        return !done;
    }
};
 
template <typename T, typename S>
class generatorIter {
    std::function<const yield<T>(S &)> lambda;
    S state;
    yield<T> current_yield;
 
public:
    explicit generatorIter(std::function<const yield<T>(S &)> lambda, const S &state, bool dummy = false) : lambda(lambda), state(state), current_yield({false, {}}) {
        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.done;
    }
};
 
template <typename T, typename S>
class generator : public iterable<generatorIter<T, S>> {
    S state;
    std::function<const yield<T>(S &)> lambda;
 
    struct countedState {
        long count;
        S state;
    };
 
public:
    generator(const S &initial, std::function<const yield<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 yield<T> next() {
        return lambda(state);
    }
 
    inline long count() {
        return this
            ->map<long>([](const T &state) {
                (void)state;
                return 1;
            })
            .reduce(0, std::plus<long>());
    }
 
    array<T> collect() {
        array<T> result;
        for (auto i : *this) {
            result.push(i);
        }
        return result;
    }
 
    long consume() {
        long count = 0;
        for (auto _ : *this) {
            count++;
        }
        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.done) {
                break;
            }
            result.push(item.value);
        }
 
        return result;
    }
 
    template <typename U>
    generator<U, S> map(std::function<U(const T &)> mapLambda) {
        auto lambda = this->lambda;
 
        return generator<U, S>(state, [lambda, mapLambda](S &state) {
            auto item = lambda(state);
            if (item.done) {
                return yield<U>{true};
            } else {
                return yield<U>{false, mapLambda(item.value)};
            }
        });
    }
 
    generator filter(std::function<T(const T &)> filterLambda) {
        auto lambda = this->lambda;
 
        return generator(state, [lambda, filterLambda](S &state) {
            auto val = lambda(state);
            while (!val.done) {
                if (filterLambda(val.value)) {
                    return val;
                }
                val = lambda(state);
            }
 
            return val;
        });
    }
 
    T reduce(const T &defaultValue, std::function<T(const T &, const T &)> reduceLambda) {
        auto result = lambda(state);
 
        if (result.done) {
            return defaultValue;
        }
 
        auto value = result.value;
 
        while (true) {
            result = lambda(state);
            if (result.done) {
                break;
            }
            value = reduceLambda(value, result.value);
        }
 
        return value;
    }
 
    generator &forEach(std::function<void(const T &)> newLambda) {
        auto lambda = this->lambda;
        this->lambda = [lambda, newLambda](S &state) {
            auto item = lambda(state);
            if (!item.done) {
                newLambda(item.value);
            }
            return item;
        };
 
        return *this;
    }
 
    generator<T, countedState> skip(long count) {
        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.done) {
                    return item;
                }
            }
            state.count = 0;
 
            return lambda(state.state);
        });
    }
 
    generator<T, countedState> limit(long count) {
        auto lambda = this->lambda;
 
        return generator<T, countedState>({count, state}, [lambda](countedState &state) {
            if (state.count <= 0) {
                return yield<T>{true};
            }
 
            auto item = lambda(state.state);
            if (item.done) {
                return item;
            }
 
            state.count--;
 
            return item;
        });
    }
 
    template <typename U, typename S2>
    generator<std::pair<T, U>, generator<U, S2>> zip(generator<U, S2> &other) {
        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.done) {
                return yield<pair_type>{true};
            }
 
            auto item2 = otherGen.next();
            if (item2.done) {
                return yield<pair_type>{true};
            }
 
            return yield<pair_type>{false, {item1.value, item2.value}};
        });
    }
 
    generator<std::pair<long, T>, std::pair<long, generator<T, S>>> enumerate() {
        return generator<std::pair<long, T>, std::pair<long, generator<T, S>>>({0, *this}, [](std::pair<long, generator<T, S>> &state) {
            auto item = state.second.next();
            if (item.done) {
                return yield<std::pair<long, T>>{true};
            }
            return yield<std::pair<long, T>>{false, {state.first++, item.value}};
        });
    }
 
    generator<T, std::pair<generator, yield<T>>> diff(generator &other) {
        return generator<T, std::pair<generator, yield<T>>>({other, other.next()}, [this](std::pair<generator, yield<T>> &state) {
            while (true) {
                auto item1 = next();
                if (item1.done) {
                    return yield<T>{true};
                }
 
                if (state.second.done) {
                    return yield<T>{false, item1.value}; // Yield the item from this generator, as the other generator is done
                }
 
                if (item1.value != state.second.value) {
                    return yield<T>{false, 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) {
        return generator<array<T>, generator>(*this, [chunkSize](generator &state) {
            array<T> chunk;
            for (long i = 0; i < chunkSize; i++) {
                auto item = state.next();
                if (item.done) {
                    if (chunk.length() == 0) {
                        return yield<array<T>>{true}; // No more items, end the generator
                    }
                    break; // The chunk has data, yield it
                }
                chunk.push(item.value);
            }
            return yield<array<T>>{false, chunk}; // Return the current chunk
        });
    }
 
    generator<std::pair<T, yield<T>>, std::pair<yield<T>, generator>> peek() {
        return generator<std::pair<T, yield<T>>, std::pair<yield<T>, generator>>({next(), *this}, [](std::pair<yield<T>, generator> &state) {
            const auto prevValue = state.first;
            auto &gen = state.second;
 
            if (prevValue.done) {
                // No more items, end the generator
                return yield<std::pair<T, yield<T>>>{true};
            }
            auto nextValue = gen.next();
            state.first = nextValue;
 
            return yield<std::pair<T, yield<T>>>{false, std::pair<T, yield<T>>(prevValue.value, nextValue)};
        });
    }
};
 
template <typename T>
using deref_type = std::remove_const_t<std::remove_reference_t<decltype(*std::declval<decltype(std::declval<T>().begin())>())>>;
 
template <typename T>
using iter_type = std::remove_const_t<decltype(std::declval<T>().begin())>;
 
template <typename T>
generator<deref_type<T>, iter_type<T>> generatorFrom(const T &list) {
    return generator<deref_type<T>, iter_type<T>>(list.begin(), [&list](iter_type<T> &iter) {
        if (iter != list.end()) {
            auto ret = yield<deref_type<T>>{false, *iter};
            ++iter; // Move to the next item
            return ret;
        }
 
        return yield<deref_type<T>>{true};
    });
}
 
template <typename T>
generator<deref_type<T>, std::pair<T, iter_type<T>>> generatorFrom(const T &&list) {
    return generator<deref_type<T>, std::pair<T, iter_type<T>>>({list, list.begin()}, [](std::pair<T, iter_type<T>> &state) {
        if (state.second != state.first.end()) {
            auto ret = yield<deref_type<T>>{false, *state.second};
            ++state.second; // Move to the next item
            return ret;
        }
 
        return yield<deref_type<T>>{true};
    });
}
 
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}, [&list](std::pair<array<T>, long> &state) {
        if (state.second < state.first.length()) {
            return yield<T>{false, state.first[state.second++]};
        }
 
        return yield<T>{true};
    });
}
 
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) {
        if (iter != map.end()) {
            auto ret = yield<std::pair<K, V>>{false, *iter};
            ++iter; // Move to the next item
            return ret;
        }
 
        return yield<std::pair<K, V>>{true};
    });
}
 
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) {
        if (state.first != state.second.end()) {
            auto ret = yield<std::pair<K, V>>{false, *state.first};
            ++state.first; // Move to the next item
            return ret;
        }
 
        return yield<std::pair<K, V>>{true};
    });
}
 
} // namespace core
} // namespace z