Skip to main content

linkrs_query/optimizer/heuristic/
rewrite_rule.rs

1//! Rewriting the definition of a trait
2//!
3//! This module provides the definition of a trait that encompasses heuristic rules for rewriting code.
4//! Heuristic rules do not rely on cost calculations and always produce either better or equivalent plans.
5//!
6//! The current implementation uses types that are independent of the planner layer and no longer relies on the optimizer module.
7
8use crate::optimizer::heuristic::context::RewriteContext;
9use crate::optimizer::heuristic::pattern::Pattern;
10use crate::optimizer::heuristic::result::{RewriteResult, TransformResult};
11use crate::optimizer::heuristic::rule::RewriteRule;
12use crate::planning::plan::PlanNodeEnum;
13
14/// Heuristic rewriting rule trait (compatibility layer)
15///
16/// A simplified interface provided for simple rewriting rules that do not require the full context of the optimizer.
17/// These rules always generate better or equivalent plans, without the need for cost calculations.
18///
19/// For rules that require access to the optimizer’s status or cost information, please use the `RewriteRule` trait.
20pub trait HeuristicRule: std::fmt::Debug + Send + Sync {
21    /// Rule Name
22    fn name(&self) -> &'static str;
23
24    /// Check whether it matches the current planned node.
25    fn matches(&self, node: &PlanNodeEnum) -> bool;
26
27    /// Apply the rule for rewriting the text.
28    ///
29    /// # Parameters
30    /// `ctx`: Rewrite the context.
31    /// `node`: The current planned node
32    ///
33    /// # Return
34    /// - `Ok(Some(node))`: rewrite succeeded, returns the new node
35    /// - `Ok(None)`: no match, keeps the original node
36    /// - `Err(e)`: rewrite failed
37    fn apply(
38        &self,
39        ctx: &mut RewriteContext,
40        node: &PlanNodeEnum,
41    ) -> RewriteResult<Option<PlanNodeEnum>>;
42}
43
44/// Create a wrapper that adapts `HeuristicRule` to `RewriteRule`.
45#[derive(Debug)]
46pub struct HeuristicRuleAdapter<T: HeuristicRule> {
47    inner: T,
48}
49
50impl<T: HeuristicRule> HeuristicRuleAdapter<T> {
51    pub fn new(rule: T) -> Self {
52        Self { inner: rule }
53    }
54
55    pub fn into_inner(self) -> T {
56        self.inner
57    }
58}
59
60impl<T: HeuristicRule> RewriteRule for HeuristicRuleAdapter<T> {
61    fn name(&self) -> &'static str {
62        self.inner.name()
63    }
64
65    fn pattern(&self) -> Pattern {
66        // Heuristic rules use wildcard patterns, and precise matching is performed by the matches method.
67        Pattern::new()
68    }
69
70    fn apply(
71        &self,
72        ctx: &mut RewriteContext,
73        node: &PlanNodeEnum,
74    ) -> RewriteResult<Option<TransformResult>> {
75        match self.inner.apply(ctx, node)? {
76            Some(new_node) => {
77                let mut result = TransformResult::new();
78                result.add_new_node(new_node);
79                Ok(Some(result))
80            }
81            None => Ok(None),
82        }
83    }
84
85    fn matches(&self, node: &PlanNodeEnum) -> bool {
86        self.inner.matches(node)
87    }
88}
89
90/// Provide a trait for the adapter constructor of HeuristicRule
91pub trait IntoOptRule: HeuristicRule + Sized {
92    fn into_opt_rule(self) -> HeuristicRuleAdapter<Self> {
93        HeuristicRuleAdapter::new(self)
94    }
95}
96
97impl<T: HeuristicRule + Sized> IntoOptRule for T {}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use crate::planning::plan::core::nodes::access::graph_scan_node::ScanVerticesNode;
103
104    #[derive(Debug)]
105    struct TestHeuristicRule;
106
107    impl HeuristicRule for TestHeuristicRule {
108        fn name(&self) -> &'static str {
109            "TestHeuristicRule"
110        }
111
112        fn matches(&self, node: &PlanNodeEnum) -> bool {
113            node.is_scan_vertices()
114        }
115
116        fn apply(
117            &self,
118            _ctx: &mut RewriteContext,
119            node: &PlanNodeEnum,
120        ) -> RewriteResult<Option<PlanNodeEnum>> {
121            if self.matches(node) {
122                // Return to the original node (the actual rule would involve more complex logic.)
123                Ok(Some(node.clone()))
124            } else {
125                Ok(None)
126            }
127        }
128    }
129
130    #[test]
131    fn test_heuristic_rule_adapter() {
132        let rule = TestHeuristicRule;
133        let adapter = rule.into_opt_rule();
134
135        assert_eq!(adapter.name(), "TestHeuristicRule");
136
137        let node = PlanNodeEnum::ScanVertices(ScanVerticesNode::new(1, "default"));
138        assert!(adapter.matches(&node));
139    }
140
141    #[test]
142    fn test_heuristic_rule_apply() {
143        let rule = TestHeuristicRule;
144        let mut ctx = RewriteContext::new();
145        let node = PlanNodeEnum::ScanVertices(ScanVerticesNode::new(1, "default"));
146
147        let result = rule
148            .apply(&mut ctx, &node)
149            .expect("Failed to apply rewrite rule");
150        assert!(result.is_some());
151    }
152}