simple/
simple.rs

1use egg::*;
2use egg_stats::{recorders, LoggingScheduler};
3use std::path::Path;
4use std::{fs::OpenOptions, time::Instant};
5
6define_language! {
7    enum SimpleLanguage {
8        Num(i32),
9        "+" = Add([Id; 2]),
10        "*" = Mul([Id; 2]),
11        "/" = Div([Id; 2]),
12        "<<" = Lsh([Id; 2]),
13        Symbol(Symbol),
14    }
15}
16
17fn make_rules() -> Vec<Rewrite<SimpleLanguage, ()>> {
18    vec![
19        rewrite!("commute-add"; "(+ ?a ?b)" => "(+ ?b ?a)"),
20        rewrite!("assoc-add"; "(+ ?a (+ ?b ?c))" => "(+ (+ ?a ?b) ?c)"),
21        rewrite!("commute-mul"; "(* ?a ?b)" => "(* ?b ?a)"),
22        rewrite!("assoc-mul"; "(* ?a (* ?b ?c))" => "(* (* ?a ?b) ?c)"),
23        rewrite!("add-0"; "(+ ?a 0)" => "?a"),
24        rewrite!("mul-0"; "(* ?a 0)" => "0"),
25        rewrite!("mul-1"; "(* ?a 1)" => "?a"),
26        rewrite!("lsh"; "(/ ?a 2)" => "(<< ?a 1)"),
27        rewrite!("lsh-rev"; "(<< ?a 1)" => "(/ ?a 2)"),
28        rewrite!("div-same"; "(/ ?a ?a)" => "1"),
29        rewrite!("div-mul-distr"; "(/ (* ?a ?b) ?c)" => "(* ?a (/ ?b ?c))"),
30        rewrite!("div-mul-distr-rev"; "(* ?a (/ ?b ?c))" => "(/ (* ?a ?b) ?c)"),
31    ]
32}
33
34/// parse an expression, simplify it using egg, and pretty print it back out
35fn simplify_with(
36    s: &str,
37    scheduler: impl RewriteScheduler<SimpleLanguage, ()> + 'static,
38    path: impl AsRef<Path>,
39) {
40    // parse the expression, the type annotation tells it which Language to use
41    let expr: RecExpr<SimpleLanguage> = s.parse().unwrap();
42
43    let mut egraph = EGraph::new(());
44    let root = egraph.add_expr(&expr);
45
46    // simplify the expression using a Runner, which creates an e-graph with
47    // the given expression and runs the given rules over it
48    Runner::default()
49        .with_scheduler(
50            LoggingScheduler::from(scheduler)
51                .with_out_file(
52                    OpenOptions::new()
53                        .write(true)
54                        .create(true)
55                        .truncate(true)
56                        .open(path.as_ref())
57                        .unwrap(),
58                )
59                .with_logging_enabled(true)
60                .with_recorder(recorders::Timestamp::new(Instant::now()))
61                .with_recorder(recorders::NumberENodes)
62                .with_recorder(recorders::NumberEClasses)
63                .with_recorder(recorders::BestProgram::new_with(|| AstSize, root)),
64        )
65        .with_egraph(egraph)
66        .run(&make_rules());
67
68    println!("Wrote {:?}", path.as_ref());
69}
70
71fn main() {
72    let expr = "(/ (* 1 (* 2 ?a)) 2)";
73    simplify_with(expr, BackoffScheduler::default(), "backoff.csv");
74
75    simplify_with(expr, SimpleScheduler, "simple.csv");
76}