libutil: prototype serialization mechanism on coroutines

Change-Id: I361d89ff556354f6930d9204f55117565f2f7f20

src/libstore/daemon.cc

@@ -162,8 +162,7 @@ struct TunnelSink : Sink
     TunnelSink(Sink & to) : to(to) { }
     void operator () (std::string_view data)
     {
-        to << STDERR_WRITE;
+        to << STDERR_WRITE << data;
-        writeString(data, to);
     }
 };
 

src/libstore/remote-store.cc

@@ -1001,7 +1001,7 @@ std::exception_ptr RemoteStore::Connection::processStderr(Sink * sink, Source *
             if (!source) throw Error("no source");
             size_t len = readNum<size_t>(from);
             auto buf = std::make_unique<char[]>(len);
-            writeString({(const char *) buf.get(), source->read(buf.get(), len)}, to);
+            to << std::string_view((const char *) buf.get(), source->read(buf.get(), len));
             to.flush();
         }
 

src/libutil/serialise.cc

@@ -315,55 +315,43 @@ void writePadding(size_t len, Sink & sink)
 }
 
 
-void writeString(std::string_view data, Sink & sink)
+WireFormatGenerator SerializingTransform::operator()(std::string_view s)
 {
-    sink << data.size();
+    co_yield s.size();
-    sink(data);
+    co_yield raw(s.begin(), s.size());
-    writePadding(data.size(), sink);
+    if (s.size() % 8) {
+        std::array<char, 8> pad{};
+        co_yield raw(pad.begin(), 8 - s.size() % 8);
+    }
 }
 
-
+WireFormatGenerator SerializingTransform::operator()(const Strings & ss)
-Sink & operator << (Sink & sink, std::string_view s)
 {
-    writeString(s, sink);
+    co_yield ss.size();
-    return sink;
+    for (const auto & s : ss)
+        co_yield std::string_view(s);
 }
 
-
+WireFormatGenerator SerializingTransform::operator()(const StringSet & ss)
-template<class T> void writeStrings(const T & ss, Sink & sink)
 {
-    sink << ss.size();
+    co_yield ss.size();
-    for (auto & i : ss)
+    for (const auto & s : ss)
-        sink << i;
+        co_yield std::string_view(s);
 }
 
-Sink & operator << (Sink & sink, const Strings & s)
+WireFormatGenerator SerializingTransform::operator()(const Error & ex)
-{
-    writeStrings(s, sink);
-    return sink;
-}
-
-Sink & operator << (Sink & sink, const StringSet & s)
-{
-    writeStrings(s, sink);
-    return sink;
-}
-
-Sink & operator << (Sink & sink, const Error & ex)
 {
     auto & info = ex.info();
-    sink
+    co_yield "Error";
-        << "Error"
+    co_yield info.level;
-        << info.level
+    co_yield "Error"; // removed
-        << "Error" // removed
+    co_yield info.msg.str();
-        << info.msg.str()
+    co_yield 0; // FIXME: info.errPos
-        << 0 // FIXME: info.errPos
+    co_yield info.traces.size();
-        << info.traces.size();
     for (auto & trace : info.traces) {
-        sink << 0; // FIXME: trace.pos
+        co_yield 0; // FIXME: trace.pos
-        sink << trace.hint.str();
+        co_yield trace.hint.str();
     }
-    return sink;
 }
 
 

src/libutil/serialise.hh

@@ -1,8 +1,10 @@
 #pragma once
 ///@file
 
+#include <concepts>
 #include <memory>
 
+#include "generator.hh"
 #include "types.hh"
 #include "util.hh"
 
@@ -372,13 +374,28 @@ std::unique_ptr<Source> sinkToSource(
         throw EndOfFile("coroutine has finished");
     });
 
+struct SerializingTransform;
+using WireFormatGenerator = Generator<std::span<const char>, SerializingTransform>;
 
