rusty/rs_iter.hpp

#ifndef _RC_ITER_HPP_
#define _RC_ITER_HPP_

#include <cstdint>
#include <functional>
#include <optional>
#include <type_traits>

namespace rs {

template<typename... ARGS>
struct InvokeType {
    template<typename F>
    auto Call(F f, ARGS... args) { return f(std::forward<ARGS>(args)...); }
    template<typename F>
    void CallNoReturn(F f, ARGS... args) { f(std::forward<ARGS>(args)...); }
    template<typename F, typename S>
    auto CallFold(F f, S s, ARGS... args) { return f(std::forward<S>(s), std::forward<ARGS>(args)...); }
};

template<typename U, typename V>
struct InvokeType<std::pair<U, V>> {
    template<typename F>
    auto Call(F f, std::pair<U, V> arg) { return f(arg.first, arg.second); }
    template<typename F>
    void CallNoReturn(F f, std::pair<U, V> arg) { f(arg.first, arg.second); }
    template<typename F, typename S>
    auto CallFold(F f, S s, std::pair<U, V> arg) { return f(std::forward<S>(s), arg.first, arg.second); }
};

template<typename... ARGS>
struct InvokeType<std::tuple<ARGS...>> {
    using sequence_type = std::index_sequence_for<ARGS...>;
    template<typename F>
    auto Call(F f, std::tuple<ARGS...> arg) { return UnpackCall(std::forward<F>(f), arg, sequence_type{}); }
    template<typename F>
    void CallNoReturn(F f, std::tuple<ARGS...> arg) { UnpackCallNoReturn(std::forward<F>(f), arg, sequence_type{}); }
    template<typename F, typename S>
    auto CallFold(F f, S s, std::tuple<ARGS...> arg) { return UnpackCallFold(std::forward<F>(f), std::forward<S>(s), arg, sequence_type{}); }

    template<typename F, size_t... I>
    auto UnpackCall(F f, std::tuple<ARGS...>& arg, std::index_sequence<I...>) { return f(std::get<I>(arg)...); }
    template<typename F, size_t... I>
    void UnpackCallNoReturn(F f, std::tuple<ARGS...>& arg, std::index_sequence<I...>) { f(std::get<I>(arg)...); }
    template<typename F, typename S, size_t... I>
    auto UnpackCallFold(F f, S s, std::tuple<ARGS...>& arg, std::index_sequence<I...>) { return f(s, std::get<I>(arg)...); }
};

template <typename T>
struct StdFuncType {
    using type = void;
    using return_type = void;
};
template <typename RET, typename CLASS, typename... ARGS>
struct StdFuncType<RET (CLASS::*)(ARGS...) const> {
    using type = std::function<RET(ARGS...)>;
    using return_type = RET;
};
template <typename RET, typename CLASS, typename... ARGS>
struct StdFuncType<RET (CLASS::*)(ARGS...)> {
    using type = std::function<RET(ARGS...)>;
    using return_type = RET;
};

template <typename HANDLER>
struct FunctionType {
    using type = typename StdFuncType<decltype(&HANDLER::operator())>::type;
    using return_type = typename StdFuncType<decltype(&HANDLER::operator())>::return_type;
};
template <typename RET, typename... ARGS>
struct FunctionType<RET (*)(ARGS...)> {
    using type = std::function<RET(ARGS...)>;
    using return_type = RET;
};
template <typename RET, typename... ARGS>
struct FunctionType<RET(ARGS...)> {
    using type = std::function<RET(ARGS...)>;
    using return_type = RET;
};

template <typename T, template<typename... TARGS> typename TP>
struct TemplateChecker {
    static constexpr bool value = false;
};
template <template<typename... TARGS> typename TP, typename... ARGS>
struct TemplateChecker<TP<ARGS...>, TP> {
    static constexpr bool value = true;
};

template <typename T>
struct iterator {
    using value_type = T;

    std::optional<T> next() {
        return iter();
    }

    size_t count() {
        size_t result = 0;
        for (auto n = next(); n.has_value(); n = next())
            result++;
        return result;
    }

    std::optional<T> nth(size_t n) {
        while (true) {
            std::optional<T> value = next();
            if (n-- == 0 || !value.has_value())
                return value;
        }
    }

    std::optional<T> last() {
        std::optional<T> value;
        for (auto n = next(); n.has_value(); n = next())
            value = n;
        return value;
    }

    auto step_by(int32_t s) {
        return iterator<T>{[s, c = s, it = std::move(iter)]() mutable -> auto {
            for (; c < s; c++)
                it();
            c = 1;
            return it();
        }};
    }

    auto chain(iterator<T>&& i) {
        return iterator<T>{[it1 = std::move(iter), it2 = std::move(i.iter), s = true]() mutable -> auto {
            if (s) {
                auto v = it1(true);
                if (!v.has_value()) {
                    s = false;
                    return it2(true);
                }
            } else
                return it2(true);
        }};
    }

