add `getAllocatedMemorySize` for F14 sets, fix corner cases

Summary:
This diff adds a method to calculate the amount of memory allocated by F14
sets. It mimics the node structure in `std::unordered_map` to calculate the
results. For certain keys and hashers, the libstdc++ implementation of
`std::unordered_map` caches a hash code in the node. libcpp seems to do this
for everything.

https://github.com/gcc-mirror/gcc/blob/9ce5a00d2d3b8b107cb1ca0e23bf1e31a095f9b6/libstdc%2B%2B-v3/include/bits/hashtable.h#L45-L49
https://github.com/llvm-mirror/libcxx/blob/58bcf28c63e0cf2ee5828e7811829c98bd244dda/include/__hash_table#L117

Depends on D7715573.

Reviewed By: nbronson

Differential Revision: D7728281

fbshipit-source-id: 0aa3ab60c746d3ab6bdb7c879f2f987110e0329d

folly/container/F14Map.h

@@ -47,13 +47,6 @@ namespace folly {
 
 namespace f14 {
 namespace detail {
-template <typename K>
-struct CachedHashOverhead : std::integral_constant<std::size_t, 0> {};
-
-template <typename... Args>
-struct CachedHashOverhead<std::basic_string<Args...>>
-    : std::integral_constant<std::size_t, sizeof(std::size_t)> {};
-
 template <typename K, typename M, typename H, typename E, typename A>
 class F14BasicMap : public std::unordered_map<K, M, H, E, A> {
   using Super = std::unordered_map<K, M, H, E, A>;
@@ -62,14 +55,12 @@ class F14BasicMap : public std::unordered_map<K, M, H, E, A> {
   using Super::Super;
   F14BasicMap() : Super() {}
 
-  // Accounts for allocated memory only, does not include sizeof(*this).
+  // Approximates allocated memory, does not include sizeof(*this), also does
+  // not account for memory allocator fragmentation
   typename Super::size_type getAllocatedMemorySize() const {
-    auto nodeSize = sizeof(typename Super::pointer) +
-        CachedHashOverhead<typename Super::key_type>::value +
-        sizeof(typename Super::value_type);
     auto bc = this->bucket_count();
     return (bc == 1 ? 0 : bc) * sizeof(typename Super::pointer) +
-        this->size() * nodeSize;
+        this->size() * sizeof(StdNodeReplica<K, typename Super::value_type, H>);
   }
 };
 } // namespace detail

folly/container/F14Set.h

@@ -39,18 +39,39 @@
 
 namespace folly {
 
+namespace f14 {
+namespace detail {
 template <typename K, typename H, typename E, typename A>
-class F14NodeSet : public std::unordered_set<K, H, E, A> {
+class F14BasicSet : public std::unordered_set<K, H, E, A> {
   using Super = std::unordered_set<K, H, E, A>;
 
+ public:
+  using Super::Super;
+  F14BasicSet() : Super() {}
+
+  // Approximates allocated memory, does not include sizeof(*this), also does
+  // not account for memory allocator fragmentation
+  typename Super::size_type getAllocatedMemorySize() const {
+    auto bc = this->bucket_count();
+    return (bc == 1 ? 0 : bc) * sizeof(typename Super::pointer) +
+        this->size() * sizeof(StdNodeReplica<K, typename Super::value_type, H>);
+  }
+};
+} // namespace detail
+} // namespace f14
+
+template <typename K, typename H, typename E, typename A>
+class F14NodeSet : public f14::detail::F14BasicSet<K, H, E, A> {
+  using Super = f14::detail::F14BasicSet<K, H, E, A>;
+
  public:
   using Super::Super;
   F14NodeSet() : Super() {}
 };
 
 template <typename K, typename H, typename E, typename A>
-class F14ValueSet : public std::unordered_set<K, H, E, A> {
-  using Super = std::unordered_set<K, H, E, A>;
+class F14ValueSet : public f14::detail::F14BasicSet<K, H, E, A> {
+  using Super = f14::detail::F14BasicSet<K, H, E, A>;
 
  public:
   using Super::Super;
@@ -58,8 +79,8 @@ class F14ValueSet : public std::unordered_set<K, H, E, A> {
 };
 
 template <typename K, typename H, typename E, typename A>
-class F14VectorSet : public std::unordered_set<K, H, E, A> {
-  using Super = std::unordered_set<K, H, E, A>;
+class F14VectorSet : public f14::detail::F14BasicSet<K, H, E, A> {
+  using Super = f14::detail::F14BasicSet<K, H, E, A>;
 
  public:
   using Super::Super;
@@ -239,6 +260,11 @@ class F14BasicSet {
     return table_.max_size();
   }
 
+  // Accounts for allocated memory only, does not include sizeof(*this).
+  std::size_t getAllocatedMemorySize() const {
+    return table_.getAllocatedMemorySize();
+  }
+
   F14TableStats computeStats() const {
     return table_.computeStats();
   }

folly/container/detail/F14IntrinsicsAvailability.h

@@ -26,7 +26,7 @@
 # else
 #  define FOLLY_F14_VECTOR_INTRINSICS_AVAILABLE 0
 #  pragma message                                       \
-    "Vector intrinsics unavailable on this platform, " \
+    "Vector intrinsics unavailable on this platform, "  \
     "falling back to std::unordered_map / set"
 # endif
 #endif

folly/container/detail/F14Table.h

@@ -35,6 +35,7 @@
 #include <folly/Portability.h>
 #include <folly/ScopeGuard.h>
 #include <folly/Traits.h>
+#include <folly/functional/Invoke.h>
 #include <folly/lang/Assume.h>
 #include <folly/lang/Exception.h>
 #include <folly/lang/Launder.h>
@@ -56,7 +57,6 @@
 #endif
 
 namespace folly {
-
 struct F14TableStats {
   char const* policy;
   std::size_t size{0};
@@ -88,6 +88,54 @@ struct F14TableStats {
   }
 };
 
+namespace f14 {
+namespace detail {
+
+#if defined(_LIBCPP_VERSION)
+
+template <typename K, typename V, typename H>
+struct StdNodeReplica {
+  void* next;
+  std::size_t hash;
+  V value;
+};
+
+#else
+
+template <typename H>
+struct StdIsFastHash : std::true_type {};
+template <>
+struct StdIsFastHash<std::hash<long double>> : std::false_type {};
+template <typename... Args>
+struct StdIsFastHash<std::hash<std::basic_string<Args...>>> : std::false_type {
+};
+
+// TODO: add specialization for std::basic_string_view
+
+// mimic internal node of unordered containers in STL to estimate the size
+template <typename K, typename V, typename H, typename Enable = void>
+struct StdNodeReplica {
+  void* next;
+  V value;
+};
+template <typename K, typename V, typename H>
+struct StdNodeReplica<
+    K,
+    V,
+    H,
+    std::enable_if_t<
+        !StdIsFastHash<H>::value ||
+        !folly::is_nothrow_invocable<H, K>::value>> {
+  void* next;
+  V value;
+  std::size_t hash;
+};
+
+#endif
+
+} // namespace detail
+} // namespace f14
+
 #if FOLLY_F14_VECTOR_INTRINSICS_AVAILABLE
 namespace f14 {
 namespace detail {

folly/container/test/F14MapTest.cpp

@@ -17,6 +17,7 @@
 #include <folly/container/F14Map.h>
 
 #include <folly/Conv.h>
+#include <folly/FBString.h>
 #include <folly/container/test/F14TestUtil.h>
 #include <folly/portability/GTest.h>
 
@@ -59,7 +60,7 @@ void testAllocatedMemorySize() {
   TMap<K, V, DefaultHasher<K>, DefaultKeyEqual<K>, A> m;
   EXPECT_EQ(A::getAllocatedMemorySize(), m.getAllocatedMemorySize());
 
-  for (size_t i = 0; i < 1; ++i) {
+  for (size_t i = 0; i < 1000; ++i) {
     m.insert(std::make_pair(folly::to<K>(i), V{}));
     EXPECT_EQ(A::getAllocatedMemorySize(), m.getAllocatedMemorySize());
   }
@@ -75,9 +76,13 @@ void runAllocatedMemorySizeTest() {
 } // namespace
 
 TEST(F14Map, getAllocatedMemorySize) {
+  runAllocatedMemorySizeTest<bool, bool>();
   runAllocatedMemorySizeTest<int, int>();
+  runAllocatedMemorySizeTest<bool, std::string>();
+  runAllocatedMemorySizeTest<long double, std::string>();
   runAllocatedMemorySizeTest<std::string, int>();
   runAllocatedMemorySizeTest<std::string, std::string>();
+  runAllocatedMemorySizeTest<folly::fbstring, long>();
 }
 
 ///////////////////////////////////

folly/container/test/F14SetTest.cpp

@@ -15,6 +15,9 @@
  */
 
 #include <folly/container/F14Set.h>
+
+#include <folly/Conv.h>
+#include <folly/FBString.h>
 #include <folly/container/test/F14TestUtil.h>
 #include <folly/portability/GTest.h>
 
@@ -42,6 +45,42 @@ TEST(F14Set, customSwap) {
   testCustomSwap<folly::F14FastSet>();
 }
 
+namespace {
+template <
+    template <typename, typename, typename, typename> class TSet,
+    typename K>
+void testAllocatedMemorySize() {
+  using namespace folly::f14;
+  using A = SwapTrackingAlloc<K>;
+
+  A::resetTracking();
+  TSet<K, DefaultHasher<K>, DefaultKeyEqual<K>, A> m;
+  EXPECT_EQ(A::getAllocatedMemorySize(), m.getAllocatedMemorySize());
+
+  for (size_t i = 0; i < 1000; ++i) {
+    m.insert(folly::to<K>(i));
+    EXPECT_EQ(A::getAllocatedMemorySize(), m.getAllocatedMemorySize());
+  }
+}
+
+template <typename K>
+void runAllocatedMemorySizeTest() {
+  testAllocatedMemorySize<folly::F14ValueSet, K>();
+  testAllocatedMemorySize<folly::F14NodeSet, K>();
+  testAllocatedMemorySize<folly::F14VectorSet, K>();
+  testAllocatedMemorySize<folly::F14FastSet, K>();
+}
+} // namespace
+
+TEST(F14Set, getAllocatedMemorySize) {
+  runAllocatedMemorySizeTest<bool>();
+  runAllocatedMemorySizeTest<int>();
+  runAllocatedMemorySizeTest<long>();
+  runAllocatedMemorySizeTest<long double>();
+  runAllocatedMemorySizeTest<std::string>();
+  runAllocatedMemorySizeTest<folly::fbstring>();
+}
+
 ///////////////////////////////////
 #if FOLLY_F14_VECTOR_INTRINSICS_AVAILABLE
 ///////////////////////////////////