[base] Implement C++17's InsertOrAssign and TryEmplace

This change implements C++17's std::map::insert_or_assign and
std::map::try_emplace as free standing functions.

Bug: 752720
Change-Id: Ic74997465503b589cdf10f346289ee4f312140f9
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1631650
Commit-Queue: Jan Wilken Dörrie <jdoerrie@chromium.org>
Reviewed-by: Daniel Cheng <dcheng@chromium.org>
Cr-Commit-Position: refs/heads/master@{#665851}

base/stl_util.h

@@ -23,6 +23,7 @@
 
 #include "base/logging.h"
 #include "base/optional.h"
+#include "base/template_util.h"
 
 namespace base {
 
@@ -39,6 +40,11 @@ void IterateAndEraseIf(Container& container, Predicate pred) {
   }
 }
 
+template <typename Iter>
+constexpr bool IsRandomAccessIter =
+    std::is_same<typename std::iterator_traits<Iter>::iterator_category,
+                 std::random_access_iterator_tag>::value;
+
 }  // namespace internal
 
 // C++14 implementation of C++17's std::size():
@@ -134,6 +140,37 @@ STLCount(const Container& container, const T& val) {
   return std::count(container.begin(), container.end(), val);
 }
 
+// O(1) implementation of const casting an iterator for any sequence,
+// associative or unordered associative container in the STL.
+//
+// Reference: https://stackoverflow.com/a/10669041
+template <typename Container,
+          typename ConstIter,
+          std::enable_if_t<!internal::IsRandomAccessIter<ConstIter>>* = nullptr>
+constexpr auto ConstCastIterator(Container& c, ConstIter it) {
+  return c.erase(it, it);
+}
+
+// Explicit overload for std::forward_list where erase() is spelt erase_after().
+template <typename T, typename Allocator>
+constexpr auto ConstCastIterator(
+    std::forward_list<T, Allocator>& c,
+    typename std::forward_list<T, Allocator>::const_iterator it) {
+  return c.erase_after(it, it);
+}
+
+// Specialized O(1) const casting for random access iterators. This is
+// necessary, because erase() is either not available (e.g. array-like
+// containers), or has O(n) complexity (e.g. std::deque or std::vector).
+template <typename Container,
+          typename ConstIter,
+          std::enable_if_t<internal::IsRandomAccessIter<ConstIter>>* = nullptr>
+constexpr auto ConstCastIterator(Container& c, ConstIter it) {
+  using std::begin;
+  using std::cbegin;
+  return begin(c) + (it - cbegin(c));
+}
+
 // Test to see if a set or map contains a particular key.
 // Returns true if the key is in the collection.
 template <typename Collection, typename Key>
@@ -154,6 +191,94 @@ class HasKeyType {
   static constexpr bool value = decltype(test<Collection>(nullptr))::value;
 };
 
