[util] implement THashMap::insert_or_assign method

IGNIETFERRO-821

ref:1e4e3b22bff7e806800218d58b1061626ebaa10f

2 files changed, 71 insertions, 0 deletions

diff --git a/util/generic/hash.h b/util/generic/hash.h
index e46db21fa9..81311cc1ee 100644
--- a/util/generic/hash.h
+++ b/util/generic/hash.h
@@ -1570,6 +1570,24 @@ public:
         return rep.insert_unique(obj);
     }
 
+    template <class M>
+    std::pair<iterator, bool> insert_or_assign(const Key& k, M&& value) {
+        auto result = try_emplace(k, std::forward<M>(value));
+        if (!result.second) {
+            result.first->second = std::forward<M>(value);
+        }
+        return result;
+    }
+
+    template <class M>
+    std::pair<iterator, bool> insert_or_assign(Key&& k, M&& value) {
+        auto result = try_emplace(std::move(k), std::forward<M>(value));
+        if (!result.second) {
+            result.first->second = std::forward<M>(value);
+        }
+        return result;
+    }
+
     template <typename... Args>
     std::pair<iterator, bool> emplace(Args&&... args) {
         return rep.emplace_unique(std::forward<Args>(args)...);
diff --git a/util/generic/hash_ut.cpp b/util/generic/hash_ut.cpp
index 68dd9d9d78..b6c46f810f 100644
--- a/util/generic/hash_ut.cpp
+++ b/util/generic/hash_ut.cpp
@@ -42,6 +42,7 @@ class THashTest: public TTestBase {
     UNIT_TEST(TestEmplaceDirect);
     UNIT_TEST(TestTryEmplace);
     UNIT_TEST(TestTryEmplaceCopyKey);
+    UNIT_TEST(TestInsertOrAssign);
     UNIT_TEST(TestHMMapEmplace);
     UNIT_TEST(TestHMMapEmplaceNoresize);
     UNIT_TEST(TestHMMapEmplaceDirect);
@@ -94,6 +95,7 @@ protected:
     void TestEmplaceDirect();
     void TestTryEmplace();
     void TestTryEmplaceCopyKey();
+    void TestInsertOrAssign();
     void TestHSetEmplace();
     void TestHSetEmplaceNoresize();
     void TestHSetEmplaceDirect();
@@ -982,6 +984,57 @@ void THashTest::TestTryEmplaceCopyKey() {
     }
 }
 
+void THashTest::TestInsertOrAssign() {
+    static int constructorCounter = 0;
+    static int assignmentCounter = 0;
+
+    struct TCountConstruct {
+        explicit TCountConstruct(int v)
+            : Value(v)
+        {
+            ++constructorCounter;
+        }
+
+        TCountConstruct& operator=(int v) {
+            Value = v;
+            ++assignmentCounter;
+            return *this;
+        }
+
+        TCountConstruct(const TCountConstruct&) = delete;
+        int Value;
+    };
+
+    THashMap<int, TCountConstruct> hash;
+    {
+        auto r = hash.insert_or_assign(TNonCopyableInt<4>(4), 1);
+        UNIT_ASSERT(r.second);
+        UNIT_ASSERT_VALUES_EQUAL(1, hash.size());
+        UNIT_ASSERT_VALUES_EQUAL(1, constructorCounter);
+        UNIT_ASSERT_VALUES_EQUAL(0, assignmentCounter);
+        UNIT_ASSERT_VALUES_EQUAL(1, r.first->second.Value);
+    }
+    {
+        auto r = hash.insert_or_assign(TNonCopyableInt<4>(4), 5);
+        UNIT_ASSERT(!r.second);
+        UNIT_ASSERT_VALUES_EQUAL(1, hash.size());
+        UNIT_ASSERT_VALUES_EQUAL(1, constructorCounter);
+        UNIT_ASSERT_VALUES_EQUAL(1, assignmentCounter);
+        UNIT_ASSERT_VALUES_EQUAL(5, r.first->second.Value);
+    }
+    {
+        constexpr int iterations = 200;
+        for (int iteration = 0; iteration < iterations; ++iteration) {
+            hash.insert_or_assign(iteration, iteration);
+        }
+        UNIT_ASSERT_VALUES_EQUAL(iterations, hash.size());
+        UNIT_ASSERT_VALUES_EQUAL(iterations, constructorCounter);
+        UNIT_ASSERT_VALUES_EQUAL(2, assignmentCounter);
+        UNIT_ASSERT_VALUES_EQUAL(4, hash.at(4).Value);
+        UNIT_ASSERT_VALUES_EQUAL(44, hash.at(44).Value);
+    }
+}
+
 void THashTest::TestHMMapEmplace() {
     using hash_t = THashMultiMap<int, TNonCopyableInt<0>>;
     hash_t hash;