get_ptr(mapOfMaps, key, key...)

Summary:
Since this is a not-uncommon pattern now:

```lang=cpp
if (auto found1 = get_ptr(m, k1)) {
  if (auto found2 = get_ptr(*found, k2)) {
    return use(*found2);
  }
}
```

This diff enables:

```lang=cpp
if (auto found = get_ptr(m, k1, k2)) {
  return use(*found);
}

```

Note that it works for const and non-const maps, returning a correspondingly mutable pointer.

Reviewed By: luciang, yfeldblum

Differential Revision: D3812701

fbshipit-source-id: 77ace9f5dca6cdc4cae7e6dfb9e5fc1075b5b571

diff --git a/folly/MapUtil.h b/folly/MapUtil.h
index 54ed01053b52a2ea571626998a04cd968781f194..063bb0f43ba24e2a7a5d35f5ff7ab255970d89e7 100644 (file)
--- a/folly/MapUtil.h
+++ b/folly/MapUtil.h
@@ -148,4 +148,41 @@ typename Map::mapped_type* get_ptr(
   return (pos != map.end() ? &pos->second : nullptr);
 }
 
+namespace detail {
+template <
+    class T,
+    size_t pathLength,
+    class = typename std::enable_if<(pathLength > 0)>::type>
+struct NestedMapType {
+  using type = typename NestedMapType<T, pathLength - 1>::type::mapped_type;
+};
+
+template <class T>
+struct NestedMapType<T, 1> {
+  using type = typename T::mapped_type;
+};
+}
+
+/**
+ * Given a map of maps and a path of keys, return a pointer to the nested value,
+ * or nullptr if the key doesn't exist in the map.
+ */
+template <class Map, class Key1, class Key2, class... Keys>
+auto get_ptr(
+    const Map& map,
+    const Key1& key1,
+    const Key2& key2,
+    const Keys&... keys) ->
+    typename detail::NestedMapType<Map, 2 + sizeof...(Keys)>::type const* {
+  auto pos = map.find(key1);
+  return pos != map.end() ? get_ptr(pos->second, key2, keys...) : nullptr;
+}
+
+template <class Map, class Key1, class Key2, class... Keys>
+auto get_ptr(Map& map, const Key1& key1, const Key2& key2, const Keys&... keys)
+    -> typename detail::NestedMapType<Map, 2 + sizeof...(Keys)>::type* {
+  auto pos = map.find(key1);
+  return pos != map.end() ? get_ptr(pos->second, key2, keys...) : nullptr;
+}
+
 }  // namespace folly

diff --git a/folly/test/MapUtilTest.cpp b/folly/test/MapUtilTest.cpp
index 7f28291f55bfb2d6aa440a585b814089ef294e10..1fa91df0f054d004d9743de460243ab99d143375 100644 (file)
--- a/folly/test/MapUtilTest.cpp
+++ b/folly/test/MapUtilTest.cpp
@@ -17,6 +17,7 @@
 #include <folly/MapUtil.h>
 
 #include <map>
+#include <unordered_map>
 
 #include <folly/portability/GTest.h>
 
@@ -98,3 +99,42 @@ TEST(MapUtil, get_ptr) {
   *get_ptr(m, 1) = 4;
   EXPECT_EQ(4, m.at(1));
 }
+
+TEST(MapUtil, get_ptr_path_simple) {
+  using std::map;
+  map<int, map<int, map<int, map<int, int>>>> m{{1, {{2, {{3, {{4, 5}}}}}}}};
+  EXPECT_EQ(5, *get_ptr(m, 1, 2, 3, 4));
+  EXPECT_TRUE(get_ptr(m, 1, 2, 3, 4));
+  EXPECT_FALSE(get_ptr(m, 1, 2, 3, 0));
+  EXPECT_TRUE(get_ptr(m, 1, 2, 3));
+  EXPECT_FALSE(get_ptr(m, 1, 2, 0));
+  EXPECT_TRUE(get_ptr(m, 1, 2));
+  EXPECT_FALSE(get_ptr(m, 1, 0));
+  EXPECT_TRUE(get_ptr(m, 1));
+  EXPECT_FALSE(get_ptr(m, 0));
+  const auto& cm = m;
+  ++*get_ptr(m, 1, 2, 3, 4);
+  EXPECT_EQ(6, *get_ptr(cm, 1, 2, 3, 4));
+  EXPECT_TRUE(get_ptr(cm, 1, 2, 3, 4));
+  EXPECT_FALSE(get_ptr(cm, 1, 2, 3, 0));
+}
+
+TEST(MapUtil, get_ptr_path_mixed) {
+  using std::map;
+  using std::unordered_map;
+  using std::string;
+  unordered_map<string, map<int, map<string, int>>> m{{"a", {{1, {{"b", 7}}}}}};
+  EXPECT_EQ(7, *get_ptr(m, "a", 1, "b"));
+  EXPECT_TRUE(get_ptr(m, "a", 1, "b"));
+  EXPECT_FALSE(get_ptr(m, "b", 1, "b"));
+  EXPECT_FALSE(get_ptr(m, "a", 2, "b"));
+  EXPECT_FALSE(get_ptr(m, "a", 1, "c"));
+  EXPECT_TRUE(get_ptr(m, "a", 1, "b"));
+  EXPECT_TRUE(get_ptr(m, "a", 1));
+  EXPECT_TRUE(get_ptr(m, "a"));
+  const auto& cm = m;
+  ++*get_ptr(m, "a", 1, "b");
+  EXPECT_EQ(8, *get_ptr(cm, "a", 1, "b"));
+  EXPECT_TRUE(get_ptr(cm, "a", 1, "b"));
+  EXPECT_FALSE(get_ptr(cm, "b", 1, "b"));
+}