use chrono::Local; use include_dir::{include_dir, Dir}; use serde::{Deserialize, Serialize}; use std::fs; use std::io; use std::io::ErrorKind; use std::path::Path; use std::path::PathBuf; use std::process::Command; pub static BUILTIN_EVAL_ASSETS: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/assets"); #[derive(Clone, Serialize, Deserialize, Debug)] pub struct BenchmarkWorkDir { pub base_path: PathBuf, pub run_dir: PathBuf, pub cwd: PathBuf, pub run_id: Option, } impl Default for BenchmarkWorkDir { fn default() -> Self { Self::new("work_dir".to_string(), Vec::new()) } } impl BenchmarkWorkDir { pub fn new(work_dir_name: String, include_dirs: Vec) -> Self { let run_dir = std::env::current_dir().unwrap().canonicalize().unwrap(); let base_path = PathBuf::from(format!("./{}", work_dir_name)); fs::create_dir_all(&base_path).unwrap(); let base_path = PathBuf::from(&base_path).canonicalize().unwrap(); // abs paths from dir-strings let dirs = Self::canonical_dirs(include_dirs); // deep copy each dir let _: Vec<_> = dirs .iter() .map(|d| Self::deep_copy(d.as_path(), base_path.as_path(), true)) .collect(); Self::copy_auto_included_dirs(&base_path); std::env::set_current_dir(&base_path).unwrap(); BenchmarkWorkDir { base_path: base_path.clone(), run_dir, cwd: base_path.clone(), run_id: None, } } pub fn init_experiment() { // create experiment folder let current_time = Local::now().format("%H:%M:%S").to_string(); let current_date = Local::now().format("%Y-%m-%d").to_string(); let exp_name = format!("{}-{}", ¤t_date, current_time); let base_path = PathBuf::from(format!("./benchmark-{}", exp_name)); fs::create_dir_all(&base_path).unwrap(); std::env::set_current_dir(&base_path).unwrap(); } pub fn canonical_dirs(include_dirs: Vec) -> Vec { include_dirs .iter() .map(|d| { let canon = d.canonicalize(); if canon.is_err() { eprintln!("{:?} can't be canonicalized", d); panic!(); } canon.unwrap() }) .collect::>() } fn copy_auto_included_dirs(dest: &Path) { let mut assets_dest = dest.to_path_buf(); assets_dest.push("assets"); if !assets_dest.exists() { fs::create_dir_all(&assets_dest).unwrap(); } BUILTIN_EVAL_ASSETS.extract(assets_dest).unwrap(); } pub fn cd(&mut self, path: PathBuf) -> anyhow::Result<&mut Self> { fs::create_dir_all(&path)?; std::env::set_current_dir(&path)?; self.cwd = path; Ok(self) } pub(crate) fn _run_dir(&mut self) -> Option { if let Some(run_id) = &self.run_id { let mut eval_dir = self.base_path.clone(); eval_dir.push(run_id); return Some(eval_dir); } None } pub fn set_eval(&mut self, eval: &str, run_id: String) { self.run_id = Some(run_id.clone()); let eval = eval.replace(":", std::path::MAIN_SEPARATOR_STR); let mut eval_dir = self.base_path.clone(); eval_dir.push(run_id); eval_dir.push(eval); self.cd(eval_dir.clone()) .unwrap_or_else(|_| panic!("Failed to execute cd into {}", eval_dir.clone().display())); } fn chop_relative_base>(path: P) -> anyhow::Result { let path = path.as_ref(); // Get the path components as an iterator let mut components = path.components(); // Check the first component if let Some(first) = components.next() { use std::path::Component; match first { Component::ParentDir => Err(anyhow::anyhow!("RelativePathBaseError: Only paths relative to the current working directory are supported.")), // If first component is "." Component::CurDir => Ok(components.collect()), // Otherwise, keep the full path _ => { // Create a new PathBuf let mut result = PathBuf::new(); // Add back the first component result.push(first); // Add all remaining components result.extend(components); Ok(result) } } } else { // Empty path Ok(PathBuf::new()) } } pub fn fs_get(&mut self, path: String) -> anyhow::Result { let p = PathBuf::from(&path); if p.exists() { return Ok(PathBuf::from(path)); } if p.is_absolute() { return Err(anyhow::anyhow!("AbsolutePathError: Only paths relative to the current working directory are supported.")); } let asset_rel_path = Self::chop_relative_base(p.clone()) .unwrap_or_else(|_| panic!("AbsolutePathError: Only paths relative to the current working directory are supported.")); let here = PathBuf::from(".").canonicalize()?; let artifact_at_root = self.base_path.clone().join(asset_rel_path); Self::deep_copy(artifact_at_root.as_path(), here.as_path(), true)?; Ok(PathBuf::from(path)) } pub(crate) fn deep_copy(src: P, dst: Q, recursive: bool) -> io::Result<()> where P: AsRef, Q: AsRef, { let src = src.as_ref(); let dst = dst.as_ref(); let mut cmd = Command::new("cp"); // Add -r flag if recursive is true if recursive { cmd.arg("-r"); } // Add source and destination paths cmd.arg(src).arg(dst); // Execute the command let output = cmd.output()?; if output.status.success() { Ok(()) } else { let error_message = String::from_utf8_lossy(&output.stderr).to_string(); Err(io::Error::new(ErrorKind::Other, error_message)) } } pub fn save(&self) { let work_dir = serde_json::to_string_pretty(&self).unwrap(); fs::write("work_dir.json", work_dir).expect("Unable to write work-dir as file"); } } impl Drop for BenchmarkWorkDir { fn drop(&mut self) { std::env::set_current_dir(&self.run_dir).unwrap(); } }