+template <typename Map, typename Key, typename Value>
+std::pair<typename Map::iterator, bool> InsertOrAssignImpl(Map& map,
+                                                           Key&& key,
+                                                           Value&& value) {
+  auto lower = map.lower_bound(key);
+  if (lower != map.end() && !map.key_comp()(key, lower->first)) {
+    // key already exists, perform assignment.
+    lower->second = std::forward<Value>(value);
+    return {lower, false};
+  }
+
+  // key did not yet exist, insert it.
+  return {map.emplace_hint(lower, std::forward<Key>(key),
+                           std::forward<Value>(value)),
+          true};
+}
+
+template <typename Map, typename Key, typename Value>
+typename Map::iterator InsertOrAssignImpl(Map& map,
+                                          typename Map::const_iterator hint,
+                                          Key&& key,
+                                          Value&& value) {
+  auto&& key_comp = map.key_comp();
+  if ((hint == map.begin() || key_comp(std::prev(hint)->first, key))) {
+    if (hint == map.end() || key_comp(key, hint->first)) {
+      // *(hint - 1) < key < *hint => key did not exist and hint is correct.
+      return map.emplace_hint(hint, std::forward<Key>(key),
+                              std::forward<Value>(value));
+    }
+
+    if (!key_comp(hint->first, key)) {
+      // key == *hint => key already exists and hint is correct.
+      auto mutable_hint = ConstCastIterator(map, hint);
+      mutable_hint->second = std::forward<Value>(value);
+      return mutable_hint;
+    }
+  }
+
+  // hint was not helpful, dispatch to hintless version.
+  return InsertOrAssignImpl(map, std::forward<Key>(key),
+                            std::forward<Value>(value))
+      .first;
+}
+
+template <typename Map, typename Key, typename... Args>
+std::pair<typename Map::iterator, bool> TryEmplaceImpl(Map& map,
+                                                       Key&& key,
+                                                       Args&&... args) {
+  auto lower = map.lower_bound(key);
+  if (lower != map.end() && !map.key_comp()(key, lower->first)) {
+    // key already exists, do nothing.
+    return {lower, false};
+  }
+
+  // key did not yet exist, insert it.
+  return {map.emplace_hint(lower, std::piecewise_construct,
+                           std::forward_as_tuple(std::forward<Key>(key)),
+                           std::forward_as_tuple(std::forward<Args>(args)...)),
+          true};
+}
+
+template <typename Map, typename Key, typename... Args>
+typename Map::iterator TryEmplaceImpl(Map& map,
+                                      typename Map::const_iterator hint,
+                                      Key&& key,
+                                      Args&&... args) {
+  auto&& key_comp = map.key_comp();
+  if ((hint == map.begin() || key_comp(std::prev(hint)->first, key))) {
+    if (hint == map.end() || key_comp(key, hint->first)) {
+      // *(hint - 1) < key < *hint => key did not exist and hint is correct.
+      return map.emplace_hint(
+          hint, std::piecewise_construct,
+          std::forward_as_tuple(std::forward<Key>(key)),
+          std::forward_as_tuple(std::forward<Args>(args)...));
+    }
+
+    if (!key_comp(hint->first, key)) {
+      // key == *hint => no-op, return correct hint.
+      return ConstCastIterator(map, hint);
+    }
+  }
+
+  // hint was not helpful, dispatch to hintless version.
+  return TryEmplaceImpl(map, std::forward<Key>(key),
+                        std::forward<Args>(args)...)
+      .first;
+}
+
 }  // namespace internal
 
 // Test to see if a collection like a vector contains a particular value.
@@ -169,6 +294,74 @@ bool ContainsValue(const Collection& collection, const Value& value) {
          std::end(collection);
 }
 
