From 6871c5cc6f02cea70038893bf795f9968255940d Mon Sep 17 00:00:00 2001 From: Tom Jackson Date: Fri, 15 Jul 2016 17:07:29 -0700 Subject: [PATCH] get_or_throw(map, key) returns references Differential Revision: D3572671 fbshipit-source-id: a80390921b41e47ed2794d48d943a9e4060c7135 --- folly/MapUtil.h | 17 +++++++++++++++-- folly/test/MapUtilTest.cpp | 7 +++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/folly/MapUtil.h b/folly/MapUtil.h index f042bd2d..143dadc2 100644 --- 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 -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(exceptionStrPrefix, key)); +} + +template +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 d31881b1..ee1e7b5b 100644 --- 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) { -- 2.34.1