From d7ecb21109af66fd2af839e4a810cad61b933c03 Mon Sep 17 00:00:00 2001
From: BlackMATov <blackmatov@gmail.com>
Date: Mon, 3 Jun 2019 00:12:40 +0700
Subject: [PATCH] add try_emplace and insert_or_assign functions #27

---
 README.md                     | 10 +++++++++
 headers/flat.hpp/flat_map.hpp | 42 +++++++++++++++++++++++++++++++++++
 untests/flat_map_tests.cpp    | 28 +++++++++++++++++++++++
 3 files changed, 80 insertions(+)

diff --git a/README.md b/README.md
index 58b3f86..eb34cd0 100644
--- a/README.md
+++ b/README.md
@@ -477,6 +477,11 @@ std::pair<iterator, bool> insert(const value_type& value);
 iterator insert(const_iterator hint, value_type&& value);
 iterator insert(const_iterator hint, const value_type& value);
 
+template < typename TT >
+std::pair<iterator, bool> insert_or_assign(key_type&& key, TT&& value);
+template < typename TT >
+std::pair<iterator, bool> insert_or_assign(const key_type& key, TT&& value);
+
 template < typename InputIter >
 void insert(InputIter first, InputIter last);
 template < typename InputIter >
@@ -490,6 +495,11 @@ std::pair<iterator, bool> emplace(Args&&... args);
 template < typename... Args >
 iterator emplace_hint(const_iterator hint, Args&&... args);
 
+template < typename... Args >
+std::pair<iterator, bool> try_emplace(key_type&& key, Args&&... args);
+template < typename... Args >
+std::pair<iterator, bool> try_emplace(const key_type& key, Args&&... args);
+
 void clear();
 iterator erase(const_iterator iter);
 iterator erase(const_iterator first, const_iterator last);
diff --git a/headers/flat.hpp/flat_map.hpp b/headers/flat.hpp/flat_map.hpp
index ac2e635..e32ae44 100644
--- a/headers/flat.hpp/flat_map.hpp
+++ b/headers/flat.hpp/flat_map.hpp
@@ -399,6 +399,28 @@ namespace flat_hpp
                 : insert(value).first;
         }
 
+        template < typename TT >
+        std::pair<iterator, bool> insert_or_assign(key_type&& key, TT&& value) {
+            iterator iter = lower_bound(key);
+            if ( iter == end() || this->operator()(key, *iter) ) {
+                iter = emplace_hint(iter, std::move(key), std::forward<TT>(value));
+                return {iter, true};
+            }
+            (*iter).second = std::forward<TT>(value);
+            return {iter, false};
+        }
+
+        template < typename TT >
+        std::pair<iterator, bool> insert_or_assign(const key_type& key, TT&& value) {
+            iterator iter = lower_bound(key);
+            if ( iter == end() || this->operator()(key, *iter) ) {
+                iter = emplace_hint(iter, key, std::forward<TT>(value));
+                return {iter, true};
+            }
+            (*iter).second = std::forward<TT>(value);
+            return {iter, false};
+        }
+
         template < typename InputIter >
         void insert(InputIter first, InputIter last) {
             insert_range_(first, last);
@@ -427,6 +449,26 @@ namespace flat_hpp
             return insert(hint, value_type(std::forward<Args>(args)...));
         }
 
+        template < typename... Args >
+        std::pair<iterator, bool> try_emplace(key_type&& key, Args&&... args) {
+            iterator iter = lower_bound(key);
+            if ( iter == end() || this->operator()(key, *iter) ) {
+                iter = emplace_hint(iter, std::move(key), std::forward<Args>(args)...);
+                return {iter, true};
+            }
+            return {iter, false};
+        }
+
+        template < typename... Args >
+        std::pair<iterator, bool> try_emplace(const key_type& key, Args&&... args) {
+            iterator iter = lower_bound(key);
+            if ( iter == end() || this->operator()(key, *iter) ) {
+                iter = emplace_hint(iter, key, std::forward<Args>(args)...);
+                return {iter, true};
+            }
+            return {iter, false};
+        }
+
         void clear()
         noexcept(noexcept(std::declval<container_type&>().clear())) {
             data_.clear();
diff --git a/untests/flat_map_tests.cpp b/untests/flat_map_tests.cpp
index e5ad3a0..5f0c406 100644
--- a/untests/flat_map_tests.cpp
+++ b/untests/flat_map_tests.cpp
@@ -369,6 +369,34 @@ TEST_CASE("flat_map") {
             s1.insert(sorted_range, {{1,3},{2,2},{2,2},{3,1}});
             REQUIRE(s1 == map_t{{1,3},{2,2},{3,1}});
         }
+
+        {
+            map_t s0;
+            auto i0 = s0.insert_or_assign(1, 4);
+            REQUIRE(i0 == std::make_pair(s0.begin(), true));
+            REQUIRE(s0 == map_t{{1,4}});
+            auto i1 = s0.insert_or_assign(1, 8);
+            REQUIRE(i1 == std::make_pair(s0.begin(), false));
+            REQUIRE(s0 == map_t{{1,8}});
+
+            const obj_t k0{2};
+            auto i2 = s0.insert_or_assign(k0, 6);
+            REQUIRE(i2 == std::make_pair(s0.begin() + 1, true));
+            REQUIRE(s0 == map_t{{1,8}, {2,6}});
+            auto i3 = s0.insert_or_assign(k0, 2);
+            REQUIRE(i3 == std::make_pair(s0.begin() + 1, false));
+            REQUIRE(s0 == map_t{{1,8}, {2,2}});
+        }
+
+        {
+            map_t s0;
+            auto i0 = s0.try_emplace(1, 4);
+            REQUIRE(i0 == std::make_pair(s0.begin(), true));
+            REQUIRE(s0 == map_t{{1,4}});
+            auto i1 = s0.try_emplace(1, 8);
+            REQUIRE(i1 == std::make_pair(s0.begin(), false));
+            REQUIRE(s0 == map_t{{1,4}});
+        }
     }
     SECTION("erasers") {
         using map_t = flat_map<int, unsigned>;