-void writePadding(size_t len, Sink & sink);
+inline void drainGenerator(WireFormatGenerator g, std::derived_from<Sink> auto & into)
-void writeString(std::string_view s, Sink & sink);
-
-inline Sink & operator << (Sink & sink, uint64_t n)
 {
-    unsigned char buf[8];
+    while (g) {
+        auto bit = g();
+        into(std::string_view(bit.data(), bit.size()));
+    }
+}
+
+struct SerializingTransform
+{
+    std::array<char, 8> buf;
+
+    static std::span<const char> raw(auto... args)
+    {
+        return std::span<const char>(args...);
+    }
+
+    std::span<const char> operator()(uint64_t n)
+    {
         buf[0] = n & 0xff;
         buf[1] = (n >> 8) & 0xff;
         buf[2] = (n >> 16) & 0xff;
@@ -387,19 +404,62 @@ inline Sink & operator << (Sink & sink, uint64_t n)
         buf[5] = (n >> 40) & 0xff;
         buf[6] = (n >> 48) & 0xff;
         buf[7] = (unsigned char) (n >> 56) & 0xff;
-    sink({(char *) buf, sizeof(buf)});
+        return {buf.begin(), 8};
+    }
+
+    // only choose this for *exactly* char spans, do not allow implicit
+    // conversions. this would cause ambiguities with strings literals,
+    // and resolving those with more string-like overloads needs a lot.
+    template<typename Span>
+        requires std::same_as<Span, std::span<char>> || std::same_as<Span, std::span<const char>>
+    std::span<const char> operator()(Span s)
+    {
+        return s;
+    }
+    WireFormatGenerator operator()(std::string_view s);
+    WireFormatGenerator operator()(const Strings & s);
+    WireFormatGenerator operator()(const StringSet & s);
+    WireFormatGenerator operator()(const Error & s);
+};
+
+inline Sink & operator<<(Sink & sink, WireFormatGenerator && g)
+{
+    while (g) {
+        auto bit = g();
+        sink(std::string_view(bit.data(), bit.size()));
+    }
     return sink;
 }
 
-Sink & operator << (Sink & in, const Error & ex);
+void writePadding(size_t len, Sink & sink);
-Sink & operator << (Sink & sink, std::string_view s);
-Sink & operator << (Sink & sink, const Strings & s);
-Sink & operator << (Sink & sink, const StringSet & s);
 
+inline Sink & operator<<(Sink & sink, uint64_t u)
+{
+    return sink << [&]() -> WireFormatGenerator { co_yield u; }();
+}
+
+inline Sink & operator<<(Sink & sink, std::string_view s)
+{
+    return sink << [&]() -> WireFormatGenerator { co_yield s; }();
+}
+
+inline Sink & operator<<(Sink & sink, const Strings & s)
+{
+    return sink << [&]() -> WireFormatGenerator { co_yield s; }();
+}
+
+inline Sink & operator<<(Sink & sink, const StringSet & s)
+{
+    return sink << [&]() -> WireFormatGenerator { co_yield s; }();
+}
+
+inline Sink & operator<<(Sink & sink, const Error & ex)
+{
+    return sink << [&]() -> WireFormatGenerator { co_yield ex; }();
+}
 
 MakeError(SerialisationError, Error);
 
-
 template<typename T>
 T readNum(Source & source)
 {

tests/unit/libutil/serialise.cc

@@ -1,15 +1,19 @@
 #include "serialise.hh"
 #include "error.hh"
 #include "fmt.hh"
+#include "generator.hh"
 #include "libexpr/pos-table.hh"
 #include "ref.hh"
 #include "types.hh"
 #include "util.hh"
 
+#include <concepts>
+#include <initializer_list>
 #include <limits.h>
 #include <gtest/gtest.h>
 
 #include <numeric>
+#include <stdexcept>
 
 namespace nix {
 
@@ -58,20 +62,29 @@ TEST(ChainSource, move)
     ASSERT_EQ(buf, "33");
 }
 
-TEST(Sink, uint64_t)
+static std::string simpleToWire(const auto & val)
 {
-    StringSink s;
+    std::string result;
-    s << 42;
+    auto g = [&] () -> WireFormatGenerator { co_yield val; }();
-    ASSERT_EQ(s.s, std::string({42, 0, 0, 0, 0, 0, 0, 0}));
+    while (g) {
+        auto bit = g();
+        result.append(bit.data(), bit.size());
+    }
+    return result;
 }
 
-TEST(Sink, string_view)
+TEST(WireFormatGenerator, uint64_t)
 {
-    StringSink s;
+    auto s = simpleToWire(42);
-    s << "";
+    ASSERT_EQ(s, std::string({42, 0, 0, 0, 0, 0, 0, 0}));
+}
+
+TEST(WireFormatGenerator, string_view)
+{
+    auto s = simpleToWire("");
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             0, 0, 0, 0, 0, 0, 0, 0,
@@ -80,11 +93,10 @@ TEST(Sink, string_view)
     );
     // clang-format on
 
-    s = {};
+    s = simpleToWire("test");
-    s << "test";
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             4, 0, 0, 0, 0, 0, 0, 0,
@@ -96,11 +108,10 @@ TEST(Sink, string_view)
     );
     // clang-format on
 