    template<typename U>
    auto zip(iterator<U>&& i) {
        return iterator<std::pair<T, U>>{[it1 = std::move(iter), it2 = std::move(i.iter)]() mutable -> auto {
            if(auto v1 = it1(); v1.has_value()) {
                if(auto v2 = it2(); v2.has_value())
                    return std::optional<std::pair<T, U>>(std::make_pair(v1.value(), v2.value()));
            }
            return std::optional<std::pair<T, U>>();
        }};
    }

    template <typename F>
    void for_each(F f) {
        for(auto v = next(); v.has_value(); v = next()) {
            InvokeType<T>().CallNoReturn(std::ref(f), v.value());
        }
    }

    template <typename F>
    auto filter(F f) {
        return iterator<T>{[f, it = std::move(iter)]() mutable -> auto {
            auto value = it();
            while(value.has_value() && !InvokeType<T>().Call(std::ref(f), value.value()))
                value = it();
            return value;
        }};
    }

    template <typename M>
    auto map(M m) {
        using RET = typename FunctionType<M>::return_type;
        return iterator<RET>{[m, it = std::move(iter)]() -> auto {
            auto t = it();
            if (t.has_value())
                return std::optional<RET>(InvokeType<T>().Call(std::ref(m), t.value()));
            return std::optional<RET>();
        }};
    }

    template <typename FM>
    auto filter_map(FM fm) {
        using OPT = typename FunctionType<FM>::return_type;
        static_assert(TemplateChecker<OPT, std::optional>::value, "filter_map: filter must return an optional value.");
        using RET = typename OPT::value_type;
        return iterator<RET>{[fm, it = std::move(iter)]() -> auto {
            auto value = it();
            while (value.has_value()) {
                if(auto mvalue = InvokeType<T>().Call(std::ref(fm), value.value()); mvalue.has_value())
                    return mvalue;
                value = it();
            }                
            return std::optional<RET>();
        }};
    }

    auto enumerate() {
        return iterator<std::pair<size_t, T>>{[it = std::move(iter), i = 0]() mutable -> auto {
            if(auto v = it(); v.has_value())
                return std::optional<std::pair<size_t, T>>(std::make_pair(i++, v.value()));
            return std::optional<std::pair<size_t, T>>();
        }};
    }

    auto skip(size_t n) {
        return iterator<T>{[n, it = std::move(iter)]() mutable -> auto {
            while (n > 0) {
                n--;
                it();
            }
            return it();
        }};
    }

    auto take(size_t n) {
        return iterator<T>{[n, it = std::move(iter)]() mutable -> auto {
            if (n > 0) {
                n--;
                return it();
            }
            return std::optional<T>();
        }};
    }

    template<typename F>
    auto skip_while(F f) {
        return iterator<T>{[f, it = std::move(iter), res = false]() mutable -> auto {
            auto v = it();
            if(res)
                return v;
            while(v.has_value() && InvokeType<T>().Call(std::ref(f), v.value()))
                v = it();
            res = true;
            return v;
        }};
    }

    template <typename F>
    auto take_while(F f) {
        return iterator<T>{[f, it = std::move(iter), res = true]() mutable -> auto {
            if (res) {
                if(auto v = it(); v.has_value() && InvokeType<T>().Call(std::ref(f), v.value()))
                    return v;
                else
                    res = false;
            }
            return std::optional<T>();
        }};
    }

    template <typename S, typename M>
    auto scan(S s, M m) {
        using RET = typename FunctionType<M>::return_type;
        return iterator<RET>{[s, m, it = std::move(iter)]() mutable -> auto {
            auto t = it();
            if(t.has_value())
                return std::optional<RET>(InvokeType<T>().CallFold(std::ref(m), std::ref(s), t.value()));
            return std::optional<RET>();
        }};
    }

    template <typename M>
    auto flat_map(M m) {
        return map(m).flatten();
    }

    auto flatten() {
        static_assert(TemplateChecker<T, rs::iterator>::value, "flatten: value type can only be iterator type.");
        using NT = typename T::value_type;
        return iterator<NT>{[inited = false, cur = std::optional<iterator<NT>>(), it = std::move(iter)]() mutable -> auto {
            if(!inited) {
                cur = it();
                inited = true;
            }
            if(!cur.has_value())
                return std::optional<NT>();
            auto v = cur.value().next();
            while(!v.has_value() && cur.has_value()) {
                cur = it();
                if(cur.has_value())
                    v = cur.value().next();
            }
            return v;
        }};
    }

    auto fuse() {
        return iterator<T>{[end = false, it = std::move(iter)]() mutable -> auto {
            if (end)
                return std::optional<T>();
            auto v = it();
            if(!it.has_value())
                end = true;
            return v;
        }};
    }

    template <typename F>
    auto inspect(F f) {
        return iterator<T>{[f, it = std::move(iter)]() -> auto {
            auto v = it();
            if(v.has_value())
                InvokeType<T>().CallNoReturn(std::ref(f), v.value());
            return v;
        }};
    }

    template<typename C>
    auto collect() {
        C container;
        for (auto v = next(); v.has_value(); v = next())
            container.collect_item(v.value());
        return container;
    }

