Fix insert/probe for oa tables

Author: pmenon

src/include/codegen/util/oa_hash_table.h

@@ -106,16 +106,26 @@ class OAHashTable {
   static void Destroy(OAHashTable &table);
 
   /**
-   * Insert a key-value pair into the hash-table.
+   * Insert a key-value pair into the hash-table. Mostly used for testing.
    *
-   * @tparam Key The datatype of the key
-   * @tparam Value The datatype of the value
    * @param hash The hash value of the key
    * @param key The key to store in the table
    * @param value The value to store in the value
    */
   template <typename Key, typename Value>
-  void Insert(uint64_t hash, Key &key, Value &value);
+  void Insert(uint64_t hash, const Key &key, const Value &value);
+
+  /**
+   * Probe a key in the hash table. Doesn't deal with duplicate values. Mostly
+   * for testing.
+   *
+   * @param hash
+   * @param key
+   * @param value
+   * @return
+   */
+  template <typename Key, typename Value>
+  bool Probe(uint64_t hash, const Key &key, Value &value);
 
   /**
    * Make room in the hash-table to store a new key-value pair. The provided
@@ -263,13 +273,13 @@ class OAHashTable {
 };
 
 template <typename Key, typename Value>
-void OAHashTable::Insert(uint64_t hash, Key &key, Value &value) {
+void OAHashTable::Insert(uint64_t hash, const Key &key, const Value &value) {
   uint64_t bucket = hash & bucket_mask_;
 
   uint64_t entry_int =
       reinterpret_cast<uint64_t>(buckets_) + bucket * entry_size_;
   while (true) {
-    HashEntry *entry = reinterpret_cast<HashEntry *>(entry_int);
+    auto *entry = reinterpret_cast<HashEntry *>(entry_int);
 
     // If entry is free, dump key and value
     if (entry->IsFree()) {
@@ -292,10 +302,49 @@ void OAHashTable::Insert(uint64_t hash, Key &key, Value &value) {
     }
 
     // Continue
-    bucket = (bucket == num_buckets_) ? 0 : bucket + 1;
-    entry_int = (bucket == num_buckets_) ? reinterpret_cast<uint64_t>(buckets_)
-                                         : entry_int + entry_size_;
+    bucket++;
+    entry_int += entry_size_;
+
+    // Wrap
+    if (bucket == num_buckets_) {
+      bucket = 0;
+      entry_int = reinterpret_cast<uint64_t>(buckets_);
+    }
+  }
+}
+
+template <typename Key, typename Value>
+bool OAHashTable::Probe(uint64_t hash, const Key &key, Value &value) {
+  uint64_t steps = 0;
+
+  uint64_t bucket = hash & bucket_mask_;
+
+  uint64_t entry_int =
+      reinterpret_cast<uint64_t>(buckets_) + (bucket * entry_size_);
+
+  while (steps++ < num_entries_) {
+    auto *entry = reinterpret_cast<HashEntry *>(entry_int);
+    // check if entry is free
+    if (!entry->IsFree() && entry->hash == hash) {
+      // hashes match, check key
+      auto *entry_key = reinterpret_cast<Key *>(entry->data);
+      if (*entry_key == key) {
+        value = *reinterpret_cast<Value *>(entry->data + sizeof(Key));
+        return true;
+      }
+    }
+
+    // Continue
+    bucket++;
+    entry_int += entry_size_;
+
+    // Wrap
+    if (bucket == num_buckets_) {
+      bucket = 0;
+      entry_int = reinterpret_cast<uint64_t>(buckets_);
+    }
   }
+  return false;
 }
 
 }  // namespace util

test/codegen/oa_hash_table_test.cpp

@@ -136,15 +136,19 @@ TEST_F(OAHashTableTest, MicroBenchmark) {
   std::vector<Key> keys;
   Value v = {6, 5, 4, 3};
 
+  std::random_device r;
+  std::default_random_engine e(r());
+  std::uniform_int_distribution<uint32_t> gen;
+
   // Create all keys
   uint32_t num_keys = 100000;
   for (uint32_t i = 0; i < num_keys; i++) {
-    keys.emplace_back(1, static_cast<uint32_t>(rand()));
+    keys.emplace_back(gen(e), gen(e));
   }
 
-  double avg_oaht = 0.0;
-  double avg_map = 0.0;
-  double avg_cuckoo = 0.0;
+  double avg_oaht_insert = 0.0, avg_oaht_probe = 0.0;
+  double avg_map_insert = 0.0, avg_map_probe = 0.0;
+  double avg_cuckoo_insert = 0.0, avg_cuckoo_probe = 0.0;
 
   // First, bench ours ...
   {
@@ -156,14 +160,29 @@ TEST_F(OAHashTableTest, MicroBenchmark) {
       Timer<std::ratio<1, 1000>> timer;
       timer.Start();
 
-      // Start
+      // Start Insert
       for (uint32_t i = 0; i < num_keys; i++) {
         ht.Insert(Hash(keys[i]), keys[i], v);
       }
-      // End
+      // End Insert
 
       timer.Stop();
-      avg_oaht += timer.GetDuration();
+      avg_oaht_insert += timer.GetDuration();
+
+      timer.Reset();
+      timer.Start();
+
+      // Start Probe
+      std::vector<Key> shuffled = keys;
+      std::random_shuffle(shuffled.begin(), shuffled.end());
+      for (uint32_t i = 0; i < num_keys; i++) {
+        Value probe_val;
+        EXPECT_TRUE(ht.Probe(Hash(shuffled[i]), shuffled[i], probe_val));
+      }
+      // End Probe
+
+      timer.Stop();
+      avg_oaht_probe += timer.GetDuration();
     }
   }
 
@@ -185,7 +204,21 @@ TEST_F(OAHashTableTest, MicroBenchmark) {
       }
 
       timer.Stop();
-      avg_map += timer.GetDuration();
+      avg_map_insert += timer.GetDuration();
+
+      timer.Reset();
+      timer.Start();
+
+      // Start Probe
+      std::vector<Key> shuffled = keys;
+      std::random_shuffle(shuffled.begin(), shuffled.end());
+      for (uint32_t i = 0; i < num_keys; i++) {
+        EXPECT_NE(ht.find(shuffled[i]), ht.end());
+      }
+      // End Probe
+
+      timer.Stop();
+      avg_map_probe += timer.GetDuration();
     }
   }
 
@@ -207,13 +240,33 @@ TEST_F(OAHashTableTest, MicroBenchmark) {
       }
 
       timer.Stop();
-      avg_cuckoo += timer.GetDuration();
+      avg_cuckoo_insert += timer.GetDuration();
+
+      timer.Reset();
+      timer.Start();
+
+      // Start Probe
+      std::vector<Key> shuffled = keys;
+      std::random_shuffle(shuffled.begin(), shuffled.end());
+      for (uint32_t i = 0; i < num_keys; i++) {
+        Value probe_val;
+        EXPECT_TRUE(map.find(shuffled[i], probe_val));
+      }
+      // End Probe
+
+      timer.Stop();
+      avg_cuckoo_probe += timer.GetDuration();
     }
   }
 
-  LOG_INFO("Avg OA_HT: %.2lf, Avg std::unordered_map: %.2lf, Avg cuckoo: %.2lf",
-           avg_oaht / (double)num_runs, avg_map / (double)num_runs,
-           avg_cuckoo / (double)num_runs);
+  LOG_INFO("OA_HT insert: %.2lf, probe: %.2lf",
+           avg_oaht_insert / (double)num_runs,
+           avg_oaht_probe / (double)num_runs);
+  LOG_INFO("std::unordered_map insert: %.2lf, probe: %.2lf",
+           avg_map_insert / (double)num_runs, avg_map_probe / (double)num_runs);
+  LOG_INFO("Cuckoo insert: %.2lf, probe: %.2lf",
+           avg_cuckoo_insert / (double)num_runs,
+           avg_cuckoo_probe / (double)num_runs);
 }
 
 }  // namespace test