-    s = {};
+    s = simpleToWire("longer string");
-    s << "longer string";
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             13, 0, 0, 0, 0, 0, 0, 0,
@@ -113,13 +124,12 @@ TEST(Sink, string_view)
     // clang-format on
 }
 
-TEST(Sink, StringSet)
+TEST(WireFormatGenerator, StringSet)
 {
-    StringSink s;
+    auto s = simpleToWire(StringSet{});
-    s << StringSet{};
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             0, 0, 0, 0, 0, 0, 0, 0,
@@ -128,11 +138,10 @@ TEST(Sink, StringSet)
     );
     // clang-format on
 
-    s = {};
+    s = simpleToWire(StringSet{"a", ""});
-    s << StringSet{"a", ""};
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             2, 0, 0, 0, 0, 0, 0, 0,
@@ -145,13 +154,12 @@ TEST(Sink, StringSet)
     // clang-format on
 }
 
-TEST(Sink, Strings)
+TEST(WireFormatGenerator, Strings)
 {
-    StringSink s;
+    auto s = simpleToWire(Strings{});
-    s << Strings{};
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             0, 0, 0, 0, 0, 0, 0, 0,
@@ -160,11 +168,10 @@ TEST(Sink, Strings)
     );
     // clang-format on
 
-    s = {};
+    s = simpleToWire(Strings{"a", ""});
-    s << Strings{"a", ""};
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             // length
             2, 0, 0, 0, 0, 0, 0, 0,
@@ -177,23 +184,22 @@ TEST(Sink, Strings)
     // clang-format on
 }
 
-TEST(Sink, Error)
+TEST(WireFormatGenerator, Error)
 {
     PosTable pt;
     auto o = pt.addOrigin(Pos::String{make_ref<std::string>("test")}, 4);
 
-    StringSink s;
+    auto s = simpleToWire(Error{{
-    s << Error{{
         .level = lvlInfo,
         .msg = HintFmt("foo"),
         .pos = pt[pt.add(o, 1)],
         .traces = {{.pos = pt[pt.add(o, 2)], .hint = HintFmt("b %1%", "foo")}},
-    }};
+    }});
     // NOTE position of the error and all traces are ignored
     // by the wire format
     // clang-format off
     ASSERT_EQ(
-        s.s,
+        s,
         std::string({
             5, 0, 0, 0, 0, 0, 0, 0, 'E', 'r', 'r', 'o', 'r', 0, 0, 0,
             3, 0, 0, 0, 0, 0, 0, 0,
@@ -209,4 +215,46 @@ TEST(Sink, Error)
     // clang-format on
 }
 
+TEST(FullFormatter, foo)
+{
+    auto gen = []() -> Generator<std::span<const char>, SerializingTransform> {
+        std::set<std::string> foo{"a", "longer string", ""};
+        co_yield 42;
+        co_yield foo;
+        co_yield std::string_view("test");
+        co_yield 7;
+    }();
+
+    std::vector<char> full;
+    while (gen) {
+        auto s = gen();
+        full.insert(full.end(), s.begin(), s.end());
+    }
+
+    ASSERT_EQ(
+        full,
+        (std::vector<char>{
+            // clang-format off
+            // 32
+            42, 0, 0, 0, 0, 0, 0, 0,
+            // foo
+            3, 0, 0, 0, 0, 0, 0, 0,
+            /// ""
+            0, 0, 0, 0, 0, 0, 0, 0,
+            /// a
+            1, 0, 0, 0, 0, 0, 0, 0,
+            'a', 0, 0, 0, 0, 0, 0, 0,
+            /// longer string
+            13, 0, 0, 0, 0, 0, 0, 0,
+            'l', 'o', 'n', 'g', 'e', 'r', ' ', 's', 't', 'r', 'i', 'n', 'g', 0, 0, 0,
+            // foo done
+            // test
+            4, 0, 0, 0, 0, 0, 0, 0,
+            't', 'e', 's', 't', 0, 0, 0, 0,
+            // 7
+            7, 0, 0, 0, 0, 0, 0, 0,
+            //clang-format on
+            }));
+}
+
 }