base::flat_map: add cxx11 style at() function

Fixed: 1026452
Change-Id: I705bde13f9f3cc6c45d5ef9365cb9c685a07c5f9
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1933849
Commit-Queue: Jiewei Qian  <qjw@chromium.org>
Reviewed-by: Jan Wilken Dörrie <jdoerrie@chromium.org>
Reviewed-by: Daniel Cheng <dcheng@chromium.org>
Reviewed-by: Giovanni Ortuño Urquidi <ortuno@chromium.org>
Cr-Commit-Position: refs/heads/master@{#719478}

base/containers/flat_map.h

@@ -101,6 +101,8 @@ struct GetKeyFromValuePairFirst {
 // Insert and accessor functions:
 //   mapped_type&         operator[](const key_type&);
 //   mapped_type&         operator[](key_type&&);
+//   mapped_type&         at(const K&);
+//   const mapped_type&   at(const K&) const;
 //   pair<iterator, bool> insert(const value_type&);
 //   pair<iterator, bool> insert(value_type&&);
 //   iterator             insert(const_iterator hint, const value_type&);
@@ -197,6 +199,12 @@ class flat_map : public ::base::internal::flat_tree<
   // Takes the first if there are duplicates in the initializer list.
   flat_map& operator=(std::initializer_list<value_type> ilist);
 
+  // Out-of-bound calls to at() will CHECK.
+  template <class K>
+  mapped_type& at(const K& key);
+  template <class K>
+  const mapped_type& at(const K& key) const;
+
   // --------------------------------------------------------------------------
   // Map-specific insert operations.
   //
@@ -272,6 +280,26 @@ auto flat_map<Key, Mapped, Compare>::operator=(
   return *this;
 }
 
+// ----------------------------------------------------------------------------
+// Lookups.
+
+template <class Key, class Mapped, class Compare>
+template <class K>
+auto flat_map<Key, Mapped, Compare>::at(const K& key) -> mapped_type& {
+  iterator found = tree::find(key);
+  CHECK(found != tree::end());
+  return found->second;
+}
+
+template <class Key, class Mapped, class Compare>
+template <class K>
+auto flat_map<Key, Mapped, Compare>::at(const K& key) const
+    -> const mapped_type& {
+  const_iterator found = tree::find(key);
+  CHECK(found != tree::cend());
+  return found->second;
+}
+
 // ----------------------------------------------------------------------------
 // Insert operations.
 

base/containers/flat_map_unittest.cc

@@ -8,6 +8,7 @@
 #include <vector>
 
 #include "base/macros.h"
+#include "base/strings/string_piece.h"
 #include "base/test/move_only_int.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -151,6 +152,33 @@ TEST(FlatMap, SubscriptMoveOnlyKey) {
   EXPECT_EQ(44, m[MoveOnlyInt(1)]);
 }
 
+// Mapped& at(const Key&)
+// const Mapped& at(const Key&) const
+TEST(FlatMap, AtFunction) {
+  base::flat_map<int, std::string> m = {{1, "a"}, {2, "b"}};
+
+  // Basic Usage.
+  EXPECT_EQ("a", m.at(1));
+  EXPECT_EQ("b", m.at(2));
+
+  // Const reference works.
+  const std::string& const_ref = base::as_const(m).at(1);
+  EXPECT_EQ("a", const_ref);
+
+  // Reference works, can operate on the string.
+  m.at(1)[0] = 'x';
+  EXPECT_EQ("x", m.at(1));
+
+  // Out-of-bounds will CHECK.
+  EXPECT_DEATH_IF_SUPPORTED(m.at(-1), "");
+  EXPECT_DEATH_IF_SUPPORTED({ m.at(-1)[0] = 'z'; }, "");
+
+  // Heterogeneous look-up works.
+  base::flat_map<std::string, int> m2 = {{"a", 1}, {"b", 2}};
+  EXPECT_EQ(1, m2.at(base::StringPiece("a")));
+  EXPECT_EQ(2, base::as_const(m2).at(base::StringPiece("b")));
+}
+
 // insert_or_assign(K&&, M&&)
 TEST(FlatMap, InsertOrAssignMoveOnlyKey) {
   base::flat_map<MoveOnlyInt, MoveOnlyInt> m;