From f5c01ee48750218ff6ed6edc63d99105ad6ccd86 Mon Sep 17 00:00:00 2001
From: Chandler Carruth <chandlerc@gmail.com>
Date: Tue, 9 Apr 2024 13:36:41 -0700
Subject: [PATCH] Add a custom main library for benchmarks. (#3872)

This largely reproduces the upstream one, but has a few advantages:

1) It initializes LLVM which can be important if we use the backends.
2) It parses commandline flags with Abseil allowing the use of Abseil
   flags in benchmarks.
3) It flips the default for tabular display of results which we use
   pretty heavily.

---------

Co-authored-by: Jon Ross-Perkins <jperkins@google.com>
Co-authored-by: Richard Smith <richard@metafoo.co.uk>
Co-authored-by: Carbon Infra Bot <carbon-external-infra@google.com>
---
 common/BUILD              | 14 +++++++++++++-
 common/benchmark_main.cpp | 31 +++++++++++++++++++++++++++++++
 toolchain/lex/BUILD       |  9 ++++++---
 3 files changed, 50 insertions(+), 4 deletions(-)
 create mode 100644 common/benchmark_main.cpp

diff --git a/common/BUILD b/common/BUILD
index 645bf9a11e23..fcf314a6bc89 100644
--- a/common/BUILD
+++ b/common/BUILD
@@ -14,6 +14,17 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "benchmark_main",
+    srcs = ["benchmark_main.cpp"],
+    deps = [
+        ":init_llvm",
+        "@abseil-cpp//absl/flags:parse",
+        "@google_benchmark//:benchmark",
+        "@llvm-project//llvm:Support",
+    ],
+)
+
 cc_library(
     name = "command_line",
     srcs = ["command_line.cpp"],
@@ -140,11 +151,12 @@ cc_binary(
     testonly = 1,
     srcs = ["hashing_benchmark.cpp"],
     deps = [
+        ":benchmark_main",
         ":check",
         ":hashing",
         "@abseil-cpp//absl/hash",
         "@abseil-cpp//absl/random",
-        "@google_benchmark//:benchmark_main",
+        "@google_benchmark//:benchmark",
         "@llvm-project//llvm:Support",
     ],
 )
diff --git a/common/benchmark_main.cpp b/common/benchmark_main.cpp
new file mode 100644
index 000000000000..0b62972ce3cd
--- /dev/null
+++ b/common/benchmark_main.cpp
@@ -0,0 +1,31 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <benchmark/benchmark.h>
+
+#include "absl/flags/parse.h"
+#include "common/init_llvm.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+
+auto main(int orig_argc, char** orig_argv) -> int {
+  // Do LLVM's initialization first, this will also transform UTF-16 to UTF-8.
+  Carbon::InitLLVM init_llvm(orig_argc, orig_argv);
+
+  // Inject a flag to override the defaults for benchmarks. This can still be
+  // disabled by user arguments.
+  llvm::SmallVector<char*> injected_argv_storage(orig_argv,
+                                                 orig_argv + orig_argc + 1);
+  char injected_flag[] = "--benchmark_counters_tabular";
+  injected_argv_storage.insert(injected_argv_storage.begin() + 1,
+                               injected_flag);
+  char** argv = injected_argv_storage.data();
+  int argc = injected_argv_storage.size() - 1;
+
+  benchmark::Initialize(&argc, argv);
+  absl::ParseCommandLine(argc, argv);
+  benchmark::RunSpecifiedBenchmarks();
+  benchmark::Shutdown();
+  return 0;
+}
diff --git a/toolchain/lex/BUILD b/toolchain/lex/BUILD
index 0e215168d8d6..73a728972742 100644
--- a/toolchain/lex/BUILD
+++ b/toolchain/lex/BUILD
@@ -85,9 +85,10 @@ cc_binary(
     srcs = ["numeric_literal_benchmark.cpp"],
     deps = [
         ":numeric_literal",
+        "//common:benchmark_main",
         "//common:check",
         "//toolchain/diagnostics:null_diagnostics",
-        "@google_benchmark//:benchmark_main",
+        "@google_benchmark//:benchmark",
     ],
 )
 
@@ -139,8 +140,9 @@ cc_binary(
     srcs = ["string_literal_benchmark.cpp"],
     deps = [
         ":string_literal",
+        "//common:benchmark_main",
         "//toolchain/diagnostics:null_diagnostics",
-        "@google_benchmark//:benchmark_main",
+        "@google_benchmark//:benchmark",
     ],
 )
 
@@ -271,12 +273,13 @@ cc_binary(
         ":lex",
         ":token_kind",
         ":tokenized_buffer",
+        "//common:benchmark_main",
         "//common:check",
         "//toolchain/base:value_store",
         "//toolchain/diagnostics:diagnostic_emitter",
         "//toolchain/diagnostics:null_diagnostics",
         "@abseil-cpp//absl/random",
-        "@google_benchmark//:benchmark_main",
+        "@google_benchmark//:benchmark",
         "@llvm-project//llvm:Support",
     ],
 )