libutil: allow construction of sources from generators

Change-Id: I78ff8d0720f06bce731e26d5e1c53b1382bbd589
This commit is contained in:
eldritch horrors 2024-06-27 21:23:50 +02:00
parent b9f91ec3c5
commit b51ea465de
2 changed files with 49 additions and 0 deletions

View file

@ -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);
/**

View file

@ -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);
}
}