get_or_throw(map, key) returns references

Differential Revision: D3572671

fbshipit-source-id: a80390921b41e47ed2794d48d943a9e4060c7135

diff --git a/folly/MapUtil.h b/folly/MapUtil.h
index f042bd2d5e86f328d4c0df508cb79a3aecc44d7a..143dadc26314e6f16c50a32d49381dbaa398d763 100644 (file)
--- a/folly/MapUtil.h
+++ b/folly/MapUtil.h
@@ -55,8 +55,21 @@ get_default(const Map& map, const typename Map::key_type& key, Func&& dflt) {
  * or throw an exception of the specified type.
  */
 template <class E = std::out_of_range, class Map>
-typename Map::mapped_type get_or_throw(
-    const Map& map, const typename Map::key_type& key,
+const typename Map::mapped_type& get_or_throw(
+    const Map& map,
+    const typename Map::key_type& key,
+    const std::string& exceptionStrPrefix = std::string()) {
+  auto pos = map.find(key);
+  if (pos != map.end()) {
+    return pos->second;
+  }
+  throw E(folly::to<std::string>(exceptionStrPrefix, key));
+}
+
+template <class E = std::out_of_range, class Map>
+typename Map::mapped_type& get_or_throw(
+    Map& map,
+    const typename Map::key_type& key,
     const std::string& exceptionStrPrefix = std::string()) {
   auto pos = map.find(key);
   if (pos != map.end()) {

diff --git a/folly/test/MapUtilTest.cpp b/folly/test/MapUtilTest.cpp
index d31881b1dd2a76441f0037effd07c051a9b834ce..ee1e7b5bd2290b05c69dce7516f210c569656023 100644 (file)
--- a/folly/test/MapUtilTest.cpp
+++ b/folly/test/MapUtilTest.cpp
@@ -42,6 +42,13 @@ TEST(MapUtil, get_or_throw) {
   m[1] = 2;
   EXPECT_EQ(2, get_or_throw(m, 1));
   EXPECT_THROW(get_or_throw(m, 2), std::out_of_range);
+  EXPECT_EQ(&m[1], &get_or_throw(m, 1));
+  get_or_throw(m, 1) = 3;
+  EXPECT_EQ(3, get_or_throw(m, 1));
+  const auto& cm = m;
+  EXPECT_EQ(&m[1], &get_or_throw(cm, 1));
+  EXPECT_EQ(3, get_or_throw(cm, 1));
+  EXPECT_THROW(get_or_throw(cm, 2), std::out_of_range);
 }
 
 TEST(MapUtil, get_or_throw_specified) {