add eraseInto to F14 maps

Summary:
It's useful to be able to move the key and value
out of a hash table as that entry is being erased. Since
std::unordered_map's extract API isn't implementable for hash tables
that perform bulk allocation, we expose similar functionality via
F14{Fast,Value,Vector,Node}Map::eraseInto.  eraseInto gives a key_type&&
and mapped_type&& to the specified callback, allowing it to move them
elsewhere before the underlying value_type is destroyed.

The keen observer will note that there is no way to move the key out
of a std::pair<key_type const,mapped_type> without relying on undefined
behavior. We already rely on this form of UB to avoid copying keys during
rehash, and perform it only when we have total control over the lifecycle
of the object and the memory backing the value_type.

eraseInto was previously implemented for F14 sets.

Reviewed By: shixiao

Differential Revision: D19777644

fbshipit-source-id: 9904224af2c6e51abeadae62cc5e66f7bf4b1ae7

folly/container/F14Map.h

@@ -477,40 +477,90 @@ class F14BasicMap {
   }
 
   FOLLY_ALWAYS_INLINE iterator erase(const_iterator pos) {
-    // If we are inlined then gcc and clang can optimize away all of the
-    // work of itemPos.advance() if our return value is discarded.
-    auto itemPos = table_.unwrapIter(pos);
-    table_.eraseIter(itemPos);
-    itemPos.advanceLikelyDead();
-    return table_.makeIter(itemPos);
+    return eraseInto(pos, [](key_type&&, mapped_type&&) {});
   }
 
   // This form avoids ambiguity when key_type has a templated constructor
   // that accepts const_iterator
   FOLLY_ALWAYS_INLINE iterator erase(iterator pos) {
+    return eraseInto(pos, [](key_type&&, mapped_type&&) {});
+  }
+
+  iterator erase(const_iterator first, const_iterator last) {
+    return eraseInto(first, last, [](key_type&&, mapped_type&&) {});
+  }
+
+  size_type erase(key_type const& key) {
+    return eraseInto(key, [](key_type&&, mapped_type&&) {});
+  }
+
+  template <typename K>
+  EnableHeterogeneousErase<K, size_type> erase(K const& key) {
+    return eraseInto(key, [](key_type&&, mapped_type&&) {});
+  }
+
+ protected:
+  template <typename BeforeDestroy>
+  FOLLY_ALWAYS_INLINE void tableEraseIterInto(
+      ItemIter pos,
+      BeforeDestroy&& beforeDestroy) {
+    table_.eraseIterInto(pos, [&](value_type&& v) {
+      auto p = Policy::moveValue(v);
+      beforeDestroy(std::move(p.first), std::move(p.second));
+    });
+  }
+
+  template <typename K, typename BeforeDestroy>
+  FOLLY_ALWAYS_INLINE std::size_t tableEraseKeyInto(
+      K const& key,
+      BeforeDestroy&& beforeDestroy) {
+    return table_.eraseKeyInto(key, [&](value_type&& v) {
+      auto p = Policy::moveValue(v);
+      beforeDestroy(std::move(p.first), std::move(p.second));
+    });
+  }
+
+ public:
+  // eraseInto contains the same overloads as erase but provides
+  // an additional callback argument which is called with an rvalue
+  // reference (not const) to the key and an rvalue reference to the
+  // mapped value directly before it is destroyed. This can be used
+  // to extract an entry out of a F14Map while avoiding a copy.
+  template <typename BeforeDestroy>
+  FOLLY_ALWAYS_INLINE iterator
+  eraseInto(const_iterator pos, BeforeDestroy&& beforeDestroy) {
+    // If we are inlined then gcc and clang can optimize away all of the
+    // work of itemPos.advance() if our return value is discarded.
     auto itemPos = table_.unwrapIter(pos);
-    table_.eraseIter(itemPos);
+    tableEraseIterInto(itemPos, beforeDestroy);
     itemPos.advanceLikelyDead();
     return table_.makeIter(itemPos);
   }
 
