length-prefixed-protocol-helper.hh

#pragma once
 
#include "lix/libutil/types.hh"
#include "lix/libutil/serialise.hh"
 
namespace nix {
 
class Store;

template<class Inner, typename T>
struct LengthPrefixedProtoHelper;

 
#define LENGTH_PREFIXED_PROTO_HELPER(Inner, T) \
    struct LengthPrefixedProtoHelper< Inner, T > \
    { \
        static T read(const Store & store, typename Inner::ReadConn conn); \
        [[nodiscard]] static WireFormatGenerator write(const Store & store, typename Inner::WriteConn conn, const T & str); \
    private: \
        template<typename U> using S = typename Inner::template Serialise<U>; \
    }
 
template<class Inner, typename T>
LENGTH_PREFIXED_PROTO_HELPER(Inner, std::vector<T>);
 
template<class Inner, typename T>
LENGTH_PREFIXED_PROTO_HELPER(Inner, std::set<T>);
 
template<class Inner, typename... Ts>
LENGTH_PREFIXED_PROTO_HELPER(Inner, std::tuple<Ts...>);
 
template<class Inner, typename K, typename V>
#define DONT_SUBSTITUTE_KV_TYPE std::map<K, V>
LENGTH_PREFIXED_PROTO_HELPER(Inner, DONT_SUBSTITUTE_KV_TYPE);
#undef DONT_SUBSTITUTE_KV_TYPE
 
template<class Inner, typename T>
std::vector<T>
LengthPrefixedProtoHelper<Inner, std::vector<T>>::read(
    const Store & store, typename Inner::ReadConn conn)
{
    std::vector<T> resSet;
    auto size = readNum<size_t>(conn.from);
    while (size--) {
        resSet.push_back(S<T>::read(store, conn));
    }
    return resSet;
}
 
template<class Inner, typename T>
WireFormatGenerator
LengthPrefixedProtoHelper<Inner, std::vector<T>>::write(
    const Store & store, typename Inner::WriteConn conn, const std::vector<T> & resSet)
{
    co_yield resSet.size();
    for (auto & key : resSet) {
        co_yield S<T>::write(store, conn, key);
    }
}
 
template<class Inner, typename T>
std::set<T>
LengthPrefixedProtoHelper<Inner, std::set<T>>::read(
    const Store & store, typename Inner::ReadConn conn)
{
    std::set<T> resSet;
    auto size = readNum<size_t>(conn.from);
    while (size--) {
        resSet.insert(S<T>::read(store, conn));
    }
    return resSet;
}
 
template<class Inner, typename T>
WireFormatGenerator
LengthPrefixedProtoHelper<Inner, std::set<T>>::write(
    const Store & store, typename Inner::WriteConn conn, const std::set<T> & resSet)
{
    co_yield resSet.size();
    for (auto & key : resSet) {
        co_yield S<T>::write(store, conn, key);
    }
}
 
template<class Inner, typename K, typename V>
std::map<K, V>
LengthPrefixedProtoHelper<Inner, std::map<K, V>>::read(
    const Store & store, typename Inner::ReadConn conn)
{
    std::map<K, V> resMap;
    auto size = readNum<size_t>(conn.from);
    while (size--) {
        auto k = S<K>::read(store, conn);
        auto v = S<V>::read(store, conn);
        resMap.insert_or_assign(std::move(k), std::move(v));
    }
    return resMap;
}
 
template<class Inner, typename K, typename V>
WireFormatGenerator
LengthPrefixedProtoHelper<Inner, std::map<K, V>>::write(
    const Store & store, typename Inner::WriteConn conn, const std::map<K, V> & resMap)
{
    co_yield resMap.size();
    for (auto & i : resMap) {
        co_yield S<K>::write(store, conn, i.first);
        co_yield S<V>::write(store, conn, i.second);
    }
}
 
template<class Inner, typename... Ts>
std::tuple<Ts...>
LengthPrefixedProtoHelper<Inner, std::tuple<Ts...>>::read(
    const Store & store, typename Inner::ReadConn conn)
{
    return std::tuple<Ts...> {
        S<Ts>::read(store, conn)...,
    };
}
 
template<class Inner, typename... Ts>
WireFormatGenerator
LengthPrefixedProtoHelper<Inner, std::tuple<Ts...>>::write(
    const Store & store, typename Inner::WriteConn conn, const std::tuple<Ts...> & res)
{
    auto fullArgs = std::apply(
        [&](auto &... rest) {
            return std::tuple<const Store &, typename Inner::WriteConn &, const Ts &...>(
                std::cref(store), conn, rest...
            );
        },
        res
    );
    return std::apply(
        []<typename... Us>(auto & store, auto conn, const Us &... args) -> WireFormatGenerator {
            (co_yield S<Us>::write(store, conn, args), ...);
        },
        fullArgs
    );
}
 
}