use egg::*;
use egg_stats::{recorders, LoggingScheduler};
use std::path::Path;
use std::{fs::OpenOptions, time::Instant};
define_language! {
enum SimpleLanguage {
Num(i32),
"+" = Add([Id; 2]),
"*" = Mul([Id; 2]),
"/" = Div([Id; 2]),
"<<" = Lsh([Id; 2]),
Symbol(Symbol),
}
}
fn make_rules() -> Vec<Rewrite<SimpleLanguage, ()>> {
vec![
rewrite!("commute-add"; "(+ ?a ?b)" => "(+ ?b ?a)"),
rewrite!("assoc-add"; "(+ ?a (+ ?b ?c))" => "(+ (+ ?a ?b) ?c)"),
rewrite!("commute-mul"; "(* ?a ?b)" => "(* ?b ?a)"),
rewrite!("assoc-mul"; "(* ?a (* ?b ?c))" => "(* (* ?a ?b) ?c)"),
rewrite!("add-0"; "(+ ?a 0)" => "?a"),
rewrite!("mul-0"; "(* ?a 0)" => "0"),
rewrite!("mul-1"; "(* ?a 1)" => "?a"),
rewrite!("lsh"; "(/ ?a 2)" => "(<< ?a 1)"),
rewrite!("lsh-rev"; "(<< ?a 1)" => "(/ ?a 2)"),
rewrite!("div-same"; "(/ ?a ?a)" => "1"),
rewrite!("div-mul-distr"; "(/ (* ?a ?b) ?c)" => "(* ?a (/ ?b ?c))"),
rewrite!("div-mul-distr-rev"; "(* ?a (/ ?b ?c))" => "(/ (* ?a ?b) ?c)"),
]
}
fn simplify_with(
s: &str,
scheduler: impl RewriteScheduler<SimpleLanguage, ()> + 'static,
path: impl AsRef<Path>,
) {
let expr: RecExpr<SimpleLanguage> = s.parse().unwrap();
let mut egraph = EGraph::new(());
let root = egraph.add_expr(&expr);
Runner::default()
.with_scheduler(
LoggingScheduler::from(scheduler)
.with_out_file(
OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(path.as_ref())
.unwrap(),
)
.with_logging_enabled(true)
.with_recorder(recorders::Timestamp::new(Instant::now()))
.with_recorder(recorders::NumberENodes)
.with_recorder(recorders::NumberEClasses)
.with_recorder(recorders::BestProgram::new_with(|| AstSize, root)),
)
.with_egraph(egraph)
.run(&make_rules());
println!("Wrote {:?}", path.as_ref());
}
fn main() {
let expr = "(/ (* 1 (* 2 ?a)) 2)";
simplify_with(expr, BackoffScheduler::default(), "backoff.csv");
simplify_with(expr, SimpleScheduler, "simple.csv");
}