From b0900a42403bf3079821f70a136fe662b5f30ecc Mon Sep 17 00:00:00 2001
From: sveitser <sveitser@gmail.com>
Date: Wed, 18 Dec 2024 13:18:53 +0100
Subject: [PATCH] Add script to benchmark compilation times

---
 .gitignore                   |   4 ++
 flake.nix                    |   2 +-
 scripts/bench-compilation.py | 103 +++++++++++++++++++++++++++++++++++
 3 files changed, 108 insertions(+), 1 deletion(-)
 create mode 100644 scripts/bench-compilation.py

diff --git a/.gitignore b/.gitignore
index a51534fddc..e8482d078b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -54,3 +54,7 @@ contracts/broadcast/*/11155111/
 
 # Contract documentation generated by `forge doc`
 docs/
+
+# cargo compilation benchmark
+bench-compilation-results.toml
+bench-compilation.log
diff --git a/flake.nix b/flake.nix
index 42e9a06f7d..917e85bb0b 100644
--- a/flake.nix
+++ b/flake.nix
@@ -218,7 +218,7 @@
             solc
             nodePackages.prettier
             solhint
-            (python3.withPackages (ps: with ps; [ black ]))
+            (python3.withPackages (ps: with ps; [ black tomlkit ]))
             yarn
           ] ++ lib.optionals stdenv.isDarwin
             [ darwin.apple_sdk.frameworks.SystemConfiguration ]
diff --git a/scripts/bench-compilation.py b/scripts/bench-compilation.py
new file mode 100644
index 0000000000..abe98a1791
--- /dev/null
+++ b/scripts/bench-compilation.py
@@ -0,0 +1,103 @@
+#!/usr/bin/env python
+import argparse
+import subprocess
+import time
+import tomlkit
+
+RESULTS = "bench-compilation-results.toml"
+
+
+def cached(func):
+    """A decorator to cache function results to a file"""
+
+    def wrapper(key, *args, **kwargs):
+        try:
+            with open(RESULTS) as f:
+                data = tomlkit.load(f)
+        except (FileNotFoundError,):
+            print(f"No cache at {RESULTS}, making a new one.")
+            data = {}
+        if key in data:
+            return data[key]
+        ret = func(key, *args, **kwargs)
+        data[key] = ret
+        with open(RESULTS, "w") as f:
+            tomlkit.dump(data, f)
+        return ret
+
+    return wrapper
+
+
+@cached
+def benchmark(key: str, setup_command: str, timed_command: str):
+    print(f"Running setup command: {setup_command}")
+    with open("bench-compilation.log", "a") as f:
+        subprocess.run(
+            setup_command,
+            shell=True,
+            stdout=f,
+            stderr=subprocess.STDOUT,
+        )
+        print(f"Running timed command: {timed_command}")
+        tic = time.time()
+        subprocess.run(
+            timed_command,
+            shell=True,
+            stdout=f,
+            stderr=subprocess.STDOUT,
+        )
+        toc = time.time()
+        return toc - tic
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description=f"""Cargo compilation benchmark
+
+This script benchmarks the compilation time of a Rust project using
+different profiles. It runs the setup command (default: "cargo clean")
+before each timed command (default: "cargo build").
+
+Results are cached into TOML file in {RESULTS} where keys are profiles and values
+are the seconds it took to run the timed command.
+
+The script generates new cargo profile for all combinations of the "strip" and
+"debug" values and runs the timed command using these profiles.""",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--timed", help="Timed command to run", default="cargo build")
+    parser.add_argument("--setup", help="Setup command to run", default="cargo clean")
+    parser.add_argument("--strip", help="Strip values", default="none,debuginfo")
+    parser.add_argument(
+        "--debug", help="Debug values", default="none,line-tables-only,limited,full"
+    )
+    args = parser.parse_args()
+
+    with open("Cargo.toml") as f:
+        cargo_toml = tomlkit.load(f)
+    profile_section = cargo_toml["profile"]
+    profiles = []
+    for strip in args.strip.split(","):
+        for debug in args.debug.split(","):
+            profile = f"debug-{debug}-strip-{strip}"
+            profiles.append(profile)
+            profile_section[profile] = {
+                "inherits": "test",
+                "debug": debug,
+                "strip": strip,
+            }
+
+    # Write the new profiles to Cargo.toml, tomlkit preserves the existing style.
+    with open("Cargo.toml", "w") as f:
+        tomlkit.dump(cargo_toml, f)
+
+    # Run the benchmarks
+    for profile in profiles:
+        command = f"{args.timed} --profile {profile}"
+        print(f"Testing profile {profile}")
+        time = benchmark(profile, args.setup, command)
+        print(f"Time: {time:.2f}s")
+
+
+if __name__ == "__main__":
+    main()