From d36642cf5d71109e8f4eaf7aead2aae4494bc3d3 Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Tue, 30 Dec 2025 03:37:04 +0000
Subject: [PATCH] test: implement profiling tools

---
 Cargo.toml                              |  4 +++
 symbolic_regression/examples/example.rs | 33 +++++++++++++++++++-
 tools/profile.sh                        | 41 +++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100755 tools/profile.sh

diff --git a/Cargo.toml b/Cargo.toml
index 1b7fc8f..2c84286 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,3 +5,7 @@ members = [
     "symbolic_regression",
     "web/wasm",
 ]
+
+[profile.profiling]
+inherits = "release"
+debug = true
diff --git a/symbolic_regression/examples/example.rs b/symbolic_regression/examples/example.rs
index 771581b..1f8b81b 100644
--- a/symbolic_regression/examples/example.rs
+++ b/symbolic_regression/examples/example.rs
@@ -6,6 +6,37 @@ use symbolic_regression::prelude::*;
 // Mirrors `SymbolicRegression.jl/example.jl`.
 
 fn main() {
+    fn usage() {
+        eprintln!("Usage: example [--niterations <n>]");
+    }
+
+    let mut args = std::env::args().skip(1);
+    let mut niterations: usize = 1_000;
+    while let Some(arg) = args.next() {
+        match arg.as_str() {
+            "-h" | "--help" => {
+                usage();
+                return;
+            }
+            "-n" | "--niterations" => {
+                let value = args.next().unwrap_or_else(|| {
+                    eprintln!("Missing value for {arg}");
+                    usage();
+                    std::process::exit(2);
+                });
+                niterations = value.parse().unwrap_or_else(|_| {
+                    eprintln!("Expected `--niterations` as an integer, got: {value}");
+                    std::process::exit(2);
+                });
+            }
+            _ => {
+                eprintln!("Unknown arg: {arg}");
+                usage();
+                std::process::exit(2);
+            }
+        }
+    }
+
     const N_FEATURES: usize = 5;
     const D: usize = 3;
     let n_rows = 100;
@@ -30,7 +61,7 @@ fn main() {
 
     let options = Options::<f32, D> {
         operators,
-        niterations: 200,
+        niterations,
         ..Default::default()
     };
 
diff --git a/tools/profile.sh b/tools/profile.sh
new file mode 100755
index 0000000..633fedc
--- /dev/null
+++ b/tools/profile.sh
@@ -0,0 +1,41 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+usage() {
+  cat <<'EOF' >&2
+Usage: ./tools/profile.sh
+
+Profiles `symbolic_regression/examples/example.rs` (built with the `profiling` Cargo profile) using Samply.
+EOF
+}
+
+if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
+  usage
+  exit 0
+fi
+if [[ $# -ne 0 ]]; then
+  echo "No args supported." >&2
+  usage
+  exit 2
+fi
+
+root_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$root_dir"
+
+export RUSTFLAGS="${RUSTFLAGS:-} -C force-frame-pointers=yes"
+cargo --config 'build.rustc-wrapper=""' build -p symbolic_regression --profile profiling --example example
+
+bin="target/profiling/examples/example"
+
+mkdir -p profiles
+timestamp="$(date +%Y%m%d-%H%M%S)"
+
+if [[ "$(uname -s)" == "Darwin" ]]; then
+  samply setup --yes >/dev/null 2>&1 || {
+    echo "Warning: \`samply setup\` failed; profiling may fail on macOS." >&2
+  }
+fi
+
+out="profiles/example-${timestamp}.json.gz"
+echo "Writing profile to: $out" >&2
+exec samply record --rate 1000 --symbol-dir "target/profiling/examples" -o "$out" -- "$bin"