From ac6c582330f61824e75ff711032943e0ee097417 Mon Sep 17 00:00:00 2001
From: Agustin Borgna <agustin.borgna@quantinuum.com>
Date: Mon, 2 Oct 2023 22:54:57 +0100
Subject: [PATCH] chore: Small code cleanups

- Remove duplicated tk1 writer fn
- Better .dot debug API
---
 src/json.rs                |  9 ++++++---
 src/utils.rs               | 13 +++++++++++--
 taso-optimiser/src/main.rs | 14 ++------------
 3 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/src/json.rs b/src/json.rs
index 37ddbb593..94cdd8a1d 100644
--- a/src/json.rs
+++ b/src/json.rs
@@ -123,21 +123,24 @@ pub fn load_tk1_json_str(json: &str) -> Result<Hugr, TK1ConvertError> {
 }
 
 /// Save a circuit to file in TK1 JSON format.
-pub fn save_tk1_json_file(path: impl AsRef<Path>, circ: &Hugr) -> Result<(), TK1ConvertError> {
+pub fn save_tk1_json_file(
+    circ: &impl Circuit,
+    path: impl AsRef<Path>,
+) -> Result<(), TK1ConvertError> {
     let file = fs::File::create(path)?;
     let writer = io::BufWriter::new(file);
     save_tk1_json_writer(circ, writer)
 }
 
 /// Save a circuit in TK1 JSON format to a writer.
-pub fn save_tk1_json_writer(circ: &Hugr, w: impl io::Write) -> Result<(), TK1ConvertError> {
+pub fn save_tk1_json_writer(circ: &impl Circuit, w: impl io::Write) -> Result<(), TK1ConvertError> {
     let serial_circ = SerialCircuit::encode(circ)?;
     serde_json::to_writer(w, &serial_circ)?;
     Ok(())
 }
 
 /// Save a circuit in TK1 JSON format to a String.
-pub fn save_tk1_json_str(circ: &Hugr) -> Result<String, TK1ConvertError> {
+pub fn save_tk1_json_str(circ: &impl Circuit) -> Result<String, TK1ConvertError> {
     let mut buf = io::BufWriter::new(Vec::new());
     save_tk1_json_writer(circ, &mut buf)?;
     let bytes = buf.into_inner().unwrap();
diff --git a/src/utils.rs b/src/utils.rs
index 124cda375..f0f8ccd22 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -36,13 +36,22 @@ pub(crate) fn build_simple_circuit(
 #[allow(dead_code)]
 #[cfg(test)]
 pub(crate) mod test {
+    #[allow(unused_imports)]
+    use hugr::HugrView;
+
     /// Open a browser page to render a dot string graph.
     ///
     /// This can be used directly on the output of `Hugr::dot_string`
     #[cfg(not(ci_run))]
-    pub(crate) fn viz_dotstr(dotstr: &str) {
+    pub(crate) fn viz_dotstr(dotstr: impl AsRef<str>) {
         let mut base: String = "https://dreampuf.github.io/GraphvizOnline/#".into();
-        base.push_str(&urlencoding::encode(dotstr));
+        base.push_str(&urlencoding::encode(dotstr.as_ref()));
         webbrowser::open(&base).unwrap();
     }
+
+    /// Open a browser page to render a HugrView's dot string graph.
+    #[cfg(not(ci_run))]
+    pub(crate) fn viz_hugr(hugr: &impl HugrView) {
+        viz_dotstr(hugr.dot_string());
+    }
 }
diff --git a/taso-optimiser/src/main.rs b/taso-optimiser/src/main.rs
index c31afb0de..1d06c646f 100644
--- a/taso-optimiser/src/main.rs
+++ b/taso-optimiser/src/main.rs
@@ -10,11 +10,9 @@ use std::path::PathBuf;
 use std::process::exit;
 
 use clap::Parser;
-use hugr::Hugr;
-use tket2::json::{load_tk1_json_file, TKETDecode};
+use tket2::json::{load_tk1_json_file, save_tk1_json_file};
 use tket2::optimiser::taso::log::TasoLogger;
 use tket2::optimiser::TasoOptimiser;
-use tket_json_rs::circuit_json::SerialCircuit;
 
 #[cfg(feature = "peak_alloc")]
 use peak_alloc::PeakAlloc;
@@ -82,14 +80,6 @@ struct CmdLineArgs {
     n_threads: Option<NonZeroUsize>,
 }
 
-fn save_tk1_json_file(path: impl AsRef<Path>, circ: &Hugr) -> Result<(), std::io::Error> {
-    let file = File::create(path)?;
-    let writer = BufWriter::new(file);
-    let serial_circ = SerialCircuit::encode(circ).unwrap();
-    serde_json::to_writer_pretty(writer, &serial_circ)?;
-    Ok(())
-}
-
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let opts = CmdLineArgs::parse();
 
@@ -129,7 +119,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let opt_circ = optimiser.optimise_with_log(&circ, taso_logger, opts.timeout, n_threads);
 
     println!("Saving result");
-    save_tk1_json_file(output_path, &opt_circ)?;
+    save_tk1_json_file(&opt_circ, output_path)?;
 
     #[cfg(feature = "peak_alloc")]
     println!("Peak memory usage: {} GB", PEAK_ALLOC.peak_usage_as_gb());