Merge pull request #5331 from edolstra/references

Add a test for RefScanSink and clean up the code

6 files changed, 126 insertions, 55 deletions

diff --git a/Makefile b/Makefile
index c7d8967c8..5040d2884 100644
--- a/Makefile
+++ b/Makefile
@@ -4,6 +4,7 @@ makefiles = \
   src/libutil/local.mk \
   src/libutil/tests/local.mk \
   src/libstore/local.mk \
+  src/libstore/tests/local.mk \
   src/libfetchers/local.mk \
   src/libmain/local.mk \
   src/libexpr/local.mk \
diff --git a/src/libstore/build/local-derivation-goal.cc b/src/libstore/build/local-derivation-goal.cc
index b1852a6bb..2ba1b3f59 100644
--- a/src/libstore/build/local-derivation-goal.cc
+++ b/src/libstore/build/local-derivation-goal.cc
@@ -2140,8 +2140,7 @@ void LocalDerivationGoal::registerOutputs()
 
         /* Pass blank Sink as we are not ready to hash data at this stage. */
         NullSink blank;
-        auto references = worker.store.parseStorePathSet(
-            scanForReferences(blank, actualPath, worker.store.printStorePathSet(referenceablePaths)));
+        auto references = scanForReferences(blank, actualPath, referenceablePaths);
 
         outputReferencesIfUnregistered.insert_or_assign(
             outputName,
diff --git a/src/libstore/references.cc b/src/libstore/references.cc
index 3a07c1411..c369b14ac 100644
--- a/src/libstore/references.cc
+++ b/src/libstore/references.cc
@@ -11,11 +11,13 @@
 namespace nix {
 
 
-static unsigned int refLength = 32; /* characters */
+static size_t refLength = 32; /* characters */
 
 
-static void search(const unsigned char * s, size_t len,
-    StringSet & hashes, StringSet & seen)
+static void search(
+    std::string_view s,
+    StringSet & hashes,
+    StringSet & seen)
 {
     static std::once_flag initialised;
     static bool isBase32[256];
@@ -25,7 +27,7 @@ static void search(const unsigned char * s, size_t len,
             isBase32[(unsigned char) base32Chars[i]] = true;
     });
 
-    for (size_t i = 0; i + refLength <= len; ) {
+    for (size_t i = 0; i + refLength <= s.size(); ) {
         int j;
         bool match = true;
         for (j = refLength - 1; j >= 0; --j)
@@ -35,7 +37,7 @@ static void search(const unsigned char * s, size_t len,
                 break;
             }
         if (!match) continue;
-        string ref((const char *) s + i, refLength);
+        std::string ref(s.substr(i, refLength));
         if (hashes.erase(ref)) {
             debug(format("found reference to '%1%' at offset '%2%'")
                   % ref % i);
@@ -46,69 +48,60 @@ static void search(const unsigned char * s, size_t len,
 }
 
 
-struct RefScanSink : Sink
+void RefScanSink::operator () (std::string_view data)
 {
-    StringSet hashes;
-    StringSet seen;
-
-    string tail;
-
-    RefScanSink() { }
-
-    void operator () (std::string_view data) override
-    {
-        /* It's possible that a reference spans the previous and current
-           fragment, so search in the concatenation of the tail of the
-           previous fragment and the start of the current fragment. */
-        string s = tail + std::string(data, 0, refLength);
-        search((const unsigned char *) s.data(), s.size(), hashes, seen);
-
-        search((const unsigned char *) data.data(), data.size(), hashes, seen);
-
-        size_t tailLen = data.size() <= refLength ? data.size() : refLength;
-        tail = std::string(tail, tail.size() < refLength - tailLen ? 0 : tail.size() - (refLength - tailLen));
-        tail.append({data.data() + data.size() - tailLen, tailLen});
-    }
-};
+    /* It's possible that a reference spans the previous and current
+       fragment, so search in the concatenation of the tail of the
+       previous fragment and the start of the current fragment. */
+    auto s = tail;
+    s.append(data.data(), refLength);
+    search(s, hashes, seen);
+
+    search(data, hashes, seen);
+
+    auto tailLen = std::min(data.size(), refLength);
+    auto rest = refLength - tailLen;
+    if (rest < tail.size())
+        tail = tail.substr(tail.size() - rest);
+    tail.append(data.data() + data.size() - tailLen, tailLen);
+}
 
 
-std::pair<PathSet, HashResult> scanForReferences(const string & path,
-    const PathSet & refs)
+std::pair<StorePathSet, HashResult> scanForReferences(
+    const string & path,
+    const StorePathSet & refs)
 {
     HashSink hashSink { htSHA256 };
     auto found = scanForReferences(hashSink, path, refs);
     auto hash = hashSink.finish();
-    return std::pair<PathSet, HashResult>(found, hash);
+    return std::pair<StorePathSet, HashResult>(found, hash);
 }
 
-PathSet scanForReferences(Sink & toTee,
-    const string & path, const PathSet & refs)
+StorePathSet scanForReferences(
+    Sink & toTee,
+    const Path & path,
+    const StorePathSet & refs)
 {
-    RefScanSink refsSink;
-    TeeSink sink { refsSink, toTee };
-    std::map<string, Path> backMap;
+    StringSet hashes;
+    std::map<std::string, StorePath> backMap;
 
     for (auto & i : refs) {
-        auto baseName = std::string(baseNameOf(i));
-        string::size_type pos = baseName.find('-');
-        if (pos == string::npos)
-            throw Error("bad reference '%1%'", i);
-        string s = string(baseName, 0, pos);
-        assert(s.size() == refLength);
-        assert(backMap.find(s) == backMap.end());
-        // parseHash(htSHA256, s);
-        refsSink.hashes.insert(s);
-        backMap[s] = i;
+        std::string hashPart(i.hashPart());
+        auto inserted = backMap.emplace(hashPart, i).second;
+        assert(inserted);
+        hashes.insert(hashPart);
     }
 
     /* Look for the hashes in the NAR dump of the path. */
+    RefScanSink refsSink(std::move(hashes));
+    TeeSink sink { refsSink, toTee };
     dumpPath(path, sink);
 
     /* Map the hashes found back to their store paths. */
-    PathSet found;
-    for (auto & i : refsSink.seen) {
-        std::map<string, Path>::iterator j;
-        if ((j = backMap.find(i)) == backMap.end()) abort();
+    StorePathSet found;
+    for (auto & i : refsSink.getResult()) {
+        auto j = backMap.find(i);
+        assert(j != backMap.end());
         found.insert(j->second);
     }
 
diff --git a/src/libstore/references.hh b/src/libstore/references.hh
index 4f12e6b21..a6119c861 100644
--- a/src/libstore/references.hh
+++ b/src/libstore/references.hh
@@ -1,13 +1,31 @@
 #pragma once
 
-#include "types.hh"
 #include "hash.hh"
+#include "path.hh"
 
 namespace nix {
 
-std::pair<PathSet, HashResult> scanForReferences(const Path & path, const PathSet & refs);
+std::pair<StorePathSet, HashResult> scanForReferences(const Path & path, const StorePathSet & refs);
 
-PathSet scanForReferences(Sink & toTee, const Path & path, const PathSet & refs);
+StorePathSet scanForReferences(Sink & toTee, const Path & path, const StorePathSet & refs);
+
+class RefScanSink : public Sink
+{
+    StringSet hashes;
+    StringSet seen;
+
+    std::string tail;
+
+public:
+
+    RefScanSink(StringSet && hashes) : hashes(hashes)
+    { }
+
+    StringSet & getResult()
+    { return seen; }
+
+    void operator () (std::string_view data) override;
+};
 
 struct RewritingSink : Sink
 {
diff --git a/src/libstore/tests/local.mk b/src/libstore/tests/local.mk
new file mode 100644
index 000000000..6ae9a0285
--- /dev/null
+++ b/src/libstore/tests/local.mk
@@ -0,0 +1,15 @@
+check: libstore-tests_RUN
+
+programs += libstore-tests
+
+libstore-tests_DIR := $(d)
+
+libstore-tests_INSTALL_DIR :=
+
+libstore-tests_SOURCES := $(wildcard $(d)/*.cc)
+
+libstore-tests_CXXFLAGS += -I src/libstore -I src/libutil
+
+libstore-tests_LIBS = libstore
+
+libstore-tests_LDFLAGS := $(GTEST_LIBS)
diff --git a/src/libstore/tests/references.cc b/src/libstore/tests/references.cc
new file mode 100644
index 000000000..d91d1cedd
--- /dev/null
+++ b/src/libstore/tests/references.cc
@@ -0,0 +1,45 @@
+#include "references.hh"
+
+#include <gtest/gtest.h>
+
+namespace nix {
+
+TEST(references, scan)
+{
+    std::string hash1 = "dc04vv14dak1c1r48qa0m23vr9jy8sm0";
+    std::string hash2 = "zc842j0rz61mjsp3h3wp5ly71ak6qgdn";
+
+    {
+        RefScanSink scanner(StringSet{hash1});
+        auto s = "foobar";
+        scanner(s);
+        ASSERT_EQ(scanner.getResult(), StringSet{});
+    }
+
+    {
+        RefScanSink scanner(StringSet{hash1});
+        auto s = "foobar" + hash1 + "xyzzy";
+        scanner(s);
+        ASSERT_EQ(scanner.getResult(), StringSet{hash1});
+    }
+
+    {
+        RefScanSink scanner(StringSet{hash1, hash2});
+        auto s = "foobar" + hash1 + "xyzzy" + hash2;
+        scanner(((std::string_view) s).substr(0, 10));
+        scanner(((std::string_view) s).substr(10, 5));
+        scanner(((std::string_view) s).substr(15, 5));
+        scanner(((std::string_view) s).substr(20));
+        ASSERT_EQ(scanner.getResult(), StringSet({hash1, hash2}));
+    }
+
+    {
+        RefScanSink scanner(StringSet{hash1, hash2});
+        auto s = "foobar" + hash1 + "xyzzy" + hash2;
+        for (auto & i : s)
+            scanner(std::string(1, i));
+        ASSERT_EQ(scanner.getResult(), StringSet({hash1, hash2}));
+    }
+}
+
+}