Add `BM_EraseIf` benchmark.

PiperOrigin-RevId: 618970135
Change-Id: Ifb9d0b425904d5cb37d80ec28ab7845957209313

2 files changed, 61 insertions, 0 deletions

diff --git a/absl/container/BUILD.bazel b/absl/container/BUILD.bazel
index 78ad60e0..12e27c21 100644
--- a/absl/container/BUILD.bazel
+++ b/absl/container/BUILD.bazel
@@ -726,6 +726,7 @@ cc_binary(
         ":hash_function_defaults",
         ":raw_hash_set",
         "//absl/base:raw_logging_internal",
+        "//absl/random",
         "//absl/strings:str_format",
         "@com_github_google_benchmark//:benchmark_main",
     ],
diff --git a/absl/container/internal/raw_hash_set_benchmark.cc b/absl/container/internal/raw_hash_set_benchmark.cc
index 0fa9c712..05f62d2f 100644
--- a/absl/container/internal/raw_hash_set_benchmark.cc
+++ b/absl/container/internal/raw_hash_set_benchmark.cc
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <algorithm>
 #include <array>
 #include <cmath>
 #include <cstddef>
@@ -27,6 +28,7 @@
 #include "absl/container/internal/container_memory.h"
 #include "absl/container/internal/hash_function_defaults.h"
 #include "absl/container/internal/raw_hash_set.h"
+#include "absl/random/random.h"
 #include "absl/strings/str_format.h"
 #include "benchmark/benchmark.h"
 
@@ -575,6 +577,64 @@ void BM_Resize(benchmark::State& state) {
 }
 BENCHMARK(BM_Resize);
 
+void BM_EraseIf(benchmark::State& state) {
+  int64_t num_elements = state.range(0);
+  size_t num_erased = static_cast<size_t>(state.range(1));
+
+  constexpr size_t kRepetitions = 64;
+
+  absl::BitGen rng;
+
+  std::vector<std::vector<int64_t>> keys(kRepetitions);
+  std::vector<IntTable> tables;
+  std::vector<int64_t> threshold;
+  for (auto& k : keys) {
+    tables.push_back(IntTable());
+    auto& table = tables.back();
+    for (int64_t i = 0; i < num_elements; i++) {
+      // We use random keys to reduce noise.
+      k.push_back(absl::Uniform<int64_t>(rng, 0, ~int64_t{}));
+      if (!table.insert(k.back()).second) {
+        k.pop_back();
+        --i;  // duplicated value, retrying
+      }
+    }
+    std::sort(k.begin(), k.end());
+    threshold.push_back(k[num_erased]);
+  }
+
+  while (state.KeepRunningBatch(static_cast<int64_t>(kRepetitions) *
+                                std::max(num_elements, int64_t{1}))) {
+    benchmark::DoNotOptimize(tables);
+    for (size_t t_id = 0; t_id < kRepetitions; t_id++) {
+      auto& table = tables[t_id];
+      benchmark::DoNotOptimize(num_erased);
+      auto pred = [t = threshold[t_id]](int64_t key) { return key < t; };
+      benchmark::DoNotOptimize(pred);
+      benchmark::DoNotOptimize(table);
+      absl::container_internal::EraseIf(pred, &table);
+    }
+    state.PauseTiming();
+    for (size_t t_id = 0; t_id < kRepetitions; t_id++) {
+      auto& k = keys[t_id];
+      auto& table = tables[t_id];
+      for (size_t i = 0; i < num_erased; i++) {
+        table.insert(k[i]);
+      }
+    }
+    state.ResumeTiming();
+  }
+}
+
+BENCHMARK(BM_EraseIf)
+    ->ArgNames({"num_elements", "num_erased"})
+    ->ArgPair(10, 0)
+    ->ArgPair(1000, 0)
+    ->ArgPair(10, 5)
+    ->ArgPair(1000, 500)
+    ->ArgPair(10, 10)
+    ->ArgPair(1000, 1000);
+
 }  // namespace
 }  // namespace container_internal
 ABSL_NAMESPACE_END