    template<typename C>
    C& collect(C& container) {
        for (auto v = next(); v.has_value(); v = next())
            container.collect_item(v.value());
        return container;
    }

    template<typename C, typename F>
    auto partition(F f) {
        std::pair<C, C> result;
        for (auto v = next(); v.has_value(); v = next()) {
            if(InvokeType<T>().CallNoReturn(std::ref(f), v.value()))
                result.first.collect_item(v.value());
            else
                result.second.collect_item(v.value());
        }
        return result;
    }

    template <typename S, typename F>
    auto fold(S s, F f) {
        using RET = typename FunctionType<F>::return_type;
        static_assert(std::is_same<S, RET>::value, "fold: function should return same type as value.");
        for (auto v = next(); v.has_value(); v = next()) {
            s = InvokeType<T>().CallFold(std::ref(f), s, v.value());
        }
        return s;
    }

    template <typename F>
    bool all(F f) {
        for (auto v = next(); v.has_value(); v = next()) {
            if(!InvokeType<T>().Call(std::ref(f), v.value()))
                return false;
        }
        return true;
    }

    template <typename F>
    bool any(F f) {
        for (auto v = next(); v.has_value(); v = next()) {
            if(InvokeType<T>().Call(std::ref(f), v.value()))
                return true;
        }
        return false;
    }

    template <typename F>
    auto find(F f) {
        for (auto v = next(); v.has_value(); v = next()) {
            if(InvokeType<T>().Call(std::ref(f), v.value()))
                return v;
        }
        return std::optional<T>();
    }

    template <typename F>
    auto find_map(F f) {
        using OPT = typename FunctionType<F>::return_type;
        static_assert(TemplateChecker<OPT, std::optional>::value, "find_map: filter must return an optional value.");
        using RET = typename OPT::value_type;
        for (auto v = next(); v.has_value(); v = next()) {
            if(auto mv = InvokeType<T>().Call(std::ref(f), v.value()); mv.has_value())
                return mv;
        }
        return std::optional<RET>();
    }

    template <typename F>
    auto position(F f) {
        size_t i = 0;
        for (auto v = next(); v.has_value(); v = next(), i++) {
            if(InvokeType<T>().Call(std::ref(f), v.value()))
                return std::optional<size_t>(i);
        }
        return std::optional<size_t>();
    }

    auto max() {
        auto res = next();
        if(!res.has_value())
            return res;
        for (auto v = next(); v.has_value(); v = next()) {
            if(v.value() >= res.value())
                res = v;
        }
        return res;
    }

    auto min() {
        auto res = next();
        if(!res.has_value())
            return res;
        for (auto v = next(); v.has_value(); v = next()) {
            if(v.value() < res.value())
                res = v;
        }
        return res;
    }

    template <typename F>
    auto max_by_key(F f) {
        auto res = next();
        if(!res.has_value())
            return res;
        for (auto v = next(); v.has_value(); v = next()) {
            if(InvokeType<T>().Call(std::ref(f), v.value()) >= InvokeType<T>().Call(std::ref(f), res.value()))
                res = v;
        }
        return res;
    }

    template <typename F>
    auto min_by_key(F f) {
        auto res = next();
        if(!res.has_value())
            return res;
        for (auto v = next(); v.has_value(); v = next()) {
            if(InvokeType<T>().Call(std::ref(f), v.value()) < InvokeType<T>().Call(std::ref(f), res.value()))
                res = v;
        }
        return res;
    }

    auto unzip() {
        static_assert(TemplateChecker<T, std::pair>::value, "unzip: iterator type must be pair");
        using TP1 = typename T::first_type;
        using TP2 = typename T::second_type;
        std::pair<iterator<TP1>, iterator<TP2>> result;
        result.first.iter = [it = iter]()-> auto {
            auto v = it();
            return v.has_value() ? std::optional<TP1>(v.value().first): std::optional<TP1>();
        };
        result.second.iter = [it = iter]()-> auto {
            auto v = it();
            return v.has_value() ? std::optional<TP2>(v.value().second): std::optional<TP2>();
        };
        return result;
    }

    auto cycle() {
        return iterator<T>{[it = iter, it2 = iter]() mutable -> auto {
            auto v = it(false);
            if(v.has_value())
                return v;
            it = it2;
            return it(false);
        }};
    }

    auto sum() {
        auto res = T(0);
        for (auto v = next(); v.has_value(); v = next())
            res += v.value();
        return res;
    }

    auto product() {
        auto res = T(1);
        for (auto v = next(); v.has_value(); v = next())
            res *= v.value();
        return res;
    }

    std::function<std::optional<T>()> iter;
};

template<typename T>
auto range(T start, T end) {
    return iterator<T>{[start, end]() mutable -> auto {
        if(start >= end)
            return std::optional<T>();
        return std::optional<T>(start++);
    }};
}

template<typename T>
auto repeat(T value) {
    return iterator<T>{[value]() -> auto {
        return std::optional<T>(value);
    }};
}

}  // namespace rs
#endif