libutil: make ChainSource *more*

- make it chain more than two sources together
- make it own its parts
- make it properly moveable
- test it

this is the first step towards eliminating a bunch of produce-to-sink
functions we have right now (such as readFile, dumpPath, etc).

Change-Id: I518eb7a7f9e1e1a9199a410074c83b77b38c8a06
This commit is contained in:
eldritch horrors 2024-03-18 16:48:14 +01:00
parent b627b0b432
commit e94033c4af
4 changed files with 105 additions and 24 deletions

View file

@ -1331,8 +1331,7 @@ StorePath LocalStore::addToStoreFromDump(Source & source0, std::string_view name
if (!inMemory) { if (!inMemory) {
/* Drain what we pulled so far, and then keep on pulling */ /* Drain what we pulled so far, and then keep on pulling */
StringSource dumpSource { dump }; ChainSource bothSource { StringSource{dump}, std::move(source) };
ChainSource bothSource { dumpSource, source };
std::tie(tempDir, tempDirFd) = createTempDirInStore(); std::tie(tempDir, tempDirFd) = createTempDirInStore();
delTempDir = std::make_unique<AutoDelete>(tempDir); delTempDir = std::make_unique<AutoDelete>(tempDir);

View file

@ -449,18 +449,4 @@ void StringSink::operator () (std::string_view data)
s.append(data); s.append(data);
} }
size_t ChainSource::read(char * data, size_t len)
{
if (useSecond) {
return source2.read(data, len);
} else {
try {
return source1.read(data, len);
} catch (EndOfFile &) {
useSecond = true;
return this->read(data, len);
}
}
}
} }

View file

@ -306,18 +306,58 @@ struct LambdaSource : Source
}; };
/** /**
* Chain two sources together so after the first is exhausted, the second is * Chain a number of sources together, exhausting them all in turn.
* used
*/ */
template<typename... Sources>
requires (std::derived_from<Sources, Source> && ...)
struct ChainSource : Source struct ChainSource : Source
{ {
Source & source1, & source2; private:
bool useSecond = false; std::tuple<Sources...> sources;
ChainSource(Source & s1, Source & s2) std::array<Source *, sizeof...(Sources)> ptrs;
: source1(s1), source2(s2) size_t sourceIdx = 0;
{ }
size_t read(char * data, size_t len) override; template<size_t... N>
void fillPtrs(std::index_sequence<N...>)
{
((ptrs[N] = &std::get<N>(sources)), ...);
}
public:
ChainSource(Sources && ... sources)
: sources(std::move(sources)...)
{
fillPtrs(std::index_sequence_for<Sources...>{});
}
ChainSource(ChainSource && other)
: sources(std::move(other.sources))
, sourceIdx(other.sourceIdx)
{
fillPtrs(std::index_sequence_for<Sources...>{});
other.sourceIdx = sizeof...(Sources);
}
ChainSource & operator=(ChainSource && other)
{
std::swap(sources, other.sources);
// since Sources... are the same the tuple type and offsets
// are the same, so pointers remain valid on both sides.
std::swap(sourceIdx, other.sourceIdx);
return *this;
}
size_t read(char * data, size_t len) override
{
if (sourceIdx == sizeof...(Sources))
throw EndOfFile("reached end of chained sources");
try {
return ptrs[sourceIdx]->read(data, len);
} catch (EndOfFile &) {
sourceIdx++;
return this->read(data, len);
}
}
}; };
std::unique_ptr<FinishSink> sourceToSink(std::function<void(Source &)> fun); std::unique_ptr<FinishSink> sourceToSink(std::function<void(Source &)> fun);

View file

@ -0,0 +1,56 @@
#include "serialise.hh"
#include "types.hh"
#include <limits.h>
#include <gtest/gtest.h>
#include <numeric>
namespace nix {
TEST(ChainSource, single)
{
ChainSource s{StringSource{"test"}};
ASSERT_EQ(s.drain(), "test");
}
TEST(ChainSource, multiple)
{
ChainSource s{StringSource{"1"}, StringSource{""}, StringSource{"3"}};
ASSERT_EQ(s.drain(), "13");
}
TEST(ChainSource, chunk)
{
std::string buf(2, ' ');
ChainSource s{StringSource{"111"}, StringSource{""}, StringSource{"333"}};
s(buf.data(), buf.size());
ASSERT_EQ(buf, "11");
s(buf.data(), buf.size());
ASSERT_EQ(buf, "13");
s(buf.data(), buf.size());
ASSERT_EQ(buf, "33");
ASSERT_THROW(s(buf.data(), buf.size()), EndOfFile);
}
TEST(ChainSource, move)
{
std::string buf(2, ' ');
ChainSource s1{StringSource{"111"}, StringSource{""}, StringSource{"333"}};
s1(buf.data(), buf.size());
ASSERT_EQ(buf, "11");
ChainSource s2 = std::move(s1);
ASSERT_THROW(s1(buf.data(), buf.size()), EndOfFile);
s2(buf.data(), buf.size());
ASSERT_EQ(buf, "13");
s1 = std::move(s2);
ASSERT_THROW(s2(buf.data(), buf.size()), EndOfFile);
s1(buf.data(), buf.size());
ASSERT_EQ(buf, "33");
}
}