diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h b/compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h --- a/compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h +++ b/compiler-rt/lib/sanitizer_common/sanitizer_addrhashmap.h @@ -89,6 +89,12 @@ bool create_; }; + typedef void (*ForEachCallback)(const uptr key, const T &val, void *arg); + // ForEach acquires an exclusive lock on each bucket while iterating over + // elements. Callers should use additional locking at the callsite for a + // consistent view of map. + void ForEach(ForEachCallback cb, void *arg); + private: friend class Handle; Bucket *table_; @@ -98,6 +104,33 @@ uptr calcHash(uptr addr); }; +template +void AddrHashMap::ForEach(ForEachCallback cb, void *arg) { + for (uptr n = 0; n < kSize; n++) { + Bucket *bucket = &table_[n]; + + Lock lock(&bucket->mtx); + + for (uptr i = 0; i < kBucketSize; i++) { + Cell *c = &bucket->cells[i]; + uptr addr1 = atomic_load(&c->addr, memory_order_relaxed); + if (addr1 != 0) + cb(addr1, c->val, arg); + } + + // Iterate over any additional cells. + if (AddBucket *add = + (AddBucket *)atomic_load(&bucket->add, memory_order_relaxed)) { + for (uptr i = 0; i < add->size; i++) { + Cell *c = &add->cells[i]; + uptr addr1 = atomic_load(&c->addr, memory_order_relaxed); + if (addr1 != 0) + cb(addr1, c->val, arg); + } + } + } +} + template AddrHashMap::Handle::Handle(AddrHashMap *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 --- a/compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt +++ b/compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt @@ -9,6 +9,7 @@ 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 --- /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 + +#include "gtest/gtest.h" + +namespace __sanitizer { + +struct Value { + int payload; + inline bool operator==(const Value& rhs) const { + return payload == rhs.payload; + } +}; + +using MapTy = AddrHashMap; +using HandleTy = MapTy::Handle; +using RefMapTy = std::unordered_map; + +static void ExistsInReferenceMap(const uptr key, const Value& val, void* arg) { + RefMapTy* ref = reinterpret_cast(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