Skip to main content

graphrecords_query/optimizer/
pattern.rs

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);