From 867b291cad2a9115fe129dd338974397d048d5a6 Mon Sep 17 00:00:00 2001
From: Tudor Bosman <tudorb@fb.com>
Date: Mon, 2 Jun 2014 12:08:05 -0700
Subject: [PATCH] Hasher and equality comparison for IOBuf

Test Plan: test added

Reviewed By: davejwatson@fb.com

Subscribers: folly@lists

FB internal diff: D1359469
---
 folly/io/IOBuf.cpp          | 41 +++++++++++++++++++++++++
 folly/io/IOBuf.h            | 27 ++++++++++++++++
 folly/io/test/IOBufTest.cpp | 61 +++++++++++++++++++++++++++++++++++++
 3 files changed, 129 insertions(+)

diff --git a/folly/io/IOBuf.cpp b/folly/io/IOBuf.cpp
index 28a41be5..b369eedf 100644
--- a/folly/io/IOBuf.cpp
+++ b/folly/io/IOBuf.cpp
@@ -23,6 +23,8 @@
 #include "folly/Malloc.h"
 #include "folly/Memory.h"
 #include "folly/ScopeGuard.h"
+#include "folly/SpookyHashV2.h"
+#include "folly/io/Cursor.h"
 
 #include <stdexcept>
 #include <assert.h>
@@ -893,4 +895,43 @@ folly::fbvector<struct iovec> IOBuf::getIov() const {
   return iov;
 }
 
+size_t IOBufHash::operator()(const IOBuf& buf) const {
+  folly::hash::SpookyHashV2 hasher;
+  hasher.Init(0, 0);
+  io::Cursor cursor(&buf);
+  for (;;) {
+    auto p = cursor.peek();
+    if (p.second == 0) {
+      break;
+    }
+    hasher.Update(p.first, p.second);
+    cursor.skip(p.second);
+  }
+  uint64_t h1;
+  uint64_t h2;
+  hasher.Final(&h1, &h2);
+  return h1;
+}
+
+bool IOBufEqual::operator()(const IOBuf& a, const IOBuf& b) const {
+  io::Cursor ca(&a);
+  io::Cursor cb(&b);
+  for (;;) {
+    auto pa = ca.peek();
+    auto pb = cb.peek();
+    if (pa.second == 0 && pb.second == 0) {
+      return true;
+    } else if (pa.second == 0 || pb.second == 0) {
+      return false;
+    }
+    size_t n = std::min(pa.second, pb.second);
+    DCHECK_GT(n, 0);
+    if (memcmp(pa.first, pb.first, n)) {
+      return false;
+    }
+    ca.skip(n);
+    cb.skip(n);
+  }
+}
+
 } // folly
diff --git a/folly/io/IOBuf.h b/folly/io/IOBuf.h
index 40229343..2895c0d5 100644
--- a/folly/io/IOBuf.h
+++ b/folly/io/IOBuf.h
@@ -1265,6 +1265,33 @@ class IOBuf {
   }
 };
 
+/**
+ * Hasher for IOBuf objects. Hashes the entire chain using SpookyHashV2.
+ */
+struct IOBufHash {
+  size_t operator()(const IOBuf& buf) const;
+  size_t operator()(const std::unique_ptr<IOBuf>& buf) const {
+    return buf ? (*this)(*buf) : 0;
+  }
+};
+
+/**
+ * Equality predicate for IOBuf objects. Compares data in the entire chain.
+ */
+struct IOBufEqual {
+  bool operator()(const IOBuf& a, const IOBuf& b) const;
+  bool operator()(const std::unique_ptr<IOBuf>& a,
+                  const std::unique_ptr<IOBuf>& b) const {
+    if (!a && !b) {
+      return true;
+    } else if (!a || !b) {
+      return false;
+    } else {
+      return (*this)(*a, *b);
+    }
+  }
+};
+
 template <class UniquePtr>
 typename std::enable_if<detail::IsUniquePtrToSL<UniquePtr>::value,
                         std::unique_ptr<IOBuf>>::type
diff --git a/folly/io/test/IOBufTest.cpp b/folly/io/test/IOBufTest.cpp
index 868289be..0a1a17c6 100644
--- a/folly/io/test/IOBufTest.cpp
+++ b/folly/io/test/IOBufTest.cpp
@@ -1006,6 +1006,67 @@ TEST(IOBuf, move) {
   EXPECT_FALSE(outerBuf.isShared());
 }
 
+namespace {
+std::unique_ptr<IOBuf> fromStr(StringPiece sp) {
+  return IOBuf::copyBuffer(ByteRange(sp));
+}
+}  // namespace
+
+TEST(IOBuf, HashAndEqual) {
+  folly::IOBufEqual eq;
+  folly::IOBufHash hash;
+
+  EXPECT_TRUE(eq(nullptr, nullptr));
+  EXPECT_EQ(0, hash(nullptr));
+
+  auto empty = IOBuf::create(0);
+
+  EXPECT_TRUE(eq(*empty, *empty));
+  EXPECT_TRUE(eq(empty, empty));
+
+  EXPECT_FALSE(eq(nullptr, empty));
+  EXPECT_FALSE(eq(empty, nullptr));
+
+  EXPECT_EQ(hash(*empty), hash(empty));
+  EXPECT_NE(0, hash(empty));
+
+  auto a = fromStr("hello");
+
+  EXPECT_TRUE(eq(*a, *a));
+  EXPECT_TRUE(eq(a, a));
+
+  EXPECT_FALSE(eq(nullptr, a));
+  EXPECT_FALSE(eq(a, nullptr));
+
+  EXPECT_EQ(hash(*a), hash(a));
+  EXPECT_NE(0, hash(a));
+
+  auto b = fromStr("hello");
+
+  EXPECT_TRUE(eq(*a, *b));
+  EXPECT_TRUE(eq(a, b));
+
+  EXPECT_EQ(hash(a), hash(b));
+
+  auto c = fromStr("hellow");
+
+  EXPECT_FALSE(eq(a, c));
+  EXPECT_NE(hash(a), hash(c));
+
+  auto d = fromStr("world");
+
+  EXPECT_FALSE(eq(a, d));
+  EXPECT_NE(hash(a), hash(d));
+
+  auto e = fromStr("helloworld");
+  auto f = fromStr("hello");
+  f->prependChain(fromStr("wo"));
+  f->prependChain(fromStr("rld"));
+
+  EXPECT_TRUE(eq(e, f));
+  EXPECT_EQ(hash(e), hash(f));
+}
+
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   google::ParseCommandLineFlags(&argc, &argv, true);
-- 
2.34.1