use std::path::PathBuf;
use std::time::{Duration, Instant};
fn mean_stdev(data: &[f64]) -> (f64, f64) {
assert_ne!(data.len(), 0);
let sum = data.iter().sum::<f64>();
let n = data.len() as f64;
let mean = sum / n;
let variance = data
.iter()
.map(|value| {
let diff = mean - (*value as f64);
diff * diff
})
.sum::<f64>()
/ n;
(mean, variance.sqrt())
}
pub fn env_var<T>(s: &str) -> Option<T>
where
T: std::str::FromStr,
T::Err: std::fmt::Debug,
{
use std::env::VarError;
match std::env::var(s) {
Err(VarError::NotPresent) => None,
Err(VarError::NotUnicode(_)) => panic!("Environment variable {} isn't unicode", s),
Ok(v) if v.is_empty() => None,
Ok(v) => match v.parse() {
Ok(v) => Some(v),
Err(err) => panic!("Couldn't parse environment variable {}={}, {:?}", s, v, err),
},
}
}
pub struct Reporter<T> {
name: String,
times: Option<Vec<f64>>,
result: T,
}
impl<T> Reporter<T> {
pub fn into_inner(self) -> T {
let _ = self.name;
let _ = self.times;
self.result
}
#[cfg(not(feature = "reports"))]
pub fn report<R>(self, to_report: impl FnOnce(&T) -> &R) -> T {
if let Some(dir) = env_var::<PathBuf>("EGG_BENCH_DIR") {
eprintln!(
"EGG_BENCH_DIR is set to '{:?}', but the 'reports' feature is not enabled",
dir
);
}
to_report(&self.result);
self.result
}
#[cfg(feature = "reports")]
pub fn report<R>(self, to_report: impl FnOnce(&T) -> &R) -> T
where
R: serde::Serialize,
{
let directory = match env_var::<PathBuf>("EGG_BENCH_DIR") {
None => {
eprintln!("EGG_BENCH_DIR not set, skipping reporting");
return self.result;
}
Some(dir) => {
assert!(dir.is_dir(), "EGG_BENCH_DIR is not a directory: {:?}", dir);
dir
}
};
let filename = format!("{}.json", self.name);
let path = directory.join(&filename);
let file = std::fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(&path)
.unwrap_or_else(|err| panic!("Failed to open {:?}: {}", path, err));
let report = serde_json::json!({
"name": &self.name,
"times": self.times.as_deref(),
"data": to_report(&self.result),
});
serde_json::to_writer_pretty(file, &report)
.unwrap_or_else(|err| panic!("Failed to serialize report to {:?}: {}", path, err));
println!("Wrote report to {:?}", path);
self.result
}
}
pub fn run<T>(name: impl Into<String>, mut f: impl FnMut() -> T) -> Reporter<T> {
let name = name.into();
let seconds: f64 = match env_var("EGG_BENCH") {
Some(s) => s,
None => {
return Reporter {
name,
times: None,
result: f(),
}
}
};
let duration = Duration::from_secs_f64(seconds);
let start = Instant::now();
let mut times = vec![];
println!("benching {} for {} seconds...", name, seconds);
let result = loop {
let i = Instant::now();
let result = f();
times.push(i.elapsed().as_secs_f64());
if start.elapsed() > duration {
break result;
}
};
let (mean, stdev) = mean_stdev(×);
println!("bench {}:", name);
println!(" n = {}", times.len());
println!(" μ = {}", mean);
println!(" σ = {}", stdev);
Reporter {
name,
times: Some(times),
result,
}
}
#[macro_export]
macro_rules! test_fn {
(
$(#[$meta:meta])*
$name:ident, $rules:expr,
$start:literal
=>
$($goal:literal),+ $(,)?
$(@check $check_fn:expr)?
) => {
$crate::test_fn! {
$(#[$meta])*
$name, $rules,
runner = $crate::Runner::<_, _, ()>::default(),
$start => $( $goal ),+
$(@check $check_fn)?
}
};
(
$(#[$meta:meta])*
$name:ident, $rules:expr,
runner = $runner:expr,
$start:literal
=>
$($goal:literal),+ $(,)?
$(@check $check_fn:expr)?
) => {
$(#[$meta])*
#[test]
fn $name() {
let _ = env_logger::builder().is_test(true).try_init();
let name = stringify!($name);
let start: $crate::RecExpr<_> = $start.parse().unwrap();
let rules = $rules;
let runner: $crate::Runner<_, _, ()> = $crate::test::run(name, || {
let mut runner = $runner.with_expr(&start);
if let Some(lim) = $crate::test::env_var("EGG_NODE_LIMIT") {
runner = runner.with_node_limit(lim)
}
if let Some(lim) = $crate::test::env_var("EGG_ITER_LIMIT") {
runner = runner.with_iter_limit(lim)
}
if let Some(lim) = $crate::test::env_var("EGG_TIME_LIMIT") {
runner = runner.with_time_limit(std::time::Duration::from_secs(lim))
}
runner.run(&rules)
}).report(|r| &r.iterations);
runner.print_report();
let goals = &[$(
$goal.parse().unwrap()
),+];
let id = runner.egraph.find(*runner.roots.last().unwrap());
runner.egraph.check_goals(id, goals);
$( ($check_fn)(runner) )?
}
};
}