get_default and get_ref_default variants taking functions

Summary:[Folly] `get_default` and `get_ref_default` variants taking functions.

Useful if the default value is computationally expensive to construct or requires IO.

Reviewed By: andriigrynenko, spacedentist

Differential Revision: D3189247

fb-gh-sync-id: 51c64293f8712d7590348d53cbfd892a5efd9e82
fbshipit-source-id: 51c64293f8712d7590348d53cbfd892a5efd9e82

diff --git a/folly/MapUtil.h b/folly/MapUtil.h
index 0df5d788868bd6d7cd34fcd3da059e11e96a6cf7..f042bd2d5e86f328d4c0df508cb79a3aecc44d7a 100644 (file)
--- a/folly/MapUtil.h
+++ b/folly/MapUtil.h
@@ -34,6 +34,22 @@ typename Map::mapped_type get_default(
   return (pos != map.end() ? pos->second : dflt);
 }
 
+/**
+ * Give a map and a key, return the value corresponding to the key in the map,
+ * or a given default value if the key doesn't exist in the map.
+ */
+template <
+    class Map,
+    typename Func,
+    typename = typename std::enable_if<std::is_convertible<
+        typename std::result_of<Func()>::type,
+        typename Map::mapped_type>::value>::type>
+typename Map::mapped_type
+get_default(const Map& map, const typename Map::key_type& key, Func&& dflt) {
+  auto pos = map.find(key);
+  return pos != map.end() ? pos->second : dflt();
+}
+
 /**
  * Given a map and a key, return the value corresponding to the key in the map,
  * or throw an exception of the specified type.
@@ -77,6 +93,25 @@ const typename Map::mapped_type& get_ref_default(
   return (pos != map.end() ? pos->second : dflt);
 }
 
+/**
+ * Given a map and a key, return a reference to the value corresponding to the
+ * key in the map, or the given default reference if the key doesn't exist in
+ * the map.
+ */
+template <
+    class Map,
+    typename Func,
+    typename = typename std::enable_if<std::is_convertible<
+        typename std::result_of<Func()>::type,
+        const typename Map::mapped_type&>::value>::type>
+const typename Map::mapped_type& get_ref_default(
+    const Map& map,
+    const typename Map::key_type& key,
+    Func&& dflt) {
+  auto pos = map.find(key);
+  return (pos != map.end() ? pos->second : dflt());
+}
+
 /**
  * Given a map and a key, return a pointer to the value corresponding to the
  * key in the map, or nullptr if the key doesn't exist in the map.

diff --git a/folly/test/MapUtilTest.cpp b/folly/test/MapUtilTest.cpp
index 4e21fd0a2eaf0320691a35c99b23c05bcfd9cf00..d31881b1dd2a76441f0037effd07c051a9b834ce 100644 (file)
--- a/folly/test/MapUtilTest.cpp
+++ b/folly/test/MapUtilTest.cpp
@@ -29,6 +29,14 @@ TEST(MapUtil, get_default) {
   EXPECT_EQ(0, get_default(m, 3));
 }
 
+TEST(MapUtil, get_default_function) {
+  std::map<int, int> m;
+  m[1] = 2;
+  EXPECT_EQ(2, get_default(m, 1, [] { return 42; }));
+  EXPECT_EQ(42, get_default(m, 2, [] { return 42; }));
+  EXPECT_EQ(0, get_default(m, 3));
+}
+
 TEST(MapUtil, get_or_throw) {
   std::map<int, int> m;
   m[1] = 2;
@@ -57,6 +65,19 @@ TEST(MapUtil, get_ref_default) {
   const int i = 42;
   EXPECT_EQ(2, get_ref_default(m, 1, i));
   EXPECT_EQ(42, get_ref_default(m, 2, i));
+  EXPECT_EQ(std::addressof(i), std::addressof(get_ref_default(m, 2, i)));
+}
+
+TEST(MapUtil, get_ref_default_function) {
+  std::map<int, int> m;
+  m[1] = 2;
+  const int i = 42;
+  EXPECT_EQ(2, get_ref_default(m, 1, [&i]() -> const int& { return i; }));
+  EXPECT_EQ(42, get_ref_default(m, 2, [&i]() -> const int& { return i; }));
+  EXPECT_EQ(
+      std::addressof(i),
+      std::addressof(
+          get_ref_default(m, 2, [&i]() -> const int& { return i; })));
 }
 
 TEST(MapUtil, get_ptr) {