Skip to main content

graphrecords_query/optimizer/
rule.rs

1use super::{engine::Session, plan::PlanNode, stats::Stats};
2use crate::Operand;
3use std::marker::PhantomData;
4
5pub(super) type ErasedRule<O> =
6    Box<dyn for<'a> Fn(O, &Session<'a>) -> Transformed<O> + Send + Sync>;
7
8pub struct Transformed<T> {
9    value: T,
10    changed: bool,
11}
12
13impl<T> Transformed<T> {
14    #[must_use]
15    pub const fn changed(value: T) -> Self {
16        Self {
17            value,
18            changed: true,
19        }
20    }
21
22    #[must_use]
23    pub const fn unchanged(value: T) -> Self {
24        Self {
25            value,
26            changed: false,
27        }
28    }
29
30    #[must_use]
31    pub const fn value(&self) -> &T {
32        &self.value
33    }
34
35    #[must_use]
36    pub const fn is_changed(&self) -> bool {
37        self.changed
38    }
39
40    #[must_use]
41    pub fn into_parts(self) -> (T, bool) {
42        (self.value, self.changed)
43    }
44}
45
46pub trait Rule<O: Operand>: 'static + Send + Sync {
47    fn apply(&self, operand: O, stats: &Stats) -> Transformed<O>;
48}
49
50#[must_use]
51pub fn rule<C, O, F>(rewrite: F) -> impl Rule<O>
52where
53    C: PlanNode,
54    O: Operand + 'static,
55    F: Fn(&C, &Stats) -> Option<O> + Send + Sync + 'static,
56{
57    ContextRule {
58        rewrite,
59        matched: PhantomData,
60    }
61}
62
63struct ContextRule<C, O, F> {
64    rewrite: F,
65    matched: PhantomData<fn() -> (C, O)>,
66}
67
68impl<C, O, F> Rule<O> for ContextRule<C, O, F>
69where
70    C: PlanNode,
71    O: Operand + 'static,
72    F: Fn(&C, &Stats) -> Option<O> + Send + Sync + 'static,
73{
74    fn apply(&self, operand: O, stats: &Stats) -> Transformed<O> {
75        let Some(context) = operand.as_plan_node().downcast::<C>() else {
76            return Transformed::unchanged(operand);
77        };
78
79        match (self.rewrite)(context, stats) {
80            Some(rewritten) => Transformed::changed(rewritten),
81            None => Transformed::unchanged(operand),
82        }
83    }
84}