Make `RewritingSink` accept a map of rewrites

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

3 files changed, 242 insertions, 0 deletions

diff --git a/src/libutil/references.cc b/src/libutil/references.cc
new file mode 100644
index 000000000..74003584a
--- /dev/null
+++ b/src/libutil/references.cc
@@ -0,0 +1,140 @@
+#include "references.hh"
+#include "hash.hh"
+#include "util.hh"
+#include "archive.hh"
+
+#include <map>
+#include <cstdlib>
+#include <mutex>
+#include <algorithm>
+
+
+namespace nix {
+
+
+static size_t refLength = 32; /* characters */
+
+
+static void search(
+    std::string_view s,
+    StringSet & hashes,
+    StringSet & seen)
+{
+    static std::once_flag initialised;
+    static bool isBase32[256];
+    std::call_once(initialised, [](){
+        for (unsigned int i = 0; i < 256; ++i) isBase32[i] = false;
+        for (unsigned int i = 0; i < base32Chars.size(); ++i)
+            isBase32[(unsigned char) base32Chars[i]] = true;
+    });
+
+    for (size_t i = 0; i + refLength <= s.size(); ) {
+        int j;
+        bool match = true;
+        for (j = refLength - 1; j >= 0; --j)
+            if (!isBase32[(unsigned char) s[i + j]]) {
+                i += j + 1;
+                match = false;
+                break;
+            }
+        if (!match) continue;
+        std::string ref(s.substr(i, refLength));
+        if (hashes.erase(ref)) {
+            debug("found reference to '%1%' at offset '%2%'", ref, i);
+            seen.insert(ref);
+        }
+        ++i;
+    }
+}
+
+
+void RefScanSink::operator () (std::string_view data)
+{
+    /* 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;
+    auto tailLen = std::min(data.size(), refLength);
+    s.append(data.data(), tailLen);
+    search(s, hashes, seen);
+
+    search(data, hashes, seen);
+
+    auto rest = refLength - tailLen;
+    if (rest < tail.size())
+        tail = tail.substr(tail.size() - rest);
+    tail.append(data.data() + data.size() - tailLen, tailLen);
+}
+
+
+RewritingSink::RewritingSink(const std::string & from, const std::string & to, Sink & nextSink)
+    : RewritingSink({{from, to}}, nextSink)
+{
+}
+
+RewritingSink::RewritingSink(const StringMap & rewrites, Sink & nextSink)
+    : rewrites(rewrites), nextSink(nextSink)
+{
+    long unsigned int maxRewriteSize = 0;
+    for (auto & [from, to] : rewrites) {
+        assert(from.size() == to.size());
+        maxRewriteSize = std::max(maxRewriteSize, from.size());
+    }
+    this->maxRewriteSize = maxRewriteSize;
+}
+
+void RewritingSink::operator () (std::string_view data)
+{
+    std::string s(prev);
+    s.append(data);
+
+    s = rewriteStrings(s, rewrites);
+
+    prev = s.size() < maxRewriteSize
+        ? s
+        : maxRewriteSize == 0
+            ? ""
+            : std::string(s, s.size() - maxRewriteSize + 1, maxRewriteSize - 1);
+
+    auto consumed = s.size() - prev.size();
+
+    pos += consumed;
+
+    if (consumed) nextSink(s.substr(0, consumed));
+}
+
+void RewritingSink::flush()
+{
+    if (prev.empty()) return;
+    pos += prev.size();
+    nextSink(prev);
+    prev.clear();
+}
+
+HashModuloSink::HashModuloSink(HashType ht, const std::string & modulus)
+    : hashSink(ht)
+    , rewritingSink(modulus, std::string(modulus.size(), 0), hashSink)
+{
+}
+
+void HashModuloSink::operator () (std::string_view data)
+{
+    rewritingSink(data);
+}
+
+HashResult HashModuloSink::finish()
+{
+    rewritingSink.flush();
+
+    /* Hash the positions of the self-references. This ensures that a
+       NAR with self-references and a NAR with some of the
+       self-references already zeroed out do not produce a hash
+       collision. FIXME: proof. */
+    for (auto & pos : rewritingSink.matches)
+        hashSink(fmt("|%d", pos));
+
+    auto h = hashSink.finish();
+    return {h.first, rewritingSink.pos};
+}
+
+}
diff --git a/src/libutil/references.hh b/src/libutil/references.hh
new file mode 100644
index 000000000..ffd730e7b
--- /dev/null
+++ b/src/libutil/references.hh
@@ -0,0 +1,56 @@
+#pragma once
+///@file
+
+#include "hash.hh"
+
+namespace nix {
+
+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
+{
+    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;
+
+    void flush();
+};
+
+struct HashModuloSink : AbstractHashSink
+{
+    HashSink hashSink;
+    RewritingSink rewritingSink;
+
+    HashModuloSink(HashType ht, const std::string & modulus);
+
+    void operator () (std::string_view data) override;
+
+    HashResult finish() 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", {}}
+    )
+);
+
+}
+