From ac8af37872a4551a3960bfc4a48d1d6e351e0bff Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Wed, 20 Jul 2022 13:35:32 +0300
Subject: [PATCH] Add size() and empty() to WaitFreeHashMap.

---
 tdutils/td/utils/WaitFreeHashMap.h | 25 +++++++++++++++++++++++++
 tdutils/test/WaitFreeHashMap.cpp   |  6 ++++++
 2 files changed, 31 insertions(+)

diff --git a/tdutils/td/utils/WaitFreeHashMap.h b/tdutils/td/utils/WaitFreeHashMap.h
index 0d32e0bd5..14f6af8d3 100644
--- a/tdutils/td/utils/WaitFreeHashMap.h
+++ b/tdutils/td/utils/WaitFreeHashMap.h
@@ -61,6 +61,31 @@ class WaitFreeHashMap {
   size_t erase(const KeyT &key) {
     return get_storage(key).erase(key);
   }
+
+  size_t size() const {
+    if (wait_free_storage_ == nullptr) {
+      return default_map_.size();
+    }
+
+    size_t result = 0;
+    for (size_t i = 0; i < MAX_STORAGE_COUNT; i++) {
+      result += wait_free_storage_->maps_[i].size();
+    }
+    return result;
+  }
+
+  bool empty() const {
+    if (wait_free_storage_ == nullptr) {
+      return default_map_.empty();
+    }
+
+    for (size_t i = 0; i < MAX_STORAGE_COUNT; i++) {
+      if (!wait_free_storage_->maps_[i].empty()) {
+        return false;
+      }
+    }
+    return true;
+  }
 };
 
 }  // namespace td
diff --git a/tdutils/test/WaitFreeHashMap.cpp b/tdutils/test/WaitFreeHashMap.cpp
index 8865d5ffd..8c1b0c1db 100644
--- a/tdutils/test/WaitFreeHashMap.cpp
+++ b/tdutils/test/WaitFreeHashMap.cpp
@@ -30,6 +30,8 @@ TEST(WaitFreeHashMap, stress_test) {
     reference[key] = value;
     map.set(key, value);
     ASSERT_EQ(reference[key], map.get(key));
+    ASSERT_EQ(reference.size(), map.size());
+    ASSERT_EQ(reference.empty(), map.empty());
   });
 
   add_step(2000, [&] {
@@ -37,6 +39,8 @@ TEST(WaitFreeHashMap, stress_test) {
     auto ref_it = reference.find(key);
     auto ref_value = ref_it == reference.end() ? 0 : ref_it->second;
     ASSERT_EQ(ref_value, map.get(key));
+    ASSERT_EQ(reference.size(), map.size());
+    ASSERT_EQ(reference.empty(), map.empty());
   });
 
   add_step(500, [&] {
@@ -44,6 +48,8 @@ TEST(WaitFreeHashMap, stress_test) {
     size_t reference_erased_count = reference.erase(key);
     size_t map_erased_count = map.erase(key);
     ASSERT_EQ(reference_erased_count, map_erased_count);
+    ASSERT_EQ(reference.size(), map.size());
+    ASSERT_EQ(reference.empty(), map.empty());
   });
 
   td::RandomSteps runner(std::move(steps));