add tsp benchmark flag

1 files changed, 13 insertions, 5 deletions

diff --git a/src/main.rs b/src/main.rs
index c8f5bbc..993e490 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -48,6 +48,10 @@ struct Args {
     #[arg(short = '0', long)]
     stdout0: bool,
 
+    /// Output total tour length to stderr
+    #[arg(short = 'b', long)]
+    benchmark: bool,
+
     images: Vec<PathBuf>,
 }
 
@@ -74,12 +78,12 @@ fn hash_file(p: &PathBuf) -> Result<[u8; 32]> {
         .unwrap())
 }
 
-fn process_embedder<E>(mut e: E, args: &Args, cfg: &Config) -> Result<Vec<PathBuf>>
+fn process_embedder<E>(mut e: E, args: &Args, cfg: &Config) -> Result<(Vec<PathBuf>, f64)>
 where
     E: BatchEmbedder,
 {
     if args.images.is_empty() {
-        return Ok(Vec::new());
+        return Ok((Vec::new(), 0.));
     }
 
     let db = sled::open(cfg.base_dirs.place_cache_file("embeddings.db")?)?;
@@ -121,9 +125,9 @@ where
     }
 
     let embeds: Vec<_> = embeds.into_iter().map(|e| e.unwrap()).collect();
-    let tsp_path = tsp(&embeds);
+    let (tsp_path, total_dist) = tsp(&embeds);
 
-    Ok(tsp_path.iter().map(|i| args.images[*i].clone()).collect())
+    Ok((tsp_path.iter().map(|i| args.images[*i].clone()).collect(), total_dist))
 }
 
 fn copy_into(tsp: &[PathBuf], target: &PathBuf, use_symlinks: bool) -> Result<()> {
@@ -153,13 +157,17 @@ fn main() -> Result<()> {
     let cfg = get_config()?;
     let args = Args::parse();
 
-    let tsp_path = match args.embedder {
+    let (tsp_path, total_dist) = match args.embedder {
         Embedder::Brightness => process_embedder(BrightnessEmbedder, &args, &cfg),
         Embedder::Hue => process_embedder(HueEmbedder, &args, &cfg),
         Embedder::Color => process_embedder(ColorEmbedder, &args, &cfg),
         Embedder::Content => process_embedder(ContentEmbedder::new(&cfg), &args, &cfg),
     }?;
 
+    if args.benchmark {
+        eprintln!("Found tour with length: {}", total_dist);
+    }
+
     if let Some(p) = args.symlink_dir {
         copy_into(&tsp_path, &p, true)?
     }