[compiler-rt] 09256fe - [sanitizer] Add DenseMap::forEach

[Vitaly Buka via llvm-commits]

Author: Vitaly Buka
Date: 2021-11-23T15:23:25-08:00
New Revision: 09256fe980ddc46be36ab4460ae1850aa46f094e

URL: https://github.com/llvm/llvm-project/commit/09256fe980ddc46be36ab4460ae1850aa46f094e
DIFF: https://github.com/llvm/llvm-project/commit/09256fe980ddc46be36ab4460ae1850aa46f094e.diff

LOG: [sanitizer] Add DenseMap::forEach

Added: 
    

Modified: 
    compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h
    compiler-rt/lib/sanitizer_common/tests/sanitizer_dense_map_test.cpp

Removed: 
    


################################################################################
diff  --git a/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h b/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h
index 3fa6af76ce29c..6ddd8b1edc8a3 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h
@@ -250,6 +250,22 @@ class DenseMapBase {
     return true;
   }
 
+  /// Iterate over active entries of the container.
+  ///
+  /// Function can return fast to stop the process.
+  template <class Fn>
+  void forEach(Fn fn) {
+    const KeyT EmptyKey = getEmptyKey(), TombstoneKey = getTombstoneKey();
+    for (auto *P = getBuckets(), *E = getBucketsEnd(); P != E; ++P) {
+      const KeyT K = P->getFirst();
+      if (!KeyInfoT::isEqual(K, EmptyKey) &&
+          !KeyInfoT::isEqual(K, TombstoneKey)) {
+        if (!fn(*P))
+          return;
+      }
+    }
+  }
+
  protected:
   DenseMapBase() = default;
 

diff  --git a/compiler-rt/lib/sanitizer_common/tests/sanitizer_dense_map_test.cpp b/compiler-rt/lib/sanitizer_common/tests/sanitizer_dense_map_test.cpp
index c4cf746f35bf7..fcc905ea87b8d 100644
--- a/compiler-rt/lib/sanitizer_common/tests/sanitizer_dense_map_test.cpp
+++ b/compiler-rt/lib/sanitizer_common/tests/sanitizer_dense_map_test.cpp
@@ -299,6 +299,29 @@ TYPED_TEST(DenseMapTest, SwapTest) {
     EXPECT_EQ(this->getValue(i), this->Map[this->getKey(i)]);
 }
 
+// A more complex iteration test
+TYPED_TEST(DenseMapTest, IterationTest) {
+  int visited[100];
+  std::map<typename TypeParam::key_type, unsigned> visitedIndex;
+
+  // Insert 100 numbers into the map
+  for (int i = 0; i < 100; ++i) {
+    visited[i] = 0;
+    visitedIndex[this->getKey(i)] = i;
+
+    this->Map[this->getKey(i)] = this->getValue(i);
+  }
+
+  // Iterate over all numbers and mark each one found.
+  this->Map.forEach([&](const typename TypeParam::value_type &kv) {
+    ++visited[visitedIndex[kv.first]];
+    return true;
+  });
+
+  // Ensure every number was visited.
+  for (int i = 0; i < 100; ++i) ASSERT_EQ(1, visited[i]);
+}
+
 namespace {
 // Simple class that counts how many moves and copy happens when growing a map
 struct CountCopyAndMove {