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}