3 files changed, 42 insertions, 0 deletions
diff --git a/tdutils/td/utils/WaitFreeHashMap.h b/tdutils/td/utils/WaitFreeHashMap.h
index 592ac96d9..956a1a865 100644
--- a/tdutils/td/utils/WaitFreeHashMap.h
+++ b/tdutils/td/utils/WaitFreeHashMap.h
@@ -161,6 +161,21 @@ class WaitFreeHashMap {
     }
   }
 
+  template <class F>
+  bool remove_if(const F &f) {
+    if (wait_free_storage_ == nullptr) {
+      return default_map_.remove_if(f);
+    }
+
+    bool is_removed = false;
+    for (auto &it : wait_free_storage_->maps_) {
+      if (it.remove_if(f)) {
+        is_removed = true;
+      }
+    }
+    return is_removed;
+  }
+
   size_t calc_size() const {
     if (wait_free_storage_ == nullptr) {
       return default_map_.size();
diff --git a/tdutils/td/utils/algorithm.h b/tdutils/td/utils/algorithm.h
index a87de8ecb..6107b1d29 100644
--- a/tdutils/td/utils/algorithm.h
+++ b/tdutils/td/utils/algorithm.h
@@ -301,4 +301,12 @@ bool table_remove_if(FlatHashTable<NodeT, HashT, EqT> &table, FuncT &&func) {
   return table.remove_if(func);
 }
 
+template <class KeyT, class ValueT, class HashT, class EqT>
+class WaitFreeHashMap;
+
+template <class KeyT, class ValueT, class HashT, class EqT, class FuncT>
+bool table_remove_if(WaitFreeHashMap<KeyT, ValueT, HashT, EqT> &table, const FuncT &func) {
+  return table.remove_if(func);
+}
+
 }  // namespace td
diff --git a/tdutils/test/WaitFreeHashMap.cpp b/tdutils/test/WaitFreeHashMap.cpp
index ee7e57df6..c73cfda47 100644
--- a/tdutils/test/WaitFreeHashMap.cpp
+++ b/tdutils/test/WaitFreeHashMap.cpp
@@ -93,3 +93,22 @@ TEST(WaitFreeHashMap, stress_test) {
     }
   }
 }
+
+TEST(WaitFreeHashMap, remove_if) {
+  td::WaitFreeHashMap<td::uint64, td::uint64> map;
+  bool is_removed;
+  for (td::uint64 i = 0; i < 10000; i++) {
+    map[2 * i + 1] = 0;
+    map[2 * i + 2] = 1;
+    is_removed = map.remove_if([i](const auto &num) { return num.first <= i; });
+    CHECK(!is_removed);
+    CHECK(map.calc_size() == i + 2);
+    is_removed = map.remove_if([i](const auto &num) { return num.first <= i + 1; });
+    CHECK(is_removed);
+    CHECK(map.calc_size() == i + 1);
+  }
+  is_removed = map.remove_if([](const auto &num) { return num.first <= 19999; });
+  CHECK(map.calc_size() == 1);
+  is_removed = map.remove_if([](const auto &num) { return num.first <= 20000; });
+  CHECK(map.calc_size() == 0);
+}