From a617227a368368596b0f6bcdb9acbcd29b95ad00 Mon Sep 17 00:00:00 2001
From: Haobo Xu <haobo@fb.com>
Date: Tue, 19 Nov 2013 23:24:12 -0800
Subject: [PATCH] [RocksDB] fix prefix_test

Summary: user comparator needs to work if either input is prefix only.

Test Plan: ./prefix_test --write_buffer_size=100000 --total_prefixes=10000 --items_per_prefix=10

Reviewers: dhruba, igor

Reviewed By: igor

CC: leveldb

Differential Revision: https://reviews.facebook.net/D14241
---
 db/prefix_test.cc | 31 ++++++++++++++++++++++++-------
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/db/prefix_test.cc b/db/prefix_test.cc
index 71016612a3..6c7fc1697f 100644
--- a/db/prefix_test.cc
+++ b/db/prefix_test.cc
@@ -41,26 +41,43 @@ inline Slice TestKeyToSlice(const TestKey& test_key) {
 }
 
 inline const TestKey* SliceToTestKey(const Slice& slice) {
-  assert(slice.size() == sizeof(TestKey));
   return (const TestKey*)slice.data();
 }
 
 class TestKeyComparator : public Comparator {
  public:
+
+  // Compare needs to be aware of the possibility of a and/or b is
+  // prefix only
   virtual int Compare(const Slice& a, const Slice& b) const {
     const TestKey* key_a = SliceToTestKey(a);
     const TestKey* key_b = SliceToTestKey(b);
     if (key_a->prefix != key_b->prefix) {
       if (key_a->prefix < key_b->prefix) return -1;
       if (key_a->prefix > key_b->prefix) return 1;
-      assert(false);
     } else {
-      if (key_a->sorted < key_b->sorted) return -1;
-      if (key_a->sorted > key_b->sorted) return 1;
-      if (key_a->sorted == key_b->sorted) return 0;
-      assert(false);
+      ASSERT_TRUE(key_a->prefix == key_b->prefix);
+      // note, both a and b could be prefix only
+      if (a.size() != b.size()) {
+        // one of them is prefix
+        ASSERT_TRUE(
+          (a.size() == sizeof(uint64_t) && b.size() == sizeof(TestKey)) ||
+          (b.size() == sizeof(uint64_t) && a.size() == sizeof(TestKey)));
+        if (a.size() < b.size()) return -1;
+        if (a.size() > b.size()) return 1;
+      } else {
+        // both a and b are prefix
+        if (a.size() == sizeof(uint64_t)) {
+          return 0;
+        }
+
+        // both a and b are whole key
+        ASSERT_TRUE(a.size() == sizeof(TestKey) && b.size() == sizeof(TestKey));
+        if (key_a->sorted < key_b->sorted) return -1;
+        if (key_a->sorted > key_b->sorted) return 1;
+        if (key_a->sorted == key_b->sorted) return 0;
+      }
     }
-    assert(false);
     return 0;
   }