From 85aacbee64680e60b22e8a729756b0a3882f8b8f Mon Sep 17 00:00:00 2001
From: John Ericson <John.Ericson@Obsidian.Systems>
Date: Thu, 13 Aug 2020 14:47:53 +0000
Subject: [PATCH] Use `TeeSink` and `TeeSouce` in a few more places

---
 src/libstore/binary-cache-store.cc | 12 ++++--------
 src/libstore/local-store.cc        |  6 +-----
 src/libstore/store-api.cc          |  6 +++---
 src/libutil/archive.cc             |  6 +-----
 src/libutil/serialise.hh           | 11 +++++++++++
 5 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/src/libstore/binary-cache-store.cc b/src/libstore/binary-cache-store.cc
index 86877dd1a..c360b9dda 100644
--- a/src/libstore/binary-cache-store.cc
+++ b/src/libstore/binary-cache-store.cc
@@ -312,14 +312,10 @@ void BinaryCacheStore::narFromPath(const StorePath & storePath, Sink & sink)
 {
     auto info = queryPathInfo(storePath).cast<const NarInfo>();
 
-    uint64_t narSize = 0;
+    LengthSink narSize;
+    TeeSink tee { sink, narSize };
 
-    LambdaSink wrapperSink([&](const unsigned char * data, size_t len) {
-        sink(data, len);
-        narSize += len;
-    });
-
-    auto decompressor = makeDecompressionSink(info->compression, wrapperSink);
+    auto decompressor = makeDecompressionSink(info->compression, tee);
 
     try {
         getFile(info->url, *decompressor);
@@ -331,7 +327,7 @@ void BinaryCacheStore::narFromPath(const StorePath & storePath, Sink & sink)
 
     stats.narRead++;
     //stats.narReadCompressedBytes += nar->size(); // FIXME
-    stats.narReadBytes += narSize;
+    stats.narReadBytes += narSize.length;
 }
 
 void BinaryCacheStore::queryPathInfoUncached(const StorePath & storePath,
diff --git a/src/libstore/local-store.cc b/src/libstore/local-store.cc
index 3c66a4dfd..e552bdc59 100644
--- a/src/libstore/local-store.cc
+++ b/src/libstore/local-store.cc
@@ -1021,11 +1021,7 @@ void LocalStore::addToStore(const ValidPathInfo & info, Source & source,
             else
                 hashSink = std::make_unique<HashModuloSink>(htSHA256, std::string(info.path.hashPart()));
 
-            LambdaSource wrapperSource([&](unsigned char * data, size_t len) -> size_t {
-                size_t n = source.read(data, len);
-                (*hashSink)(data, n);
-                return n;
-            });
+            TeeSource wrapperSource { source, *hashSink };
 
             restorePath(realPath, wrapperSource);
 
diff --git a/src/libstore/store-api.cc b/src/libstore/store-api.cc
index 2837910b9..3d07e2d38 100644
--- a/src/libstore/store-api.cc
+++ b/src/libstore/store-api.cc
@@ -746,12 +746,12 @@ void copyStorePath(ref<Store> srcStore, ref<Store> dstStore,
     }
 
     auto source = sinkToSource([&](Sink & sink) {
-        LambdaSink wrapperSink([&](const unsigned char * data, size_t len) {
-            sink(data, len);
+        LambdaSink progressSink([&](const unsigned char * data, size_t len) {
             total += len;
             act.progress(total, info->narSize);
         });
-        srcStore->narFromPath(storePath, wrapperSink);
+        TeeSink tee { sink, progressSink };
+        srcStore->narFromPath(storePath, tee);
     }, [&]() {
            throw EndOfFile("NAR for '%s' fetched from '%s' is incomplete", srcStore->printStorePath(storePath), srcStore->getUri());
     });
diff --git a/src/libutil/archive.cc b/src/libutil/archive.cc
index ce7cf9754..14399dea3 100644
--- a/src/libutil/archive.cc
+++ b/src/libutil/archive.cc
@@ -366,11 +366,7 @@ void copyNAR(Source & source, Sink & sink)
 
     ParseSink parseSink; /* null sink; just parse the NAR */
 
-    LambdaSource wrapper([&](unsigned char * data, size_t len) {
-        auto n = source.read(data, len);
-        sink(data, n);
-        return n;
-    });
+    TeeSource wrapper { source, sink };
 
     parseDump(parseSink, wrapper);
 }
diff --git a/src/libutil/serialise.hh b/src/libutil/serialise.hh
index c29c6b29b..69ae0874a 100644
--- a/src/libutil/serialise.hh
+++ b/src/libutil/serialise.hh
@@ -225,6 +225,17 @@ struct SizedSource : Source
     }
 };
 
+/* A sink that that just counts the number of bytes given to it */
+struct LengthSink : Sink
+{
+    uint64_t length = 0;
+
+    virtual void operator () (const unsigned char * _, size_t len)
+    {
+        length += len;
+    }
+};
+
 /* Convert a function into a sink. */
 struct LambdaSink : Sink
 {