Make `RewritingSink` accept a map of rewrites

Giving it the same semantics as `rewriteStrings`.
Also add some tests for it

7 files changed, 166 insertions, 85 deletions

diff --git a/src/libexpr/primops.cc b/src/libexpr/primops.cc
index cfae1e5f8..572e76ec6 100644
--- a/src/libexpr/primops.cc
+++ b/src/libexpr/primops.cc
@@ -6,7 +6,7 @@
 #include "globals.hh"
 #include "json-to-value.hh"
 #include "names.hh"
-#include "references.hh"
+#include "path-references.hh"
 #include "store-api.hh"
 #include "util.hh"
 #include "value-to-json.hh"
diff --git a/src/libstore/build/local-derivation-goal.cc b/src/libstore/build/local-derivation-goal.cc
index 05d6685da..6f7ab8a3d 100644
--- a/src/libstore/build/local-derivation-goal.cc
+++ b/src/libstore/build/local-derivation-goal.cc
@@ -4,7 +4,7 @@
 #include "worker.hh"
 #include "builtins.hh"
 #include "builtins/buildenv.hh"
-#include "references.hh"
+#include "path-references.hh"
 #include "finally.hh"
 #include "util.hh"
 #include "archive.hh"
diff --git a/src/libstore/path-references.cc b/src/libstore/path-references.cc
new file mode 100644
index 000000000..33cf66ce3
--- /dev/null
+++ b/src/libstore/path-references.cc
@@ -0,0 +1,73 @@
+#include "path-references.hh"
+#include "hash.hh"
+#include "util.hh"
+#include "archive.hh"
+
+#include <map>
+#include <cstdlib>
+#include <mutex>
+#include <algorithm>
+
+
+namespace nix {
+
+
+PathRefScanSink::PathRefScanSink(StringSet && hashes, std::map<std::string, StorePath> && backMap)
+    : RefScanSink(std::move(hashes))
+    , backMap(std::move(backMap))
+{ }
+
+PathRefScanSink PathRefScanSink::fromPaths(const StorePathSet & refs)
+{
+    StringSet hashes;
+    std::map<std::string, StorePath> backMap;
+
+    for (auto & i : refs) {
+        std::string hashPart(i.hashPart());
+        auto inserted = backMap.emplace(hashPart, i).second;
+        assert(inserted);
+        hashes.insert(hashPart);
+    }
+
+    return PathRefScanSink(std::move(hashes), std::move(backMap));
+}
+
+StorePathSet PathRefScanSink::getResultPaths()
+{
+    /* Map the hashes found back to their store paths. */
+    StorePathSet found;
+    for (auto & i : getResult()) {
+        auto j = backMap.find(i);
+        assert(j != backMap.end());
+        found.insert(j->second);
+    }
+
+    return found;
+}
+
+
+std::pair<StorePathSet, HashResult> scanForReferences(
+    const std::string & path,
+    const StorePathSet & refs)
+{
+    HashSink hashSink { htSHA256 };
+    auto found = scanForReferences(hashSink, path, refs);
+    auto hash = hashSink.finish();
+    return std::pair<StorePathSet, HashResult>(found, hash);
+}
+
+StorePathSet scanForReferences(
+    Sink & toTee,
+    const Path & path,
+    const StorePathSet & refs)
+{
+    PathRefScanSink refsSink = PathRefScanSink::fromPaths(refs);
+    TeeSink sink { refsSink, toTee };
+
+    /* Look for the hashes in the NAR dump of the path. */
+    dumpPath(path, sink);
+
+    return refsSink.getResultPaths();
+}
+
+}
diff --git a/src/libstore/path-references.hh b/src/libstore/path-references.hh
new file mode 100644
index 000000000..7b44e3261
--- /dev/null
+++ b/src/libstore/path-references.hh
@@ -0,0 +1,25 @@
+#pragma once
+
+#include "references.hh"
+#include "path.hh"
+
+namespace nix {
+
+std::pair<StorePathSet, HashResult> scanForReferences(const Path & path, const StorePathSet & refs);
+
+StorePathSet scanForReferences(Sink & toTee, const Path & path, const StorePathSet & refs);
+
+class PathRefScanSink : public RefScanSink
+{
+    std::map<std::string, StorePath> backMap;
+
+    PathRefScanSink(StringSet && hashes, std::map<std::string, StorePath> && backMap);
+
+public:
+
+    static PathRefScanSink fromPaths(const StorePathSet & refs);
+
+    StorePathSet getResultPaths();
+};
+
+}
diff --git a/src/libstore/references.cc b/src/libutil/references.cc
index 345f4528b..74003584a 100644
--- a/src/libstore/references.cc
+++ b/src/libutil/references.cc
@@ -6,6 +6,7 @@
 #include <map>
 #include <cstdlib>
 #include <mutex>