+// Implementation of C++17's std::map::insert_or_assign as a free function.
+template <typename Map, typename Value>
+std::pair<typename Map::iterator, bool>
+InsertOrAssign(Map& map, const typename Map::key_type& key, Value&& value) {
+  return internal::InsertOrAssignImpl(map, key, std::forward<Value>(value));
+}
+
+template <typename Map, typename Value>
+std::pair<typename Map::iterator, bool>
+InsertOrAssign(Map& map, typename Map::key_type&& key, Value&& value) {
+  return internal::InsertOrAssignImpl(map, std::move(key),
+                                      std::forward<Value>(value));
+}
+
+// Implementation of C++17's std::map::insert_or_assign with hint as a free
+// function.
+template <typename Map, typename Value>
+typename Map::iterator InsertOrAssign(Map& map,
+                                      typename Map::const_iterator hint,
+                                      const typename Map::key_type& key,
+                                      Value&& value) {
+  return internal::InsertOrAssignImpl(map, hint, key,
+                                      std::forward<Value>(value));
+}
+
+template <typename Map, typename Value>
+typename Map::iterator InsertOrAssign(Map& map,
+                                      typename Map::const_iterator hint,
+                                      typename Map::key_type&& key,
+                                      Value&& value) {
+  return internal::InsertOrAssignImpl(map, hint, std::move(key),
+                                      std::forward<Value>(value));
+}
+
+// Implementation of C++17's std::map::try_emplace as a free function.
+template <typename Map, typename... Args>
+std::pair<typename Map::iterator, bool>
+TryEmplace(Map& map, const typename Map::key_type& key, Args&&... args) {
+  return internal::TryEmplaceImpl(map, key, std::forward<Args>(args)...);
+}
+
+template <typename Map, typename... Args>
+std::pair<typename Map::iterator, bool> TryEmplace(Map& map,
+                                                   typename Map::key_type&& key,
+                                                   Args&&... args) {
+  return internal::TryEmplaceImpl(map, std::move(key),
+                                  std::forward<Args>(args)...);
+}
+
+// Implementation of C++17's std::map::try_emplace with hint as a free
+// function.
+template <typename Map, typename... Args>
+typename Map::iterator TryEmplace(Map& map,
+                                  typename Map::const_iterator hint,
+                                  const typename Map::key_type& key,
+                                  Args&&... args) {
+  return internal::TryEmplaceImpl(map, hint, key, std::forward<Args>(args)...);
+}
+
+template <typename Map, typename... Args>
+typename Map::iterator TryEmplace(Map& map,
+                                  typename Map::const_iterator hint,
+                                  typename Map::key_type&& key,
+                                  Args&&... args) {
+  return internal::TryEmplaceImpl(map, hint, std::move(key),
+                                  std::forward<Args>(args)...);
+}
+
 // Returns true if the container is sorted.
 template <typename Container>
 bool STLIsSorted(const Container& cont) {

base/stl_util_unittest.cc

@@ -29,6 +29,9 @@
 
 namespace {
 
+using ::testing::IsNull;
+using ::testing::Pair;
+
 // Used as test case to ensure the various base::STLXxx functions don't require
 // more than operators "<" and "==" on values stored in containers.
 class ComparableValue {
@@ -86,6 +89,23 @@ void RunEraseIfTest() {
   }
 }
 
+template <typename Container>
+void RunConstCastIteratorTest() {
+  using std::begin;
+  using std::cbegin;
+
+  Container c = {1, 2, 3, 4, 5};
+  auto c_it = std::next(cbegin(c), 3);
+  auto it = base::ConstCastIterator(c, c_it);
+  static_assert(std::is_same<decltype(cbegin(std::declval<Container&>())),
+                             decltype(c_it)>::value,
+                "c_it is not a constant iterator.");
+  static_assert(std::is_same<decltype(begin(std::declval<Container&>())),
+                             decltype(it)>::value,
+                "it is not a iterator.");
+  EXPECT_EQ(c_it, it);
+}
+
 struct CustomIntHash {
   size_t operator()(int elem) const { return std::hash<int>()(elem) + 1; }
 };
@@ -280,6 +300,25 @@ TEST(STLUtilTest, GetUnderlyingContainer) {
   }
 }
 
+TEST(STLUtilTest, ConstCastIterator) {
+  // Sequence Containers
+  RunConstCastIteratorTest<std::forward_list<int>>();
+  RunConstCastIteratorTest<std::list<int>>();
+  RunConstCastIteratorTest<std::deque<int>>();
+  RunConstCastIteratorTest<std::vector<int>>();
+  RunConstCastIteratorTest<std::array<int, 5>>();
+  RunConstCastIteratorTest<std::initializer_list<int>>();
+  RunConstCastIteratorTest<int[5]>();
+
+  // Associative Containers
+  RunConstCastIteratorTest<std::set<int>>();
+  RunConstCastIteratorTest<std::multiset<int>>();
+
+  // Unordered Associative Containers
+  RunConstCastIteratorTest<std::unordered_set<int>>();
+  RunConstCastIteratorTest<std::unordered_multiset<int>>();
+}
+
 TEST(STLUtilTest, STLIsSorted) {
   {
     std::set<int> set;
@@ -606,6 +645,104 @@ TEST(ContainsValue, OrdinaryArrays) {
   EXPECT_TRUE(ContainsValue(allowed_chars_including_nul, 0));
 }
 
+TEST(STLUtilTest, InsertOrAssign) {
+  std::map<std::string, int> my_map;
+  auto result = InsertOrAssign(my_map, "Hello", 42);
+  EXPECT_THAT(*result.first, Pair("Hello", 42));
+  EXPECT_TRUE(result.second);
+
+  result = InsertOrAssign(my_map, "Hello", 43);
+  EXPECT_THAT(*result.first, Pair("Hello", 43));
+  EXPECT_FALSE(result.second);
+}
+
+TEST(STLUtilTest, InsertOrAssignHint) {
+  std::map<std::string, int> my_map;
+  auto result = InsertOrAssign(my_map, my_map.end(), "Hello", 42);
+  EXPECT_THAT(*result, Pair("Hello", 42));
+
+  result = InsertOrAssign(my_map, my_map.begin(), "Hello", 43);
+  EXPECT_THAT(*result, Pair("Hello", 43));
+}
+
+TEST(STLUtilTest, InsertOrAssignWrongHints) {
+  std::map<int, int> my_map;
+  // Since we insert keys in sorted order, my_map.begin() will be a wrong hint
+  // after the first iteration. Check that insertion happens anyway.
+  for (int i = 0; i < 10; ++i) {
+    SCOPED_TRACE(i);
+    auto result = InsertOrAssign(my_map, my_map.begin(), i, i);
+    EXPECT_THAT(*result, Pair(i, i));
+  }
+
+  // Overwrite the keys we just inserted. Since we no longer insert into the
+  // map, my_map.end() will be a wrong hint for all iterations but the last.
+  for (int i = 0; i < 10; ++i) {
+    SCOPED_TRACE(10 + i);
+    auto result = InsertOrAssign(my_map, my_map.end(), i, 10 + i);
+    EXPECT_THAT(*result, Pair(i, 10 + i));
+  }
+}
+
+TEST(STLUtilTest, TryEmplace) {
+  std::map<std::string, std::unique_ptr<int>> my_map;
+  auto result = TryEmplace(my_map, "Hello", nullptr);
+  EXPECT_THAT(*result.first, Pair("Hello", IsNull()));
+  EXPECT_TRUE(result.second);
+
+  auto new_value = std::make_unique<int>(42);
+  result = TryEmplace(my_map, "Hello", std::move(new_value));
+  EXPECT_THAT(*result.first, Pair("Hello", IsNull()));
+  EXPECT_FALSE(result.second);
+  // |new_value| should not be touched following a failed insertion.
+  ASSERT_NE(nullptr, new_value);
+  EXPECT_EQ(42, *new_value);
+
+  result = TryEmplace(my_map, "World", std::move(new_value));
+  EXPECT_EQ("World", result.first->first);
+  EXPECT_EQ(42, *result.first->second);
+  EXPECT_TRUE(result.second);
+  EXPECT_EQ(nullptr, new_value);
+}
+
+TEST(STLUtilTest, TryEmplaceHint) {
+  std::map<std::string, std::unique_ptr<int>> my_map;
+  auto result = TryEmplace(my_map, my_map.begin(), "Hello", nullptr);
+  EXPECT_THAT(*result, Pair("Hello", IsNull()));
+
+  auto new_value = std::make_unique<int>(42);
+  result = TryEmplace(my_map, result, "Hello", std::move(new_value));
+  EXPECT_THAT(*result, Pair("Hello", IsNull()));
+  // |new_value| should not be touched following a failed insertion.
+  ASSERT_NE(nullptr, new_value);
+  EXPECT_EQ(42, *new_value);
+
+  result = TryEmplace(my_map, result, "World", std::move(new_value));
+  EXPECT_EQ("World", result->first);
+  EXPECT_EQ(42, *result->second);
+  EXPECT_EQ(nullptr, new_value);
+}
+
+TEST(STLUtilTest, TryEmplaceWrongHints) {
+  std::map<int, int> my_map;
+  // Since we emplace keys in sorted order, my_map.begin() will be a wrong hint
+  // after the first iteration. Check that emplacement happens anyway.
+  for (int i = 0; i < 10; ++i) {
+    SCOPED_TRACE(i);
+    auto result = TryEmplace(my_map, my_map.begin(), i, i);
+    EXPECT_THAT(*result, Pair(i, i));
+  }
+
+  // Fail to overwrite the keys we just inserted. Since we no longer emplace
+  // into the map, my_map.end() will be a wrong hint for all tried emplacements
+  // but the last.
+  for (int i = 0; i < 10; ++i) {
+    SCOPED_TRACE(10 + i);
+    auto result = TryEmplace(my_map, my_map.end(), i, 10 + i);
+    EXPECT_THAT(*result, Pair(i, i));
+  }
+}
+
 TEST(STLUtilTest, OptionalOrNullptr) {
   Optional<float> optional;
   EXPECT_EQ(nullptr, base::OptionalOrNullptr(optional));