Add operator[] to WaitFreeHashMap.

tdutils/td/utils/WaitFreeHashMap.h

@@ -27,16 +27,29 @@ class WaitFreeHashMap {
   };
   unique_ptr<WaitFreeStorage> wait_free_storage_;
 
-  const Storage &get_storage(const KeyT &key) const {
-    if (wait_free_storage_ == nullptr) {
-      return default_map_;
-    }
-
+  Storage &get_wait_free_storage(const KeyT &key) {
     return wait_free_storage_->maps_[randomize_hash(HashT()(key)) & (MAX_STORAGE_COUNT - 1)];
   }
 
   Storage &get_storage(const KeyT &key) {
-    return const_cast<Storage &>(static_cast<const WaitFreeHashMap *>(this)->get_storage(key));
+    if (wait_free_storage_ == nullptr) {
+      return default_map_;
+    }
+
+    return get_wait_free_storage(key);
+  }
+
+  const Storage &get_storage(const KeyT &key) const {
+    return const_cast<WaitFreeHashMap *>(this)->get_storage(key);
+  }
+
+  void split() {
+    CHECK(wait_free_storage_ == nullptr);
+    wait_free_storage_ = make_unique<WaitFreeStorage>();
+    for (auto &it : default_map_) {
+      get_wait_free_storage(it.first).emplace(it.first, std::move(it.second));
+    }
+    default_map_.clear();
   }
 
  public:
@@ -44,12 +57,7 @@ class WaitFreeHashMap {
     auto &storage = get_storage(key);
     storage[key] = std::move(value);
     if (default_map_.size() == MAX_STORAGE_SIZE) {
-      CHECK(wait_free_storage_ == nullptr);
-      wait_free_storage_ = make_unique<WaitFreeStorage>();
-      for (auto &it : default_map_) {
-        get_storage(it.first).emplace(it.first, std::move(it.second));
-      }
-      default_map_.clear();
+      split();
     }
   }
 
@@ -83,6 +91,19 @@ class WaitFreeHashMap {
     return it->second.get();
   }
 
+  ValueT &operator[](const KeyT &key) {
+    if (wait_free_storage_ == nullptr) {
+      ValueT &result = default_map_[key];
+      if (default_map_.size() != MAX_STORAGE_SIZE) {
+        return result;
+      }
+
+      split();
+    }
+
+    return get_wait_free_storage(key)[key];
+  }
+
   size_t erase(const KeyT &key) {
     return get_storage(key).erase(key);
   }

tdutils/test/WaitFreeHashMap.cpp

@@ -28,12 +28,23 @@ TEST(WaitFreeHashMap, stress_test) {
     auto key = gen_key();
     auto value = rnd();
     reference[key] = value;
-    map.set(key, value);
+    if (td::Random::fast_bool()) {
+      map.set(key, value);
+    } else {
+      map[key] = value;
+    }
     ASSERT_EQ(reference[key], map.get(key));
     ASSERT_EQ(reference.size(), map.size());
     ASSERT_EQ(reference.empty(), map.empty());
   });
 
+  add_step(200, [&] {
+    auto key = gen_key();
+    ASSERT_EQ(reference[key], map[key]);
+    ASSERT_EQ(reference.size(), map.size());
+    ASSERT_EQ(reference.empty(), map.empty());
+  });
+
   add_step(2000, [&] {
     auto key = gen_key();
     auto ref_it = reference.find(key);
@@ -56,4 +67,13 @@ TEST(WaitFreeHashMap, stress_test) {
   for (size_t i = 0; i < 1000000; i++) {
     runner.step(rnd);
   }
+
+  for (size_t test = 0; test < 1000; test++) {
+    reference = {};
+    map = {};
+
+    for (size_t i = 0; i < 100; i++) {
+      runner.step(rnd);
+    }
+  }
 }