From e94033c4af327e8a5e87a0d8d5c0b4fe1f67776d Mon Sep 17 00:00:00 2001
From: eldritch horrors <pennae@lix.systems>
Date: Mon, 18 Mar 2024 16:48:14 +0100
Subject: [PATCH] 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
---
 src/libstore/local-store.cc     |  3 +-
 src/libutil/serialise.cc        | 14 ---------
 src/libutil/serialise.hh        | 56 ++++++++++++++++++++++++++++-----
 tests/unit/libutil/serialise.cc | 56 +++++++++++++++++++++++++++++++++
 4 files changed, 105 insertions(+), 24 deletions(-)
 create mode 100644 tests/unit/libutil/serialise.cc

diff --git a/src/libstore/local-store.cc b/src/libstore/local-store.cc
index 29081244f..c8f194b92 100644
--- a/src/libstore/local-store.cc
+++ b/src/libstore/local-store.cc
@@ -1331,8 +1331,7 @@ StorePath LocalStore::addToStoreFromDump(Source & source0, std::string_view name
 
     if (!inMemory) {
         /* Drain what we pulled so far, and then keep on pulling */
-        StringSource dumpSource { dump };
-        ChainSource bothSource { dumpSource, source };
+        ChainSource bothSource { StringSource{dump}, std::move(source) };
 
         std::tie(tempDir, tempDirFd) = createTempDirInStore();
         delTempDir = std::make_unique<AutoDelete>(tempDir);
diff --git a/src/libutil/serialise.cc b/src/libutil/serialise.cc
index 6450a9651..222e88df8 100644
--- a/src/libutil/serialise.cc
+++ b/src/libutil/serialise.cc
@@ -449,18 +449,4 @@ void StringSink::operator () (std::string_view 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);
-        }
-    }
-}
-
 }
diff --git a/src/libutil/serialise.hh b/src/libutil/serialise.hh
index e6290a652..9b831fe56 100644
--- a/src/libutil/serialise.hh
+++ b/src/libutil/serialise.hh
@@ -306,18 +306,58 @@ struct LambdaSource : Source
 };
 
 /**
- * Chain two sources together so after the first is exhausted, the second is
- * used
+ * Chain a number of sources together, exhausting them all in turn.
  */
+template<typename... Sources>
+    requires (std::derived_from<Sources, Source> && ...)
 struct ChainSource : Source
 {
-    Source & source1, & source2;
-    bool useSecond = false;
-    ChainSource(Source & s1, Source & s2)
-        : source1(s1), source2(s2)
-    { }
+private:
+    std::tuple<Sources...> sources;
+    std::array<Source *, sizeof...(Sources)> ptrs;
+    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);
diff --git a/tests/unit/libutil/serialise.cc b/tests/unit/libutil/serialise.cc
new file mode 100644
index 000000000..71d99b9a7
--- /dev/null
+++ b/tests/unit/libutil/serialise.cc
@@ -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");
+}
+
+}