From 46c0f844304773328c7df6da7c400cd542835c04 Mon Sep 17 00:00:00 2001
From: Agustin Borgna <agustin.borgna@quantinuum.com>
Date: Wed, 20 Sep 2023 14:30:53 +0100
Subject: [PATCH] feat: Track max memory usage for TASO

---
 taso-optimiser/Cargo.toml  |  7 ++++++-
 taso-optimiser/src/main.rs | 10 ++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/taso-optimiser/Cargo.toml b/taso-optimiser/Cargo.toml
index eb862e14..787f5654 100644
--- a/taso-optimiser/Cargo.toml
+++ b/taso-optimiser/Cargo.toml
@@ -11,4 +11,9 @@ serde_json = "1.0"
 tket2 = { path = "../", features = ["portmatching"] }
 quantinuum-hugr = { workspace = true }
 itertools = { workspace = true }
-tket-json-rs = { workspace = true }
\ No newline at end of file
+tket-json-rs = { workspace = true }
+peak_alloc = { version = "0.2.0", optional = true }
+
+[features]
+default = ["peak_alloc"]
+peak_alloc = ["dep:peak_alloc"]
\ No newline at end of file
diff --git a/taso-optimiser/src/main.rs b/taso-optimiser/src/main.rs
index f73bd12d..09b022a1 100644
--- a/taso-optimiser/src/main.rs
+++ b/taso-optimiser/src/main.rs
@@ -8,6 +8,13 @@ use tket2::{
 };
 use tket_json_rs::circuit_json::SerialCircuit;
 
+#[cfg(feature = "peak_alloc")]
+use peak_alloc::PeakAlloc;
+
+#[cfg(feature = "peak_alloc")]
+#[global_allocator]
+static PEAK_ALLOC: PeakAlloc = PeakAlloc;
+
 /// Optimise circuits using Quartz-generated ECCs.
 ///
 /// Quartz: <https://github.com/quantum-compiler/quartz>
@@ -91,5 +98,8 @@ fn main() {
     println!("Saving result");
     save_tk1_json_file(output_path, &opt_circ).unwrap();
 
+    #[cfg(feature = "peak_alloc")]
+    println!("Peak memory usage: {} GB", PEAK_ALLOC.peak_usage_as_gb());
+
     println!("Done.")
 }