-  iterator erase(const_iterator first, const_iterator last) {
+  template <typename BeforeDestroy>
+  iterator eraseInto(
+      const_iterator first,
+      const_iterator last,
+      BeforeDestroy&& beforeDestroy) {
     auto itemFirst = table_.unwrapIter(first);
     auto itemLast = table_.unwrapIter(last);
     while (itemFirst != itemLast) {
-      table_.eraseIter(itemFirst);
+      tableEraseIterInto(itemFirst, beforeDestroy);
       itemFirst.advance();
     }
     return table_.makeIter(itemFirst);
   }
 
-  size_type erase(key_type const& key) {
-    return table_.eraseKey(key);
+  template <typename BeforeDestroy>
+  size_type eraseInto(key_type const& key, BeforeDestroy&& beforeDestroy) {
+    return tableEraseKeyInto(key, beforeDestroy);
   }
 
-  template <typename K>
-  EnableHeterogeneousErase<K, size_type> erase(K const& key) {
-    return table_.eraseKey(key);
+  template <typename K, typename BeforeDestroy>
+  EnableHeterogeneousErase<K, size_type> eraseInto(
+      K const& key,
+      BeforeDestroy&& beforeDestroy) {
+    return tableEraseKeyInto(key, beforeDestroy);
   }
 
   //// PUBLIC - Lookup
@@ -930,6 +980,7 @@ class F14VectorMapImpl : public F14BasicMap<MapPolicyWithDefaults<
   using typename Super::const_iterator;
   using typename Super::iterator;
   using typename Super::key_type;
+  using typename Super::mapped_type;
   using typename Super::value_type;
 
   F14VectorMapImpl() = default;
@@ -963,7 +1014,10 @@ class F14VectorMapImpl : public F14BasicMap<MapPolicyWithDefaults<
   }
 
  private:
-  void eraseUnderlying(typename Policy::ItemIter underlying) {
+  template <typename BeforeDestroy>
+  void eraseUnderlying(
+      typename Policy::ItemIter underlying,
+      BeforeDestroy&& beforeDestroy) {
     Alloc& a = this->table_.alloc();
     auto values = this->table_.values_;
 
@@ -971,7 +1025,7 @@ class F14VectorMapImpl : public F14BasicMap<MapPolicyWithDefaults<
     auto index = underlying.item();
     // The item still needs to be hashable during this call, so we must destroy
     // the value _afterwards_.
-    this->table_.eraseIter(underlying);
+    this->tableEraseIterInto(underlying, beforeDestroy);
     Policy::AllocTraits::destroy(a, std::addressof(values[index]));
 
     // move the last element in values_ down and fix up the inbound index
@@ -986,47 +1040,81 @@ class F14VectorMapImpl : public F14BasicMap<MapPolicyWithDefaults<
     }
   }
 
-  template <typename K>
-  std::size_t eraseUnderlyingKey(K const& key) {
+  template <typename K, typename BeforeDestroy>
+  std::size_t eraseUnderlyingKey(K const& key, BeforeDestroy&& beforeDestroy) {
     auto underlying = this->table_.find(key);
     if (underlying.atEnd()) {
       return 0;
     } else {
-      eraseUnderlying(underlying);
+      eraseUnderlying(underlying, beforeDestroy);
       return 1;
     }
   }
 
  public:
   FOLLY_ALWAYS_INLINE iterator erase(const_iterator pos) {
+    return eraseInto(pos, [](key_type&&, mapped_type&&) {});
+  }
+
+  // This form avoids ambiguity when key_type has a templated constructor
+  // that accepts const_iterator
+  FOLLY_ALWAYS_INLINE iterator erase(iterator pos) {
+    return eraseInto(pos, [](key_type&&, mapped_type&&) {});
+  }
+
+  iterator erase(const_iterator first, const_iterator last) {
+    return eraseInto(first, last, [](key_type&&, mapped_type&&) {});
+  }
+
+  std::size_t erase(key_type const& key) {
+    return eraseInto(key, [](key_type&&, mapped_type&&) {});
+  }
+
+  template <typename K>
+  EnableHeterogeneousVectorErase<K, std::size_t> erase(K const& key) {
+    return eraseInto(key, [](key_type&&, mapped_type&&) {});
+  }
+
+  template <typename BeforeDestroy>
+  FOLLY_ALWAYS_INLINE iterator
+  eraseInto(const_iterator pos, BeforeDestroy&& beforeDestroy) {
     auto index = this->table_.iterToIndex(pos);
     auto underlying = this->table_.find(VectorContainerIndexSearch{index});
-    eraseUnderlying(underlying);
+    eraseUnderlying(underlying, beforeDestroy);
     return index == 0 ? end() : this->table_.indexToIter(index - 1);
   }
 
   // This form avoids ambiguity when key_type has a templated constructor
   // that accepts const_iterator
-  FOLLY_ALWAYS_INLINE iterator erase(iterator pos) {
+  template <typename BeforeDestroy>
+  FOLLY_ALWAYS_INLINE iterator
+  eraseInto(iterator pos, BeforeDestroy&& beforeDestroy) {
     const_iterator cpos{pos};
-    return erase(cpos);
+    return eraseInto(cpos, beforeDestroy);
   }
 
-  iterator erase(const_iterator first, const_iterator last) {
+  template <typename BeforeDestroy>
+  iterator eraseInto(
+      const_iterator first,
+      const_iterator last,
+      BeforeDestroy&& beforeDestroy) {
     while (first != last) {
-      first = erase(first);
+      first = eraseInto(first, beforeDestroy);
     }
     auto index = this->table_.iterToIndex(first);
     return index == 0 ? end() : this->table_.indexToIter(index - 1);
   }
 
-  std::size_t erase(key_type const& key) {
-    return eraseUnderlyingKey(key);
+  template <typename BeforeDestroy>
+  std::size_t eraseInto(key_type const& key, BeforeDestroy&& beforeDestroy) {
+    return eraseUnderlyingKey(key, beforeDestroy);
   }
 
-  template <typename K>
-  EnableHeterogeneousVectorErase<K, std::size_t> erase(K const& key) {
-    return eraseUnderlyingKey(key);
+  template <typename K, typename BeforeDestroy>
+  EnableHeterogeneousVectorErase<K, std::size_t> eraseInto(
+      K const& key,
+      BeforeDestroy&& beforeDestroy) {
+    return eraseUnderlyingKey(key, beforeDestroy);
   }
 
   template <typename V>

folly/container/detail/F14Table.h

@@ -2153,13 +2153,6 @@ class F14Table : public Policy {
   }
 
  public:
-  // The item needs to still be hashable during this call.  If you want
-  // to intercept the value before it is destroyed (to extract it, for
-  // example), use eraseIterInto(pos, beforeDestroy).
-  void eraseIter(ItemIter pos) {
-    eraseIterInto(pos, [](value_type&&) {});
-  }
-
   // The item needs to still be hashable during this call.  If you want
   // to intercept the value before it is destroyed (to extract it, for
   // example), do so in the beforeDestroy callback.
@@ -2173,11 +2166,6 @@ class F14Table : public Policy {
     eraseImpl(pos, hp);
   }
 
-  template <typename K>
-  std::size_t eraseKey(K const& key) {
-    return eraseKeyInto(key, [](value_type&&) {});
-  }
-
   template <typename K, typename BeforeDestroy>
   std::size_t eraseKeyInto(K const& key, BeforeDestroy&& beforeDestroy) {
     if (UNLIKELY(size() == 0)) {

folly/container/test/F14MapTest.cpp

@@ -1446,6 +1446,72 @@ TEST(F14FastMap, moveOnly) {
       F14FastMap<folly::test::MoveOnlyTestInt, folly::test::MoveOnlyTestInt>>();
 }
 
+template <typename M>
+void runEraseIntoTest() {
+  M t0;
+  M t1;
+
+  auto emplaceIntoT0 = [&t0](auto&& key, auto&& value) {
+    EXPECT_FALSE(key.destroyed);
+    t0.emplace(std::move(key), std::move(value));
+  };
+  auto emplaceIntoT0Mut = [&](typename M::key_type&& key,
+                              typename M::mapped_type&& value) mutable {
+    emplaceIntoT0(std::move(key), std::move(value));
+  };
+
+  t0.emplace(10, 0);
+  t1.emplace(20, 0);
+  t1.eraseInto(t1.begin(), emplaceIntoT0);
+  EXPECT_TRUE(t1.empty());
+  EXPECT_EQ(t0.size(), 2);
+  EXPECT_TRUE(t0.find(10) != t0.end());
+  EXPECT_TRUE(t0.find(20) != t0.end());
+
+  t1.emplace(20, 0);
+  t1.emplace(30, 0);
+  t1.emplace(40, 0);
+  t1.eraseInto(t1.begin(), t1.end(), emplaceIntoT0Mut);
+  EXPECT_TRUE(t1.empty());
+  EXPECT_EQ(t0.size(), 4);
+  EXPECT_TRUE(t0.find(30) != t0.end());
+  EXPECT_TRUE(t0.find(40) != t0.end());
+
+  t1.emplace(50, 0);
+  size_t erased = t1.eraseInto(t1.find(50)->first, emplaceIntoT0);
+  EXPECT_EQ(erased, 1);
+  EXPECT_TRUE(t1.empty());
+  EXPECT_EQ(t0.size(), 5);
+  EXPECT_TRUE(t0.find(50) != t0.end());
+
+  typename M::key_type key{60};
+  erased = t1.eraseInto(key, emplaceIntoT0Mut);
+  EXPECT_EQ(erased, 0);
+  EXPECT_EQ(t0.size(), 5);
+}
+
+TEST(F14ValueMap, eraseInto) {
+  runEraseIntoTest<F14ValueMap<
+      folly::test::MoveOnlyTestInt,
+      folly::test::MoveOnlyTestInt>>();
+}
+
+TEST(F14NodeMap, eraseInto) {
+  runEraseIntoTest<
+      F14NodeMap<folly::test::MoveOnlyTestInt, folly::test::MoveOnlyTestInt>>();
+}
+
+TEST(F14VectorMap, eraseInto) {
+  runEraseIntoTest<F14VectorMap<
+      folly::test::MoveOnlyTestInt,
+      folly::test::MoveOnlyTestInt>>();
+}
+
+TEST(F14FastMap, eraseInto) {
+  runEraseIntoTest<
+      F14FastMap<folly::test::MoveOnlyTestInt, folly::test::MoveOnlyTestInt>>();
+}
+
 TEST(F14ValueMap, heterogeneousLookup) {
   using Hasher = folly::transparent<folly::hasher<folly::StringPiece>>;
   using KeyEqual = folly::transparent<std::equal_to<folly::StringPiece>>;
@@ -1617,6 +1683,13 @@ void runHeterogeneousInsertTest() {
   EXPECT_EQ(map.size(), 0);
   EXPECT_EQ(Tracked<1>::counts().dist(Counts{0, 0, 0, 0}), 0)
       << Tracked<1>::counts;
+
+  map.emplace(10, 40);
+  resetTracking();
+  map.eraseInto(10, [](auto&&, auto&&) {});
+  EXPECT_EQ(map.size(), 0);
+  EXPECT_EQ(Tracked<1>::counts().dist(Counts{0, 0, 0, 0}), 0)
+      << Tracked<1>::counts;
 }
 
 template <typename M>

folly/container/test/TrackingTypes.h

@@ -24,6 +24,7 @@
 #include <folly/Function.h>
 #include <folly/hash/Hash.h>
 #include <folly/lang/SafeAssert.h>
+#include <folly/portability/Asm.h>
 
 namespace folly {
 namespace test {
@@ -37,17 +38,21 @@ struct MoveOnlyTestInt {
   MoveOnlyTestInt(MoveOnlyTestInt&& rhs) noexcept : x(rhs.x) {}
   MoveOnlyTestInt(MoveOnlyTestInt const&) = delete;
   MoveOnlyTestInt& operator=(MoveOnlyTestInt&& rhs) noexcept {
+    FOLLY_SAFE_CHECK(!rhs.destroyed, "");
     x = rhs.x;
-    destroyed = rhs.destroyed;
     return *this;
   }
   MoveOnlyTestInt& operator=(MoveOnlyTestInt const&) = delete;
 
   ~MoveOnlyTestInt() {
+    FOLLY_SAFE_CHECK(!destroyed, "");
     destroyed = true;
+    asm_volatile_memory(); // try to keep compiler from eliding the store
   }
 
   bool operator==(MoveOnlyTestInt const& rhs) const {
+    FOLLY_SAFE_CHECK(!destroyed, "");
+    FOLLY_SAFE_CHECK(!rhs.destroyed, "");
     return x == rhs.x && destroyed == rhs.destroyed;
   }
   bool operator!=(MoveOnlyTestInt const& rhs) const {
@@ -538,6 +543,7 @@ namespace std {
 template <>
 struct hash<folly::test::MoveOnlyTestInt> {
   std::size_t operator()(folly::test::MoveOnlyTestInt const& val) const {
+    FOLLY_SAFE_CHECK(!val.destroyed, "");
     return val.x;
   }
 };