From af6abcb3d0097a7f7892fb8b54a4c6363e5c2c7f Mon Sep 17 00:00:00 2001
From: Pieter Wuille <pieter@wuille.net>
Date: Fri, 5 Nov 2021 17:48:18 -0400
Subject: [PATCH] Make bench support selecting which benchmarks to run

---
 src/bench.c                         | 13 +++++++------
 src/modules/ecdh/bench_impl.h       |  5 +++--
 src/modules/recovery/bench_impl.h   |  5 +++--
 src/modules/schnorrsig/bench_impl.h |  7 ++++---
 4 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/src/bench.c b/src/bench.c
index f46d5d60..dfc0470b 100644
--- a/src/bench.c
+++ b/src/bench.c
@@ -88,12 +88,13 @@ static void bench_sign_run(void* arg, int iters) {
 # include "modules/schnorrsig/bench_impl.h"
 #endif
 
-int main(void) {
+int main(int argc, char** argv) {
     int i;
     secp256k1_pubkey pubkey;
     secp256k1_ecdsa_signature sig;
     bench_verify_data data;
 
+    int d = argc == 1;
     int iters = get_iters(20000);
 
     /* ECDSA verification benchmark */
@@ -113,30 +114,30 @@ int main(void) {
     CHECK(secp256k1_ec_pubkey_serialize(data.ctx, data.pubkey, &data.pubkeylen, &pubkey, SECP256K1_EC_COMPRESSED) == 1);
 
     print_output_table_header_row();
-    run_benchmark("ecdsa_verify", bench_verify, NULL, NULL, &data, 10, iters);
+    if (d || have_flag(argc, argv, "ecdsa") || have_flag(argc, argv, "verify") || have_flag(argc, argv, "ecdsa_verify")) run_benchmark("ecdsa_verify", bench_verify, NULL, NULL, &data, 10, iters);
 
     secp256k1_context_destroy(data.ctx);
 
     /* ECDSA signing benchmark */
     data.ctx = secp256k1_context_create(SECP256K1_CONTEXT_SIGN);
 
-    run_benchmark("ecdsa_sign", bench_sign_run, bench_sign_setup, NULL, &data, 10, iters);
+    if (d || have_flag(argc, argv, "ecdsa") || have_flag(argc, argv, "sign") || have_flag(argc, argv, "ecdsa_sign")) run_benchmark("ecdsa_sign", bench_sign_run, bench_sign_setup, NULL, &data, 10, iters);
 
     secp256k1_context_destroy(data.ctx);
 
 #ifdef ENABLE_MODULE_ECDH
     /* ECDH benchmarks */
-    run_ecdh_bench(iters);
+    run_ecdh_bench(iters, argc, argv);
 #endif
 
 #ifdef ENABLE_MODULE_RECOVERY
     /* ECDSA recovery benchmarks */
-    run_recovery_bench(iters);
+    run_recovery_bench(iters, argc, argv);
 #endif
 
 #ifdef ENABLE_MODULE_SCHNORRSIG
     /* Schnorr signature benchmarks */
-    run_schnorrsig_bench(iters);
+    run_schnorrsig_bench(iters, argc, argv);
 #endif
 
     return 0;
diff --git a/src/modules/ecdh/bench_impl.h b/src/modules/ecdh/bench_impl.h
index 0b8b74e2..94d83346 100644
--- a/src/modules/ecdh/bench_impl.h
+++ b/src/modules/ecdh/bench_impl.h
@@ -42,13 +42,14 @@ static void bench_ecdh(void* arg, int iters) {
     }
 }
 
-void run_ecdh_bench(int iters) {
+void run_ecdh_bench(int iters, int argc, char** argv) {
     bench_ecdh_data data;
+    int d = argc == 1;
 
     /* create a context with no capabilities */
     data.ctx = secp256k1_context_create(SECP256K1_FLAGS_TYPE_CONTEXT);
 
-    run_benchmark("ecdh", bench_ecdh, bench_ecdh_setup, NULL, &data, 10, iters);
+    if (d || have_flag(argc, argv, "ecdh")) run_benchmark("ecdh", bench_ecdh, bench_ecdh_setup, NULL, &data, 10, iters);
 
     secp256k1_context_destroy(data.ctx);
 }
diff --git a/src/modules/recovery/bench_impl.h b/src/modules/recovery/bench_impl.h
index 95330b1e..4a9e8869 100644
--- a/src/modules/recovery/bench_impl.h
+++ b/src/modules/recovery/bench_impl.h
@@ -48,12 +48,13 @@ void bench_recover_setup(void* arg) {
     }
 }
 
-void run_recovery_bench(int iters) {
+void run_recovery_bench(int iters, int argc, char** argv) {
     bench_recover_data data;
+    int d = argc == 1;
 
     data.ctx = secp256k1_context_create(SECP256K1_CONTEXT_VERIFY);
 
-    run_benchmark("ecdsa_recover", bench_recover, bench_recover_setup, NULL, &data, 10, iters);
+    if (d || have_flag(argc, argv, "ecdsa") || have_flag(argc, argv, "recover") || have_flag(argc, argv, "ecdsa_recover")) run_benchmark("ecdsa_recover", bench_recover, bench_recover_setup, NULL, &data, 10, iters);
 
     secp256k1_context_destroy(data.ctx);
 }
diff --git a/src/modules/schnorrsig/bench_impl.h b/src/modules/schnorrsig/bench_impl.h
index 2f2c8db0..41f393c8 100644
--- a/src/modules/schnorrsig/bench_impl.h
+++ b/src/modules/schnorrsig/bench_impl.h
@@ -45,9 +45,10 @@ void bench_schnorrsig_verify(void* arg, int iters) {
     }
 }
 
-void run_schnorrsig_bench(int iters) {
+void run_schnorrsig_bench(int iters, int argc, char** argv) {
     int i;
     bench_schnorrsig_data data;
+    int d = argc == 1;
 
     data.ctx = secp256k1_context_create(SECP256K1_CONTEXT_VERIFY | SECP256K1_CONTEXT_SIGN);
     data.keypairs = (const secp256k1_keypair **)malloc(iters * sizeof(secp256k1_keypair *));
@@ -81,8 +82,8 @@ void run_schnorrsig_bench(int iters) {
         CHECK(secp256k1_xonly_pubkey_serialize(data.ctx, pk_char, &pk) == 1);
     }
 
-    run_benchmark("schnorrsig_sign", bench_schnorrsig_sign, NULL, NULL, (void *) &data, 10, iters);
-    run_benchmark("schnorrsig_verify", bench_schnorrsig_verify, NULL, NULL, (void *) &data, 10, iters);
+    if (d || have_flag(argc, argv, "schnorrsig") || have_flag(argc, argv, "sign") || have_flag(argc, argv, "schnorrsig_sign")) run_benchmark("schnorrsig_sign", bench_schnorrsig_sign, NULL, NULL, (void *) &data, 10, iters);
+    if (d || have_flag(argc, argv, "schnorrsig") || have_flag(argc, argv, "verify") || have_flag(argc, argv, "schnorrsig_verify")) run_benchmark("schnorrsig_verify", bench_schnorrsig_verify, NULL, NULL, (void *) &data, 10, iters);
 
     for (i = 0; i < iters; i++) {
         free((void *)data.keypairs[i]);