3 files changed, 98 insertions, 0 deletions
diff --git a/absl/container/btree_map.h b/absl/container/btree_map.h
index 470e3197084c..cbfcb58c4129 100644
--- a/absl/container/btree_map.h
+++ b/absl/container/btree_map.h
@@ -412,6 +412,20 @@ void swap(btree_map<K, V, C, A> &x, btree_map<K, V, C, A> &y) {
   return x.swap(y);
 }
 
+// absl::erase_if(absl::btree_map<>, Pred)
+//
+// Erases all elements that satisfy the predicate pred from the container.
+template <typename K, typename V, typename C, typename A, typename Pred>
+void erase_if(btree_map<K, V, C, A> &map, Pred pred) {
+  for (auto it = map.begin(); it != map.end();) {
+    if (pred(*it)) {
+      it = map.erase(it);
+    } else {
+      ++it;
+    }
+  }
+}
+
 // absl::btree_multimap
 //
 // An `absl::btree_multimap<K, V>` is an ordered associative container of
@@ -701,6 +715,20 @@ void swap(btree_multimap<K, V, C, A> &x, btree_multimap<K, V, C, A> &y) {
   return x.swap(y);
 }
 
+// absl::erase_if(absl::btree_multimap<>, Pred)
+//
+// Erases all elements that satisfy the predicate pred from the container.
+template <typename K, typename V, typename C, typename A, typename Pred>
+void erase_if(btree_multimap<K, V, C, A> &map, Pred pred) {
+  for (auto it = map.begin(); it != map.end();) {
+    if (pred(*it)) {
+      it = map.erase(it);
+    } else {
+      ++it;
+    }
+  }
+}
+
 ABSL_NAMESPACE_END
 }  // namespace absl
 
diff --git a/absl/container/btree_set.h b/absl/container/btree_set.h
index 2a4d0aceb354..127fb940d40e 100644
--- a/absl/container/btree_set.h
+++ b/absl/container/btree_set.h
@@ -360,6 +360,20 @@ void swap(btree_set<K, C, A> &x, btree_set<K, C, A> &y) {
   return x.swap(y);
 }
 
+// absl::erase_if(absl::btree_set<>, Pred)
+//
+// Erases all elements that satisfy the predicate pred from the container.
+template <typename K, typename C, typename A, typename Pred>
+void erase_if(btree_set<K, C, A> &set, Pred pred) {
+  for (auto it = set.begin(); it != set.end();) {
+    if (pred(*it)) {
+      it = set.erase(it);
+    } else {
+      ++it;
+    }
+  }
+}
+
 // absl::btree_multiset<>
 //
 // An `absl::btree_multiset<K>` is an ordered associative container of
@@ -649,6 +663,20 @@ void swap(btree_multiset<K, C, A> &x, btree_multiset<K, C, A> &y) {
   return x.swap(y);
 }
 
+// absl::erase_if(absl::btree_multiset<>, Pred)
+//
+// Erases all elements that satisfy the predicate pred from the container.
+template <typename K, typename C, typename A, typename Pred>
+void erase_if(btree_multiset<K, C, A> &set, Pred pred) {
+  for (auto it = set.begin(); it != set.end();) {
+    if (pred(*it)) {
+      it = set.erase(it);
+    } else {
+      ++it;
+    }
+  }
+}
+
 ABSL_NAMESPACE_END
 }  // namespace absl
 
diff --git a/absl/container/btree_test.cc b/absl/container/btree_test.cc
index f8aadd62ad75..8692b9c2b9d7 100644
--- a/absl/container/btree_test.cc
+++ b/absl/container/btree_test.cc
@@ -51,6 +51,7 @@ using ::absl::test_internal::InstanceTracker;
 using ::absl::test_internal::MovableOnlyInstance;
 using ::testing::ElementsAre;
 using ::testing::ElementsAreArray;
+using ::testing::IsEmpty;
 using ::testing::Pair;
 
 template <typename T, typename U>
@@ -2303,6 +2304,47 @@ TEST(Btree, EmptyTree) {
   EXPECT_GT(s.max_size(), 0);
 }
 
+bool IsEven(int k) { return k % 2 == 0; }
+
+TEST(Btree, EraseIf) {
+  // Test that erase_if works with all the container types and supports lambdas.
+  {
+    absl::btree_set<int> s = {1, 3, 5, 6, 100};
+    erase_if(s, [](int k) { return k > 3; });
+    EXPECT_THAT(s, ElementsAre(1, 3));
+  }
+  {
+    absl::btree_multiset<int> s = {1, 3, 3, 5, 6, 6, 100};
+    erase_if(s, [](int k) { return k <= 3; });
+    EXPECT_THAT(s, ElementsAre(5, 6, 6, 100));
+  }
+  {
+    absl::btree_map<int, int> m = {{1, 1}, {3, 3}, {6, 6}, {100, 100}};
+    erase_if(m, [](std::pair<const int, int> kv) { return kv.first > 3; });
+    EXPECT_THAT(m, ElementsAre(Pair(1, 1), Pair(3, 3)));
+  }
+  {
+    absl::btree_multimap<int, int> m = {{1, 1}, {3, 3}, {3, 6},
+                                        {6, 6}, {6, 7}, {100, 6}};
+    erase_if(m, [](std::pair<const int, int> kv) { return kv.second == 6; });
+    EXPECT_THAT(m, ElementsAre(Pair(1, 1), Pair(3, 3), Pair(6, 7)));
+  }
+  // Test that erasing all elements from a large set works and test support for
+  // function pointers.
+  {
+    absl::btree_set<int> s;
+    for (int i = 0; i < 1000; ++i) s.insert(2 * i);
+    erase_if(s, IsEven);
+    EXPECT_THAT(s, IsEmpty());
+  }
+  // Test that erase_if supports other format of function pointers.
+  {
+    absl::btree_set<int> s = {1, 3, 5, 6, 100};
+    erase_if(s, &IsEven);
+    EXPECT_THAT(s, ElementsAre(1, 3, 5));
+  }
+}
+
 }  // namespace
 }  // namespace container_internal
 ABSL_NAMESPACE_END