+#include <algorithm>
 
 
 namespace nix {
@@ -66,69 +67,20 @@ void RefScanSink::operator () (std::string_view data)
 }
 
 
-PathRefScanSink::PathRefScanSink(StringSet && hashes, std::map<std::string, StorePath> && backMap)
-    : RefScanSink(std::move(hashes))
-    , backMap(std::move(backMap))
-{ }
-
-PathRefScanSink PathRefScanSink::fromPaths(const StorePathSet & refs)
+RewritingSink::RewritingSink(const std::string & from, const std::string & to, Sink & nextSink)
+    : RewritingSink({{from, to}}, nextSink)
 {
-    StringSet hashes;
-    std::map<std::string, StorePath> backMap;
-
-    for (auto & i : refs) {
-        std::string hashPart(i.hashPart());
-        auto inserted = backMap.emplace(hashPart, i).second;
-        assert(inserted);
-        hashes.insert(hashPart);
-    }
-
-    return PathRefScanSink(std::move(hashes), std::move(backMap));
 }
 
-StorePathSet PathRefScanSink::getResultPaths()
+RewritingSink::RewritingSink(const StringMap & rewrites, Sink & nextSink)
+    : rewrites(rewrites), nextSink(nextSink)
 {
-    /* Map the hashes found back to their store paths. */
-    StorePathSet found;
-    for (auto & i : getResult()) {
-        auto j = backMap.find(i);
-        assert(j != backMap.end());
-        found.insert(j->second);
+    long unsigned int maxRewriteSize = 0;
+    for (auto & [from, to] : rewrites) {
+        assert(from.size() == to.size());
+        maxRewriteSize = std::max(maxRewriteSize, from.size());
     }
-
-    return found;
-}
-
-
-std::pair<StorePathSet, HashResult> scanForReferences(
-    const std::string & path,
-    const StorePathSet & refs)
-{
-    HashSink hashSink { htSHA256 };
-    auto found = scanForReferences(hashSink, path, refs);
-    auto hash = hashSink.finish();
-    return std::pair<StorePathSet, HashResult>(found, hash);
-}
-
-StorePathSet scanForReferences(
-    Sink & toTee,
-    const Path & path,
-    const StorePathSet & refs)
-{
-    PathRefScanSink refsSink = PathRefScanSink::fromPaths(refs);
-    TeeSink sink { refsSink, toTee };
-
-    /* Look for the hashes in the NAR dump of the path. */
-    dumpPath(path, sink);
-
-    return refsSink.getResultPaths();
-}
-
-
-RewritingSink::RewritingSink(const std::string & from, const std::string & to, Sink & nextSink)
-    : from(from), to(to), nextSink(nextSink)
-{
-    assert(from.size() == to.size());
+    this->maxRewriteSize = maxRewriteSize;
 }
 
 void RewritingSink::operator () (std::string_view data)
@@ -136,13 +88,13 @@ void RewritingSink::operator () (std::string_view data)
     std::string s(prev);
     s.append(data);
 
-    size_t j = 0;
-    while ((j = s.find(from, j)) != std::string::npos) {
-        matches.push_back(pos + j);
-        s.replace(j, from.size(), to);
-    }
+    s = rewriteStrings(s, rewrites);
 
-    prev = s.size() < from.size() ? s : std::string(s, s.size() - from.size() + 1, from.size() - 1);
+    prev = s.size() < maxRewriteSize
+        ? s
+        : maxRewriteSize == 0
+            ? ""
+            : std::string(s, s.size() - maxRewriteSize + 1, maxRewriteSize - 1);
 
     auto consumed = s.size() - prev.size();
 
diff --git a/src/libstore/references.hh b/src/libutil/references.hh
index 52d71b333..ffd730e7b 100644
--- a/src/libstore/references.hh
+++ b/src/libutil/references.hh
@@ -2,14 +2,9 @@
 ///@file
 
 #include "hash.hh"
-#include "path.hh"
 
 namespace nix {
 
-std::pair<StorePathSet, HashResult> scanForReferences(const Path & path, const StorePathSet & refs);
-
-StorePathSet scanForReferences(Sink & toTee, const Path & path, const StorePathSet & refs);
-
 class RefScanSink : public Sink
 {
     StringSet hashes;
@@ -28,28 +23,18 @@ public:
     void operator () (std::string_view data) override;
 };
 
-class PathRefScanSink : public RefScanSink
-{
-    std::map<std::string, StorePath> backMap;
-
-    PathRefScanSink(StringSet && hashes, std::map<std::string, StorePath> && backMap);
-
-public:
-
-    static PathRefScanSink fromPaths(const StorePathSet & refs);
-
-    StorePathSet getResultPaths();
-};
-
 struct RewritingSink : Sink
 {
-    std::string from, to, prev;
+    const StringMap rewrites;
+    long unsigned int maxRewriteSize;
+    std::string prev;
     Sink & nextSink;
     uint64_t pos = 0;
 
     std::vector<uint64_t> matches;
 
     RewritingSink(const std::string & from, const std::string & to, Sink & nextSink);
+    RewritingSink(const StringMap & rewrites, Sink & nextSink);
 
     void operator () (std::string_view data) override;
 
diff --git a/src/libutil/tests/references.cc b/src/libutil/tests/references.cc
new file mode 100644
index 000000000..a517d9aa1
--- /dev/null
+++ b/src/libutil/tests/references.cc
@@ -0,0 +1,46 @@
+#include "references.hh"
+#include <gtest/gtest.h>
+
+namespace nix {
+
+using std::string;
+
+struct RewriteParams {
+    string originalString, finalString;
+    StringMap rewrites;
+
+    friend std::ostream& operator<<(std::ostream& os, const RewriteParams& bar) {
+        StringSet strRewrites;
+        for (auto & [from, to] : bar.rewrites)
+            strRewrites.insert(from + "->" + to);
+        return os <<
+            "OriginalString: " << bar.originalString << std::endl <<
+            "Rewrites: " << concatStringsSep(",", strRewrites) << std::endl <<
+            "Expected result: " << bar.finalString;
+    }
+};
+
+class RewriteTest : public ::testing::TestWithParam<RewriteParams> {
+};
+
+TEST_P(RewriteTest, IdentityRewriteIsIdentity) {
+    RewriteParams param = GetParam();
+    StringSink rewritten;
+    auto rewriter = RewritingSink(param.rewrites, rewritten);
+    rewriter(param.originalString);
+    rewriter.flush();
+    ASSERT_EQ(rewritten.s, param.finalString);
+}
+
+INSTANTIATE_TEST_CASE_P(
+    references,
+    RewriteTest,
+    ::testing::Values(
+        RewriteParams{ "foooo", "baroo", {{"foo", "bar"}, {"bar", "baz"}}},
+        RewriteParams{ "foooo", "bazoo", {{"fou", "bar"}, {"foo", "baz"}}},
+        RewriteParams{ "foooo", "foooo", {}}
+    )
+);
+
+}
+