egg_stats/
datum.rs

1use std::borrow::Cow;
2
3pub trait Recorder<L, N>
4where
5    L: egg::Language,
6    N: egg::Analysis<L>,
7{
8    fn identifier(&self) -> Cow<'static, str>;
9
10    #[allow(unused_variables)]
11    fn record_before_search<'r>(
12        &self,
13        iteration: usize,
14        egraph: &'r egg::EGraph<L, N>,
15        rewrite: &'r egg::Rewrite<L, N>,
16    ) -> Option<String> {
17        None
18    }
19
20    #[allow(unused_variables)]
21    fn record_after_search<'r>(
22        &self,
23        iteration: usize,
24        egraph: &'r egg::EGraph<L, N>,
25        rewrite: &'r egg::Rewrite<L, N>,
26        matches: &[egg::SearchMatches<L>],
27    ) -> Option<String> {
28        None
29    }
30
31    #[allow(unused_variables)]
32    fn record_before_rewrite(
33        &self,
34        iteration: usize,
35        egraph: &mut egg::EGraph<L, N>,
36        rewrite: &egg::Rewrite<L, N>,
37        matches: &[egg::SearchMatches<L>],
38    ) -> Option<String> {
39        None
40    }
41
42    #[allow(unused_variables)]
43    fn record_after_rewrite(
44        &self,
45        iteration: usize,
46        egraph: &mut egg::EGraph<L, N>,
47        rewrite: &egg::Rewrite<L, N>,
48        n_matches: usize,
49    ) -> Option<String> {
50        None
51    }
52}
53
54pub struct BeforeSearchArgs<'a, L, N>
55where
56    L: egg::Language,
57    N: egg::Analysis<L>,
58{
59    pub iteration: usize,
60    pub egraph: &'a egg::EGraph<L, N>,
61    pub rewrite: &'a egg::Rewrite<L, N>,
62}
63
64pub struct AfterSearchArgs<'a, L, N>
65where
66    L: egg::Language,
67    N: egg::Analysis<L>,
68{
69    pub iteration: usize,
70    pub egraph: &'a egg::EGraph<L, N>,
71    pub rewrite: &'a egg::Rewrite<L, N>,
72    pub matches: &'a [egg::SearchMatches<'a, L>],
73}
74
75pub struct BeforeRewriteArgs<'a, L, N>
76where
77    L: egg::Language,
78    N: egg::Analysis<L>,
79{
80    pub iteration: usize,
81    pub egraph: &'a mut egg::EGraph<L, N>,
82    pub rewrite: &'a egg::Rewrite<L, N>,
83    pub matches: &'a [egg::SearchMatches<'a, L>],
84}
85
86pub struct AfterRewriteArgs<'a, L, N>
87where
88    L: egg::Language,
89    N: egg::Analysis<L>,
90{
91    pub iteration: usize,
92    pub egraph: &'a mut egg::EGraph<L, N>,
93    pub rewrite: &'a egg::Rewrite<L, N>,
94    pub n_matches: usize,
95}
96
97#[macro_export]
98macro_rules! impl_recorder {
99    ($type_name:ident; $($name:ident => $val:expr),*) => {
100        impl<L, N> Recorder<L, N> for $type_name
101        where
102            L: egg::Language,
103            N: egg::Analysis<L>,
104        {
105            $(impl_recorder!(impl $name, $val);)*
106        }
107    };
108
109    ($vis:vis struct $type_name:ident; $($name:ident => $val:expr),*) => {
110        $vis struct $type_name;
111        impl_recorder!($type_name; $($name => $val),*);
112    };
113
114    (impl identifier, $name:literal) => {
115        fn identifier(&self) -> Cow<'static, str> {
116            $name.into()
117        }
118    };
119
120    (impl before_search, $br:expr) => {
121        fn record_before_search<'r>(
122            &self,
123            iteration: usize,
124            egraph: &'r egg::EGraph<L, N>,
125            rewrite: &'r egg::Rewrite<L, N>,
126        ) -> Option<String> {
127            $br(self, $crate::BeforeSearchArgs {
128                iteration,
129                egraph,
130                rewrite,
131            })
132        }
133    };
134
135    (impl after_search, $br:expr) => {
136        fn record_after_search<'r>(
137            &self,
138            iteration: usize,
139            egraph: &'r egg::EGraph<L, N>,
140            rewrite: &'r egg::Rewrite<L, N>,
141            matches: &[egg::SearchMatches<L>],
142        ) -> Option<String> {
143            $br(self, $crate::AfterSearchArgs {
144                iteration,
145                egraph,
146                rewrite,
147                matches,
148            })
149        }
150    };
151
152    (impl before_rewrite, $br:expr) => {
153        fn record_before_rewrite(
154            &self,
155            iteration: usize,
156            egraph: &mut egg::EGraph<L, N>,
157            rewrite: &egg::Rewrite<L, N>,
158            matches: &[egg::SearchMatches<L>],
159        ) -> Option<String> {
160            $br(self, $crate::BeforeRewriteArgs {
161                iteration,
162                egraph,
163                rewrite,
164                matches,
165            })
166        }
167    };
168
169    (impl after_rewrite, $br:expr) => {
170        fn record_after_rewrite(
171            &self,
172            iteration: usize,
173            egraph: &mut egg::EGraph<L, N>,
174            rewrite: &egg::Rewrite<L, N>,
175            n_matches: usize,
176        ) -> Option<String> {
177            $br(self, $crate::AfterRewriteArgs {
178                iteration,
179                egraph,
180                rewrite,
181                n_matches,
182            })
183        }
184    };
185}