libutil: allow construction of sources from generators

Change-Id: I78ff8d0720f06bce731e26d5e1c53b1382bbd589

2 files changed, 49 insertions, 0 deletions

diff --git a/src/libutil/serialise.hh b/src/libutil/serialise.hh
index 2651ec979..874c67b75 100644
--- a/src/libutil/serialise.hh
+++ b/src/libutil/serialise.hh
@@ -334,6 +334,34 @@ struct ChainSource : Source
     size_t read(char * data, size_t len) override;
 };
 
+struct GeneratorSource : Source
+{
+    GeneratorSource(Generator<Bytes> && g) : g(std::move(g)) {}
+
+    virtual size_t read(char * data, size_t len)
+    {
+        // we explicitly do not poll the generator multiple times to fill the
+        // buffer, only to produce some output at all. this is allowed by the
+        // semantics of read(), only operator() must fill the buffer entirely
+        while (!buf.size()) {
+            if (auto next = g.next()) {
+                buf = *next;
+            } else {
+                throw EndOfFile("coroutine has finished");
+            }
+        }
+
+        len = std::min(len, buf.size());
+        memcpy(data, buf.data(), len);
+        buf = buf.subspan(len);
+        return len;
+    }
+
+private:
+    Generator<Bytes> g;
+    Bytes buf{};
+};
+
 std::unique_ptr<FinishSink> sourceToSink(std::function<void(Source &)> fun);
 
 /**
diff --git a/tests/unit/libutil/serialise.cc b/tests/unit/libutil/serialise.cc
index 78882ad2c..4fe010af8 100644
--- a/tests/unit/libutil/serialise.cc
+++ b/tests/unit/libutil/serialise.cc
@@ -214,4 +214,25 @@ TEST(WireFormatGenerator, exampleMessage)
             }));
 }
 
+TEST(GeneratorSource, works)
+{
+    GeneratorSource src = []() -> Generator<Bytes> {
+        co_yield std::span{"", 0};
+        co_yield std::span{"a", 1};
+        co_yield std::span{"", 0};
+        co_yield std::span{"bcd", 3};
+        co_yield std::span{"", 0};
+    }();
+
+    char buf[2];
+    ASSERT_EQ(src.read(buf, sizeof(buf)), 1);
+    ASSERT_EQ(buf[0], 'a');
+    ASSERT_EQ(src.read(buf, sizeof(buf)), 2);
+    ASSERT_EQ(buf[0], 'b');
+    ASSERT_EQ(buf[1], 'c');
+    ASSERT_EQ(src.read(buf, sizeof(buf)), 1);
+    ASSERT_EQ(buf[0], 'd');
+    ASSERT_THROW(src.read(buf, sizeof(buf)), EndOfFile);
+}
+
 }