1 files changed, 35 insertions, 6 deletions
diff --git a/absl/container/flat_hash_map_test.cc b/absl/container/flat_hash_map_test.cc
index 89ec60c9..263951f1 100644
--- a/absl/container/flat_hash_map_test.cc
+++ b/absl/container/flat_hash_map_test.cc
@@ -236,33 +236,36 @@ TEST(FlatHashMap, EraseIf) {
   // Erase all elements.
   {
     flat_hash_map<int, int> s = {{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}};
-    erase_if(s, [](std::pair<const int, int>) { return true; });
+    EXPECT_EQ(erase_if(s, [](std::pair<const int, int>) { return true; }), 5);
     EXPECT_THAT(s, IsEmpty());
   }
   // Erase no elements.
   {
     flat_hash_map<int, int> s = {{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}};
-    erase_if(s, [](std::pair<const int, int>) { return false; });
+    EXPECT_EQ(erase_if(s, [](std::pair<const int, int>) { return false; }), 0);
     EXPECT_THAT(s, UnorderedElementsAre(Pair(1, 1), Pair(2, 2), Pair(3, 3),
                                         Pair(4, 4), Pair(5, 5)));
   }
   // Erase specific elements.
   {
     flat_hash_map<int, int> s = {{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}};
-    erase_if(s,
-             [](std::pair<const int, int> kvp) { return kvp.first % 2 == 1; });
+    EXPECT_EQ(erase_if(s,
+                       [](std::pair<const int, int> kvp) {
+                         return kvp.first % 2 == 1;
+                       }),
+              3);
     EXPECT_THAT(s, UnorderedElementsAre(Pair(2, 2), Pair(4, 4)));
   }
   // Predicate is function reference.
   {
     flat_hash_map<int, int> s = {{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}};
-    erase_if(s, FirstIsEven);
+    EXPECT_EQ(erase_if(s, FirstIsEven), 2);
     EXPECT_THAT(s, UnorderedElementsAre(Pair(1, 1), Pair(3, 3), Pair(5, 5)));
   }
   // Predicate is function pointer.
   {
     flat_hash_map<int, int> s = {{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}};
-    erase_if(s, &FirstIsEven);
+    EXPECT_EQ(erase_if(s, &FirstIsEven), 2);
     EXPECT_THAT(s, UnorderedElementsAre(Pair(1, 1), Pair(3, 3), Pair(5, 5)));
   }
 }
@@ -282,6 +285,32 @@ TEST(FlatHashMap, NodeHandleMutableKeyAccess) {
 }
 #endif
 
+TEST(FlatHashMap, Reserve) {
+  // Verify that if we reserve(size() + n) then we can perform n insertions
+  // without a rehash, i.e., without invalidating any references.
+  for (size_t trial = 0; trial < 20; ++trial) {
+    for (size_t initial = 3; initial < 100; ++initial) {
+      // Fill in `initial` entries, then erase 2 of them, then reserve space for
+      // two inserts and check for reference stability while doing the inserts.
+      flat_hash_map<size_t, size_t> map;
+      for (size_t i = 0; i < initial; ++i) {
+        map[i] = i;
+      }
+      map.erase(0);
+      map.erase(1);
+      map.reserve(map.size() + 2);
+      size_t& a2 = map[2];
+      // In the event of a failure, asan will complain in one of these two
+      // assignments.
+      map[initial] = a2;
+      map[initial + 1] = a2;
+      // Fail even when not under asan:
+      size_t& a2new = map[2];
+      EXPECT_EQ(&a2, &a2new);
+    }
+  }
+}
+
 }  // namespace
 }  // namespace container_internal
 ABSL_NAMESPACE_END