egg_stats/
recorders.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
use std::{borrow::Cow, time::Instant};

use crate::{impl_recorder, AfterRewriteArgs, BeforeRewriteArgs, Recorder};

pub struct Timestamp {
    start: Instant,
}

impl Default for Timestamp {
    fn default() -> Self {
        Timestamp {
            start: Instant::now(),
        }
    }
}

impl Timestamp {
    pub fn new(start: Instant) -> Self {
        Timestamp { start }
    }
}

impl_recorder! {
    Timestamp;
    identifier => "time",
    before_rewrite => |ts: &Self, _| {
        let passed = ts.start.elapsed();
        Some(format!("{}", passed.as_millis()))
    }
}

impl_recorder! {
    pub struct NumberENodes;
    identifier => "enodes",
    before_rewrite => |_, x: BeforeRewriteArgs<_, _>| {
        Some(format!("{}", x.egraph.total_size()))
    },
    after_rewrite => |_, x: AfterRewriteArgs<_, _>| {
        Some(format!("{}", x.egraph.total_size()))
    }
}

impl_recorder! {
    pub struct NumberEClasses;
    identifier => "eclasses",
    before_rewrite => |_, x: BeforeRewriteArgs<_, _>| {
        Some(format!("{}", x.egraph.number_of_classes()))
    },
    after_rewrite => |_, x: AfterRewriteArgs<_, _>| {
        Some(format!("{}", x.egraph.number_of_classes()))
    }
}

pub struct BestProgram<C> {
    cost_fn: Box<dyn Fn() -> C>,
    root: egg::Id,
    rebuild: bool,
}

impl<C> BestProgram<C> {
    pub fn new(cost_fn: C, root: egg::Id) -> Self
    where
        C: Clone + 'static,
    {
        let cost_fn_clone = cost_fn.clone();
        BestProgram {
            cost_fn: Box::new(move || cost_fn_clone.clone()),
            root,
            rebuild: false,
        }
    }

    pub fn new_with(cost_fn: impl Fn() -> C + 'static, root: egg::Id) -> Self {
        BestProgram {
            cost_fn: Box::new(cost_fn),
            root,
            rebuild: false,
        }
    }

    pub fn with_rebuild(mut self, rebuild: bool) -> Self {
        self.rebuild = rebuild;
        self
    }
}

impl<L, N, C> Recorder<L, N> for BestProgram<C>
where
    L: egg::Language,
    N: egg::Analysis<L>,
    C: egg::CostFunction<L>,
    <C as egg::CostFunction<L>>::Cost: std::fmt::Display,
{
    fn identifier(&self) -> Cow<'static, str> {
        "cost".into()
    }

    fn record_before_rewrite(
        &self,
        _iteration: usize,
        egraph: &mut egg::EGraph<L, N>,
        _rewrite: &egg::Rewrite<L, N>,
        _matches: &[egg::SearchMatches<L>],
    ) -> Option<String> {
        let extractor = egg::Extractor::new(egraph, (self.cost_fn)());
        Some(format!("{}", extractor.find_best_cost(self.root)))
    }

    fn record_after_rewrite(
        &self,
        _iteration: usize,
        egraph: &mut egg::EGraph<L, N>,
        _rewrite: &egg::Rewrite<L, N>,
        _n_matches: usize,
    ) -> Option<String> {
        if self.rebuild {
            egraph.rebuild();
            let extractor = egg::Extractor::new(egraph, (self.cost_fn)());
            Some(format!("{}", extractor.find_best_cost(self.root)))
        } else {
            None
        }
    }
}