egg_stats/
recorders.rs

1use std::{borrow::Cow, time::Instant};
2
3use crate::{impl_recorder, AfterRewriteArgs, BeforeRewriteArgs, Recorder};
4
5pub struct Timestamp {
6    start: Instant,
7}
8
9impl Default for Timestamp {
10    fn default() -> Self {
11        Timestamp {
12            start: Instant::now(),
13        }
14    }
15}
16
17impl Timestamp {
18    pub fn new(start: Instant) -> Self {
19        Timestamp { start }
20    }
21}
22
23impl_recorder! {
24    Timestamp;
25    identifier => "time",
26    before_rewrite => |ts: &Self, _| {
27        Some(format!("{}", ts.start.elapsed().as_millis()))
28    },
29    after_rewrite => |ts: &Self, _| {
30        Some(format!("{}", ts.start.elapsed().as_millis()))
31    }
32}
33
34impl_recorder! {
35    pub struct NumberENodes;
36    identifier => "enodes",
37    before_rewrite => |_, x: BeforeRewriteArgs<_, _>| {
38        Some(format!("{}", x.egraph.total_size()))
39    },
40    after_rewrite => |_, x: AfterRewriteArgs<_, _>| {
41        Some(format!("{}", x.egraph.total_size()))
42    }
43}
44
45impl_recorder! {
46    pub struct NumberEClasses;
47    identifier => "eclasses",
48    before_rewrite => |_, x: BeforeRewriteArgs<_, _>| {
49        Some(format!("{}", x.egraph.number_of_classes()))
50    },
51    after_rewrite => |_, x: AfterRewriteArgs<_, _>| {
52        Some(format!("{}", x.egraph.number_of_classes()))
53    }
54}
55
56pub struct BestProgram<C> {
57    cost_fn: Box<dyn Fn() -> C>,
58    root: egg::Id,
59    rebuild: bool,
60}
61
62impl<C> BestProgram<C> {
63    pub fn new(cost_fn: C, root: egg::Id) -> Self
64    where
65        C: Clone + 'static,
66    {
67        let cost_fn_clone = cost_fn.clone();
68        BestProgram {
69            cost_fn: Box::new(move || cost_fn_clone.clone()),
70            root,
71            rebuild: false,
72        }
73    }
74
75    pub fn new_with(cost_fn: impl Fn() -> C + 'static, root: egg::Id) -> Self {
76        BestProgram {
77            cost_fn: Box::new(cost_fn),
78            root,
79            rebuild: false,
80        }
81    }
82
83    pub fn with_rebuild(mut self, rebuild: bool) -> Self {
84        self.rebuild = rebuild;
85        self
86    }
87}
88
89impl<L, N, C> Recorder<L, N> for BestProgram<C>
90where
91    L: egg::Language,
92    N: egg::Analysis<L>,
93    C: egg::CostFunction<L>,
94    <C as egg::CostFunction<L>>::Cost: std::fmt::Display,
95{
96    fn identifier(&self) -> Cow<'static, str> {
97        "cost".into()
98    }
99
100    fn record_before_rewrite(
101        &self,
102        _iteration: usize,
103        egraph: &mut egg::EGraph<L, N>,
104        _rewrite: &egg::Rewrite<L, N>,
105        _matches: &[egg::SearchMatches<L>],
106    ) -> Option<String> {
107        let extractor = egg::Extractor::new(egraph, (self.cost_fn)());
108        Some(format!("{}", extractor.find_best_cost(self.root)))
109    }
110
111    fn record_after_rewrite(
112        &self,
113        _iteration: usize,
114        egraph: &mut egg::EGraph<L, N>,
115        _rewrite: &egg::Rewrite<L, N>,
116        _n_matches: usize,
117    ) -> Option<String> {
118        if self.rebuild {
119            egraph.rebuild();
120            let extractor = egg::Extractor::new(egraph, (self.cost_fn)());
121            Some(format!("{}", extractor.find_best_cost(self.root)))
122        } else {
123            None
124        }
125    }
126}
egg_stats/recorders.rs

egg_stats/
recorders.rs
