linkrs_query/optimizer/heuristic/
rewrite_rule.rs1use 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
14pub trait HeuristicRule: std::fmt::Debug + Send + Sync {
21 fn name(&self) -> &'static str;
23
24 fn matches(&self, node: &PlanNodeEnum) -> bool;
26
27 fn apply(
38 &self,
39 ctx: &mut RewriteContext,
40 node: &PlanNodeEnum,
41 ) -> RewriteResult<Option<PlanNodeEnum>>;
42}
43
44#[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 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
90pub 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 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}