1use super::{
2 plan::{MatchInputs, OptimizePlan, PlanNode},
3 rule::{Rule, Transformed},
4 stats::Stats,
5};
6use crate::Operand;
7use std::marker::PhantomData;
8
9pub trait Pattern<O: Operand> {
10 type Bindings;
11
12 fn try_match(&self, operand: &O) -> Option<Self::Bindings>;
13
14 fn rewrite<F>(self, rewrite: F) -> impl Rule<O>
15 where
16 Self: Sized + Send + Sync + 'static,
17 O: 'static,
18 F: Fn(Self::Bindings, &Stats) -> Option<O> + Send + Sync + 'static,
19 {
20 PatternRule {
21 pattern: self,
22 rewrite,
23 }
24 }
25
26 fn guard<G>(self, guard: G) -> GuardedPattern<Self, G>
27 where
28 Self: Sized,
29 G: Fn(&Stats) -> bool,
30 {
31 GuardedPattern {
32 pattern: self,
33 guard,
34 }
35 }
36
37 fn or<Q>(self, other: Q) -> impl Pattern<O, Bindings = Self::Bindings>
38 where
39 Self: Sized,
40 Q: Pattern<O, Bindings = Self::Bindings>,
41 {
42 OrPattern {
43 left: self,
44 right: other,
45 }
46 }
47}
48
49struct PatternRule<P, F> {
50 pattern: P,
51 rewrite: F,
52}
53
54impl<O, P, F> Rule<O> for PatternRule<P, F>
55where
56 O: Operand,
57 P: Pattern<O> + Send + Sync + 'static,
58 F: Fn(P::Bindings, &Stats) -> Option<O> + Send + Sync + 'static,
59{
60 fn apply(&self, operand: O, stats: &Stats) -> Transformed<O> {
61 let Some(bindings) = self.pattern.try_match(&operand) else {
62 return Transformed::unchanged(operand);
63 };
64
65 match (self.rewrite)(bindings, stats) {
66 Some(rewritten) => Transformed::changed(rewritten),
67 None => Transformed::unchanged(operand),
68 }
69 }
70}
71
72pub struct GuardedPattern<P, G> {
73 pattern: P,
74 guard: G,
75}
76
77impl<P, G> GuardedPattern<P, G> {
78 pub fn rewrite<O, F>(self, rewrite: F) -> impl Rule<O>
79 where
80 O: Operand + 'static,
81 P: Pattern<O> + Send + Sync + 'static,
82 G: Fn(&Stats) -> bool + Send + Sync + 'static,
83 F: Fn(P::Bindings, &Stats) -> Option<O> + Send + Sync + 'static,
84 {
85 let guard = self.guard;
86
87 PatternRule {
88 pattern: self.pattern,
89 rewrite: move |bindings, stats: &Stats| {
90 if guard(stats) {
91 rewrite(bindings, stats)
92 } else {
93 None
94 }
95 },
96 }
97 }
98}
99
100struct OrPattern<P, Q> {
101 left: P,
102 right: Q,
103}
104
105impl<O: Operand, P: Pattern<O>, Q: Pattern<O, Bindings = P::Bindings>> Pattern<O>
106 for OrPattern<P, Q>
107{
108 type Bindings = P::Bindings;
109
110 fn try_match(&self, operand: &O) -> Option<Self::Bindings> {
111 self.left
112 .try_match(operand)
113 .or_else(|| self.right.try_match(operand))
114 }
115}
116
117pub struct NotPattern<P> {
118 inner: P,
119}
120
121#[must_use]
122pub const fn not<P>(inner: P) -> NotPattern<P> {
123 NotPattern { inner }
124}
125
126impl<O: Operand, P: Pattern<O>> Pattern<O> for NotPattern<P> {
127 type Bindings = ();
128
129 fn try_match(&self, operand: &O) -> Option<Self::Bindings> {
130 self.inner.try_match(operand).is_none().then_some(())
131 }
132}
133
134pub struct Wildcard;
135
136#[must_use]
137pub const fn any() -> Wildcard {
138 Wildcard
139}
140
141impl<O: Operand> Pattern<O> for Wildcard {
142 type Bindings = ();
143
144 fn try_match(&self, _operand: &O) -> Option<Self::Bindings> {
145 Some(())
146 }
147}
148
149pub struct Capture;
150
151#[must_use]
152pub const fn capture() -> Capture {
153 Capture
154}
155
156impl<O: Operand> Pattern<O> for Capture {
157 type Bindings = O;
158
159 fn try_match(&self, operand: &O) -> Option<Self::Bindings> {
160 Some(operand.clone())
161 }
162}
163
164pub struct Matching<C, P> {
165 patterns: P,
166 matched: PhantomData<fn() -> C>,
167}
168
169#[must_use]
170pub const fn matching<C, P>(patterns: P) -> Matching<C, P> {
171 Matching {
172 patterns,
173 matched: PhantomData,
174 }
175}
176
177impl<C, P, B> Pattern<C::Output> for Matching<C, P>
178where
179 C: PlanNode + MatchInputs + OptimizePlan,
180 P: for<'a> MatchAgainst<C::Inputs<'a>, Bindings = B>,
181{
182 type Bindings = B;
183
184 fn try_match(&self, operand: &C::Output) -> Option<Self::Bindings> {
185 let context = operand.as_plan_node().downcast::<C>()?;
186
187 self.patterns.match_against(MatchInputs::inputs(context))
188 }
189}
190
191impl<C, P> Matching<C, P> {
192 pub fn rewrite_matched<B, F>(self, rewrite: F) -> impl Rule<C::Output>
193 where
194 C: PlanNode + MatchInputs + OptimizePlan,
195 P: for<'a> MatchAgainst<C::Inputs<'a>, Bindings = B> + Send + Sync + 'static,
196 B: 'static,
197 F: Fn(&C, B, &Stats) -> Option<C::Output> + Send + Sync + 'static,
198 {
199 MatchingRewriteRule {
200 pattern: self,
201 rewrite,
202 }
203 }
204}
205
206struct MatchingRewriteRule<C, P, F> {
207 pattern: Matching<C, P>,
208 rewrite: F,
209}
210
211impl<C, P, B, F> Rule<C::Output> for MatchingRewriteRule<C, P, F>
212where
213 C: PlanNode + MatchInputs + OptimizePlan,
214 P: for<'a> MatchAgainst<C::Inputs<'a>, Bindings = B> + Send + Sync + 'static,
215 B: 'static,
216 F: Fn(&C, B, &Stats) -> Option<C::Output> + Send + Sync + 'static,
217{
218 fn apply(&self, operand: C::Output, stats: &Stats) -> Transformed<C::Output> {
219 let Some(context) = operand.as_plan_node().downcast::<C>() else {
220 return Transformed::unchanged(operand);
221 };
222
223 let Some(bindings) = self
224 .pattern
225 .patterns
226 .match_against(MatchInputs::inputs(context))
227 else {
228 return Transformed::unchanged(operand);
229 };
230
231 match (self.rewrite)(context, bindings, stats) {
232 Some(rewritten) => Transformed::changed(rewritten),
233 None => Transformed::unchanged(operand),
234 }
235 }
236}
237
238pub trait MatchAgainst<I> {
239 type Bindings;
240
241 fn match_against(&self, inputs: I) -> Option<Self::Bindings>;
242}
243
244impl MatchAgainst<()> for () {
245 type Bindings = ();
246
247 fn match_against(&self, _inputs: ()) -> Option<Self::Bindings> {
248 Some(())
249 }
250}
251
252macro_rules! impl_match_against {
253 ($($index:tt $operand:ident $pattern:ident),+) => {
254 impl<'inputs, $($operand,)+ $($pattern,)+> MatchAgainst<($(&'inputs $operand,)+)>
255 for ($($pattern,)+)
256 where
257 $($operand: Operand,)+
258 $($pattern: Pattern<$operand>,)+
259 {
260 type Bindings = ($(<$pattern as Pattern<$operand>>::Bindings,)+);
261
262 fn match_against(&self, inputs: ($(&$operand,)+)) -> Option<Self::Bindings> {
263 Some(($( self.$index.try_match(inputs.$index)?, )+))
264 }
265 }
266 };
267}
268
269impl_match_against!(0 O0 P0);
270impl_match_against!(0 O0 P0, 1 O1 P1);
271impl_match_against!(0 O0 P0, 1 O1 P1, 2 O2 P2);
272impl_match_against!(0 O0 P0, 1 O1 P1, 2 O2 P2, 3 O3 P3);
273impl_match_against!(0 O0 P0, 1 O1 P1, 2 O2 P2, 3 O3 P3, 4 O4 P4);
274impl_match_against!(0 O0 P0, 1 O1 P1, 2 O2 P2, 3 O3 P3, 4 O4 P4, 5 O5 P5);
275impl_match_against!(0 O0 P0, 1 O1 P1, 2 O2 P2, 3 O3 P3, 4 O4 P4, 5 O5 P5, 6 O6 P6);
276impl_match_against!(0 O0 P0, 1 O1 P1, 2 O2 P2, 3 O3 P3, 4 O4 P4, 5 O5 P5, 6 O6 P6, 7 O7 P7);
277impl_match_against!(
278 0 O0 P0, 1 O1 P1, 2 O2 P2, 3 O3 P3, 4 O4 P4, 5 O5 P5, 6 O6 P6, 7 O7 P7, 8 O8 P8
279);
280impl_match_against!(
281 0 O0 P0, 1 O1 P1, 2 O2 P2, 3 O3 P3, 4 O4 P4, 5 O5 P5, 6 O6 P6, 7 O7 P7, 8 O8 P8, 9 O9 P9
282);
283impl_match_against!(
284 0 O0 P0, 1 O1 P1, 2 O2 P2, 3 O3 P3, 4 O4 P4, 5 O5 P5, 6 O6 P6, 7 O7 P7, 8 O8 P8, 9 O9 P9,
285 10 O10 P10
286);
287impl_match_against!(
288 0 O0 P0, 1 O1 P1, 2 O2 P2, 3 O3 P3, 4 O4 P4, 5 O5 P5, 6 O6 P6, 7 O7 P7, 8 O8 P8, 9 O9 P9,
289 10 O10 P10, 11 O11 P11
290);