From d19470540a0769313d219295f245975af5589edb Mon Sep 17 00:00:00 2001
From: Snehasish Kumar <snehasishk@google.com>
Date: Thu, 7 Oct 2021 16:35:11 -0700
Subject: [PATCH] [sanitizer] Add a ForEach callback interface for AddrHashMap.

This change adds a ForEach method to the AddrHashMap class which can
then be used to iterate over all the key value pairs in the hash map.
I intend to use this in an upcoming change to the memprof runtime.

Added a unit test to cover basic insertion and the ForEach callback.

Differential Revision: https://reviews.llvm.org/D111368
---
 .../lib/sanitizer_common/sanitizer_addrhashmap.h   | 38 +++++++++++++
 .../lib/sanitizer_common/tests/CMakeLists.txt      |  1 +
 .../tests/sanitizer_addrhashmap_test.cpp           | 62 ++++++++++++++++++++++
 3 files changed, 101 insertions(+)
 create mode 100644 compiler-rt/lib/sanitizer_common/tests/sanitizer_addrhashmap_test.cpp

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h b/compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h
index 73b48cb..7e2fa91 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h
@@ -39,6 +39,11 @@ namespace __sanitizer {
 //   the current thread has exclusive access to the data
 //   if !h.exists() then the element never existed
 // }
+// {
+//   Map::Handle h(&m, addr, false, true);
+//   this will create a new element or return a handle to an existing element
+//   if !h.created() this thread does *not* have exclusive access to the data
+// }
 template<typename T, uptr kSize>
 class AddrHashMap {
  private:
@@ -89,6 +94,12 @@ class AddrHashMap {
     bool                   create_;
   };
 
+  typedef void (*ForEachCallback)(const uptr key, const T &val, void *arg);
+  // ForEach acquires a lock on each bucket while iterating over
+  // elements. Note that this only ensures that the structure of the hashmap is
+  // unchanged, there may be a data race to the element itself.
+  void ForEach(ForEachCallback cb, void *arg);
+
  private:
   friend class Handle;
   Bucket *table_;
@@ -98,6 +109,33 @@ class AddrHashMap {
   uptr calcHash(uptr addr);
 };
 
+template <typename T, uptr kSize>
+void AddrHashMap<T, kSize>::ForEach(ForEachCallback cb, void *arg) {
+  for (uptr n = 0; n < kSize; n++) {
+    Bucket *bucket = &table_[n];
+
+    ReadLock lock(&bucket->mtx);
+
+    for (uptr i = 0; i < kBucketSize; i++) {
+      Cell *c = &bucket->cells[i];
+      uptr addr1 = atomic_load(&c->addr, memory_order_acquire);
+      if (addr1 != 0)
+        cb(addr1, c->val, arg);
+    }
+
+    // Iterate over any additional cells.
+    if (AddBucket *add =
+            (AddBucket *)atomic_load(&bucket->add, memory_order_acquire)) {
+      for (uptr i = 0; i < add->size; i++) {
+        Cell *c = &add->cells[i];
+        uptr addr1 = atomic_load(&c->addr, memory_order_acquire);
+        if (addr1 != 0)
+          cb(addr1, c->val, arg);
+      }
+    }
+  }
+}
+
 template<typename T, uptr kSize>
 AddrHashMap<T, kSize>::Handle::Handle(AddrHashMap<T, kSize> *map, uptr addr) {
   map_ = map;
diff --git a/compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt b/compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt
index e982992..d689d60 100644
--- a/compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt
+++ b/compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt
@@ -9,6 +9,7 @@ if(APPLE)
 endif()
 
 set(SANITIZER_UNITTESTS
+  sanitizer_addrhashmap_test.cpp
   sanitizer_allocator_test.cpp
   sanitizer_atomic_test.cpp
   sanitizer_bitvector_test.cpp
diff --git a/compiler-rt/lib/sanitizer_common/tests/sanitizer_addrhashmap_test.cpp b/compiler-rt/lib/sanitizer_common/tests/sanitizer_addrhashmap_test.cpp
new file mode 100644
index 0000000..8c7574e
--- /dev/null
+++ b/compiler-rt/lib/sanitizer_common/tests/sanitizer_addrhashmap_test.cpp
@@ -0,0 +1,62 @@
+//===-- sanitizer_addrhashmap_test.cpp ------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#include "sanitizer_common/sanitizer_addrhashmap.h"
+
+#include <unordered_map>
+
+#include "gtest/gtest.h"
+
+namespace __sanitizer {
+
+struct Value {
+  int payload;
+  inline bool operator==(const Value& rhs) const {
+    return payload == rhs.payload;
+  }
+};
+
+using MapTy = AddrHashMap<Value, 11>;
+using HandleTy = MapTy::Handle;
+using RefMapTy = std::unordered_map<uptr, Value>;
+
+static void ExistsInReferenceMap(const uptr key, const Value& val, void* arg) {
+  RefMapTy* ref = reinterpret_cast<RefMapTy*>(arg);
+  const RefMapTy::iterator iter = ref->find(key);
+  ASSERT_NE(iter, ref->end());
+  EXPECT_EQ(iter->second, val);
+  ref->erase(iter);
+}
+
+TEST(AddrHashMap, Basic) {
+  // Use a reference implementation to compare with.
+  RefMapTy reference_map{
+      {0x1000, {1}},
+      {0x2000, {2}},
+      {0x3000, {3}},
+  };
+
+  MapTy m;
+
+  for (const auto& key_val : reference_map) {
+    const uptr key = key_val.first;
+    const Value val = key_val.second;
+
+    // Insert all the elements.
+    {
+      HandleTy h(&m, key);
+      ASSERT_TRUE(h.created());
+      h->payload = val.payload;
+    }
+  }
+
+  // Now check that all the elements are present.
+  m.ForEach(ExistsInReferenceMap, &reference_map);
+  EXPECT_TRUE(reference_map.empty());
+}
+
+}  // namespace __sanitizer
-- 
2.7.4