Skip to main content

linkrs_query/optimizer/heuristic/
rule_enum.rs

1//! Rewrite rule enumeration – Implementation of static distribution
2//!
3//! This module uses enumerations to implement static distribution, thereby avoiding the overhead associated with dynamic distribution.
4//! All rules are distributed in the form of enumeration variants, using the `match` mechanism.
5//!
6//! # Advantages
7//!
8//! No dynamic distribution overhead (no lookup in virtual function tables).
9//! No heap allocation (the rules are stored on the stack).
10//! Better cache locality
11//! Compilers can perform inlining optimizations.
12//!
13//! # Usage Examples
14//!
15//! ```rust
16//! use linkrs_query::optimizer::heuristic::rule_enum::RuleRegistry;
17//!
18//! // Create the rule registry
19//! let registry = RuleRegistry::default();
20//!
21//! // Iterate over the registered rules
22//! let rule_count = registry.iter().count();
23//! assert!(rule_count > 0);
24//! ```
25
26use crate::optimizer::heuristic::aggregate;
27use crate::optimizer::heuristic::constant_folding;
28use crate::optimizer::heuristic::context::RewriteContext;
29use crate::optimizer::heuristic::decorrelation;
30use crate::optimizer::heuristic::elimination;
31use crate::optimizer::heuristic::expand_pushdown;
32use crate::optimizer::heuristic::join_optimization;
33use crate::optimizer::heuristic::limit_pushdown;
34use crate::optimizer::heuristic::merge;
35use crate::optimizer::heuristic::pattern::Pattern;
36use crate::optimizer::heuristic::predicate_pushdown;
37use crate::optimizer::heuristic::projection_pushdown;
38use crate::optimizer::heuristic::result::{RewriteResult, TransformResult};
39use crate::optimizer::heuristic::rule::RewriteRule as RewriteRuleTrait;
40use crate::optimizer::heuristic::scan_edge_identity;
41use crate::optimizer::heuristic::scan_identity;
42use crate::optimizer::heuristic::slot_coverage;
43use crate::planning::plan::PlanNodeEnum;
44
45macro_rules! define_rewrite_rules {
46    (
47        $(#[$enum_meta:meta])*
48        pub enum $enum_name:ident {
49            $(
50                $(#[$variant_meta:meta])*
51                $variant_name:ident($rule_type:ty)
52            ),+ $(,)?
53        }
54    ) => {
55        $(#[$enum_meta])*
56        #[derive(Debug)]
57        pub enum $enum_name {
58            $(
59                $(#[$variant_meta])*
60                $variant_name($rule_type),
61            )+
62        }
63
64        impl $enum_name {
65            pub fn name(&self) -> &'static str {
66                match self {
67                    $(
68                        $enum_name::$variant_name(_) => {
69                            let type_name = stringify!($rule_type);
70                            if let Some(pos) = type_name.rfind("::") {
71                                &type_name[pos + 2..]
72                            } else {
73                                type_name
74                            }
75                        }
76                    )+
77                }
78            }
79
80            pub fn pattern(&self) -> Pattern {
81                match self {
82                    $(
83                        $enum_name::$variant_name(rule) => rule.pattern(),
84                    )+
85                }
86            }
87
88            pub fn apply(
89                &self,
90                ctx: &mut RewriteContext,
91                node: &PlanNodeEnum,
92            ) -> RewriteResult<Option<TransformResult>> {
93                match self {
94                    $(
95                        $enum_name::$variant_name(rule) => rule.apply(ctx, node),
96                    )+
97                }
98            }
99
100            pub fn matches(&self, node: &PlanNodeEnum) -> bool {
101                self.pattern().matches(node)
102            }
103        }
104
105        impl RewriteRuleTrait for $enum_name {
106            fn name(&self) -> &'static str {
107                self.name()
108            }
109
110            fn pattern(&self) -> Pattern {
111                self.pattern()
112            }
113
114            fn apply(
115                &self,
116                ctx: &mut RewriteContext,
117                node: &PlanNodeEnum,
118            ) -> RewriteResult<Option<TransformResult>> {
119                self.apply(ctx, node)
120            }
121        }
122    };
123}
124
125define_rewrite_rules! {
126    pub enum RewriteRule {
127        // ==================== Remove Rules ====================
128        EliminateFilter(elimination::EliminateFilterRule),
129        RemoveNoopProject(elimination::RemoveNoopProjectRule),
130        EliminateAppendVertices(elimination::EliminateAppendVerticesRule),
131        RemoveAppendVerticesBelowJoin(elimination::RemoveAppendVerticesBelowJoinRule),
132        EliminateRowCollect(elimination::EliminateRowCollectRule),
133        EliminateEmptySetOperation(elimination::EliminateEmptySetOperationRule),
134        DedupElimination(elimination::DedupEliminationRule),
135        EliminateSort(elimination::EliminateSortRule),
136
137        // ==================== Constant Folding Rule ====================
138        FoldConstants(constant_folding::FoldConstantsRule),
139
140        // ==================== Merging Rules ====================
141        CombineFilter(merge::CombineFilterRule),
142        CollapseProject(merge::CollapseProjectRule),
143        CollapseConsecutiveProject(merge::CollapseConsecutiveProjectRule),
144        MergeGetVerticesAndProject(merge::MergeGetVerticesAndProjectRule),
145        MergeGetVerticesAndDedup(merge::MergeGetVerticesAndDedupRule),
146        MergeGetNbrsAndProject(merge::MergeGetNbrsAndProjectRule),
147        MergeGetNbrsAndDedup(merge::MergeGetNbrsAndDedupRule),
148
149        // ==================== Predicate Pushdown Rules ====================
150        PushFilterDownTraverse(predicate_pushdown::PushFilterDownTraverseRule),
151        PushFilterDownExpandAll(predicate_pushdown::PushFilterDownExpandAllRule),
152        PushFilterDownNode(predicate_pushdown::PushFilterDownNodeRule),
153        PushEFilterDown(predicate_pushdown::PushEFilterDownRule),
154        PushVFilterDownScanVertices(predicate_pushdown::PushVFilterDownScanVerticesRule),
155        PushFilterDownScanVertices(predicate_pushdown::PushFilterDownScanVerticesRule),
156        PushFilterDownScanEdges(predicate_pushdown::PushFilterDownScanEdgesRule),
157        EliminateRedundantTagFilter(predicate_pushdown::EliminateRedundantTagFilterRule),
158        PushFilterDownInnerJoin(predicate_pushdown::PushFilterDownInnerJoinRule),
159        PushFilterDownLeftJoin(predicate_pushdown::PushFilterDownLeftJoinRule),
160        PushFilterDownCrossJoin(predicate_pushdown::PushFilterDownCrossJoinRule),
161        PushFilterDownGetNbrs(predicate_pushdown::PushFilterDownGetNbrsRule),
162        PushFilterDownAllPaths(predicate_pushdown::PushFilterDownAllPathsRule),
163
164        // ==================== Projection Pushdown Rules ====================
165        // NOTE: GetVertices / GetNeighbors / GetEdges rules narrow the graph
166        // operator's `projected_properties` through typed required-property
167        // analysis (optimizer/analysis/required_properties.rs). The Project
168        // node is preserved; only the columns read by the graph operator are
169        // narrowed, which keeps computed expressions, aliases, and functions
170        // sound (bare / opaque variable usage blocks the rewrite).
171        PushProjectDownGetVertices(projection_pushdown::PushProjectDownGetVerticesRule),
172        PushProjectDownGetNeighbors(projection_pushdown::PushProjectDownGetNeighborsRule),
173        PushProjectDownGetEdges(projection_pushdown::PushProjectDownGetEdgesRule),
174        PushProjectDownAppendVertices(projection_pushdown::PushProjectDownAppendVerticesRule),
175        PushProjectDownScanVertices(projection_pushdown::PushProjectDownScanVerticesRule),
176        PushProjectDownScanEdges(projection_pushdown::PushProjectDownScanEdgesRule),
177        // Enriches scan `projected_properties` with residual filter
178        // predicate columns so the columnar evaluator can serve WHERE
179        // clauses directly above scan sources.
180        EnrichScanSlotsWithFilterProps(slot_coverage::EnrichScanSlotsWithFilterPropsRule),
181
182        // ==================== Rules for Pushing Limits Down ====================
183        PushLimitDownGetVertices(limit_pushdown::PushLimitDownGetVerticesRule),
184        PushLimitDownGetEdges(limit_pushdown::PushLimitDownGetEdgesRule),
185        PushLimitDownScanVertices(limit_pushdown::PushLimitDownScanVerticesRule),
186        PushLimitDownScanEdges(limit_pushdown::PushLimitDownScanEdgesRule),
187        PushLimitDownIndexScan(limit_pushdown::PushLimitDownIndexScanRule),
188        PushTopNDownIndexScan(limit_pushdown::PushTopNDownIndexScanRule),
189        ConvertSortLimitToTopN(limit_pushdown::ConvertSortLimitToTopNRule),
190
191        // ==================== Aggregation Optimization Rules ====================
192        PushFilterDownAggregate(aggregate::PushFilterDownAggregateRule),
193
194        // ==================== Expand Pushdown Annotation Rules ====================
195        // Runs after PredicatePushdown so filters are already pushed into the
196        // scans/expands; annotates ExpandAll hops with id_only/count_only so
197        // the executor can skip full vertex/edge materialization.
198        ExpandPushdownAnnotate(expand_pushdown::ExpandPushdownAnnotateRule),
199        // Runs after ExpandPushdownAnnotate so expand seed tolerance is
200        // final; annotates ScanVertices with identity_only so the executor
201        // can emit lightweight vertex ids instead of boxed vertices.
202        ScanIdentityAnnotate(scan_identity::ScanIdentityAnnotateRule),
203        ScanEdgeIdentityAnnotate(scan_edge_identity::ScanEdgeIdentityAnnotateRule),
204
205        // ==================== Decorrelation Rules ====================
206        // Stat-free gate: converts simple deterministic PatternApply
207        // subqueries to SemiJoin/AntiJoin. Heavy cost-based judgment stays
208        // in the CBO phase (cost_based/subquery_unnesting.rs).
209        UnnestSimplePatternApply(decorrelation::UnnestSimplePatternApplyRule),
210
211        // ==================== JOIN Optimization Rules ====================
212        PushProjectDownJoin(join_optimization::PushProjectDownJoinRule),
213        ExtractEquiJoinKeys(join_optimization::ExtractEquiJoinKeysRule),
214        LeftJoinToInnerJoin(join_optimization::LeftJoinToInnerJoinRule),
215        JoinConditionSimplify(join_optimization::JoinConditionSimplifyRule),
216        JoinToExpand(join_optimization::JoinToExpandRule),
217        JoinToAppendVertices(join_optimization::JoinToAppendVerticesRule),
218        MergeConsecutiveExpand(join_optimization::MergeConsecutiveExpandRule),
219        JoinElimination(join_optimization::JoinEliminationRule),
220        IndexJoinSelection(join_optimization::IndexJoinSelectionRule),
221        JoinReorder(join_optimization::JoinReorderRule),
222    }
223}
224
225#[derive(Debug)]
226pub struct RuleRegistry {
227    rules: Vec<RewriteRule>,
228}
229
230impl RuleRegistry {
231    pub fn new() -> Self {
232        Self { rules: Vec::new() }
233    }
234
235    pub fn add(&mut self, rule: RewriteRule) {
236        self.rules.push(rule);
237    }
238
239    pub fn iter(&self) -> impl Iterator<Item = &RewriteRule> {
240        self.rules.iter()
241    }
242
243    pub fn len(&self) -> usize {
244        self.rules.len()
245    }
246
247    pub fn is_empty(&self) -> bool {
248        self.rules.is_empty()
249    }
250
251    pub fn clear(&mut self) {
252        self.rules.clear();
253    }
254
255    pub fn into_vec(self) -> Vec<RewriteRule> {
256        self.rules
257    }
258}
259
260impl Default for RuleRegistry {
261    fn default() -> Self {
262        let mut registry = Self::new();
263        registry.add(RewriteRule::EliminateFilter(
264            elimination::EliminateFilterRule::new(),
265        ));
266        registry.add(RewriteRule::FoldConstants(
267            constant_folding::FoldConstantsRule::new(),
268        ));
269        registry.add(RewriteRule::RemoveNoopProject(
270            elimination::RemoveNoopProjectRule::new(),
271        ));
272        registry.add(RewriteRule::EliminateAppendVertices(
273            elimination::EliminateAppendVerticesRule::new(),
274        ));
275        registry.add(RewriteRule::RemoveAppendVerticesBelowJoin(
276            elimination::RemoveAppendVerticesBelowJoinRule::new(),
277        ));
278        registry.add(RewriteRule::EliminateRowCollect(
279            elimination::EliminateRowCollectRule::new(),
280        ));
281        registry.add(RewriteRule::EliminateEmptySetOperation(
282            elimination::EliminateEmptySetOperationRule::new(),
283        ));
284        registry.add(RewriteRule::DedupElimination(
285            elimination::DedupEliminationRule::new(),
286        ));
287        registry.add(RewriteRule::EliminateSort(
288            elimination::EliminateSortRule::new(),
289        ));
290        registry.add(RewriteRule::CombineFilter(merge::CombineFilterRule::new()));
291        registry.add(RewriteRule::CollapseProject(
292            merge::CollapseProjectRule::new(),
293        ));
294        registry.add(RewriteRule::CollapseConsecutiveProject(
295            merge::CollapseConsecutiveProjectRule::new(),
296        ));
297        registry.add(RewriteRule::MergeGetVerticesAndProject(
298            merge::MergeGetVerticesAndProjectRule::new(),
299        ));
300        registry.add(RewriteRule::MergeGetVerticesAndDedup(
301            merge::MergeGetVerticesAndDedupRule::new(),
302        ));
303        registry.add(RewriteRule::MergeGetNbrsAndProject(
304            merge::MergeGetNbrsAndProjectRule::new(),
305        ));
306        registry.add(RewriteRule::MergeGetNbrsAndDedup(
307            merge::MergeGetNbrsAndDedupRule::new(),
308        ));
309        registry.add(RewriteRule::PushFilterDownTraverse(
310            predicate_pushdown::PushFilterDownTraverseRule::new(),
311        ));
312        registry.add(RewriteRule::PushFilterDownExpandAll(
313            predicate_pushdown::PushFilterDownExpandAllRule::new(),
314        ));
315        registry.add(RewriteRule::PushFilterDownNode(
316            predicate_pushdown::PushFilterDownNodeRule::new(),
317        ));
318        registry.add(RewriteRule::PushEFilterDown(
319            predicate_pushdown::PushEFilterDownRule::new(),
320        ));
321        registry.add(RewriteRule::PushVFilterDownScanVertices(
322            predicate_pushdown::PushVFilterDownScanVerticesRule::new(),
323        ));
324        registry.add(RewriteRule::EliminateRedundantTagFilter(
325            predicate_pushdown::EliminateRedundantTagFilterRule::new(),
326        ));
327        registry.add(RewriteRule::PushFilterDownScanVertices(
328            predicate_pushdown::PushFilterDownScanVerticesRule::new(),
329        ));
330        registry.add(RewriteRule::PushFilterDownScanEdges(
331            predicate_pushdown::PushFilterDownScanEdgesRule::new(),
332        ));
333        registry.add(RewriteRule::PushFilterDownInnerJoin(
334            predicate_pushdown::PushFilterDownInnerJoinRule::new(),
335        ));
336        registry.add(RewriteRule::PushFilterDownLeftJoin(
337            predicate_pushdown::PushFilterDownLeftJoinRule::new(),
338        ));
339        registry.add(RewriteRule::PushFilterDownCrossJoin(
340            predicate_pushdown::PushFilterDownCrossJoinRule::new(),
341        ));
342        registry.add(RewriteRule::PushFilterDownGetNbrs(
343            predicate_pushdown::PushFilterDownGetNbrsRule::new(),
344        ));
345        registry.add(RewriteRule::PushFilterDownAllPaths(
346            predicate_pushdown::PushFilterDownAllPathsRule::new(),
347        ));
348        registry.add(RewriteRule::PushProjectDownScanVertices(
349            projection_pushdown::PushProjectDownScanVerticesRule::new(),
350        ));
351        registry.add(RewriteRule::PushProjectDownScanEdges(
352            projection_pushdown::PushProjectDownScanEdgesRule::new(),
353        ));
354        registry.add(RewriteRule::PushProjectDownGetVertices(
355            projection_pushdown::PushProjectDownGetVerticesRule::new(),
356        ));
357        registry.add(RewriteRule::PushProjectDownGetNeighbors(
358            projection_pushdown::PushProjectDownGetNeighborsRule::new(),
359        ));
360        registry.add(RewriteRule::PushProjectDownGetEdges(
361            projection_pushdown::PushProjectDownGetEdgesRule::new(),
362        ));
363        registry.add(RewriteRule::PushProjectDownAppendVertices(
364            projection_pushdown::PushProjectDownAppendVerticesRule::new(),
365        ));
366        registry.add(RewriteRule::EnrichScanSlotsWithFilterProps(
367            slot_coverage::EnrichScanSlotsWithFilterPropsRule::new(),
368        ));
369        registry.add(RewriteRule::PushLimitDownGetVertices(
370            limit_pushdown::PushLimitDownGetVerticesRule::new(),
371        ));
372        registry.add(RewriteRule::PushLimitDownGetEdges(
373            limit_pushdown::PushLimitDownGetEdgesRule::new(),
374        ));
375        registry.add(RewriteRule::PushLimitDownScanVertices(
376            limit_pushdown::PushLimitDownScanVerticesRule::new(),
377        ));
378        registry.add(RewriteRule::PushLimitDownScanEdges(
379            limit_pushdown::PushLimitDownScanEdgesRule::new(),
380        ));
381        registry.add(RewriteRule::PushLimitDownIndexScan(
382            limit_pushdown::PushLimitDownIndexScanRule::new(),
383        ));
384        registry.add(RewriteRule::PushTopNDownIndexScan(
385            limit_pushdown::PushTopNDownIndexScanRule::new(),
386        ));
387        registry.add(RewriteRule::ConvertSortLimitToTopN(
388            limit_pushdown::ConvertSortLimitToTopNRule::new(),
389        ));
390        registry.add(RewriteRule::PushFilterDownAggregate(
391            aggregate::PushFilterDownAggregateRule::new(),
392        ));
393        registry.add(RewriteRule::ExpandPushdownAnnotate(
394            expand_pushdown::ExpandPushdownAnnotateRule::new(),
395        ));
396        registry.add(RewriteRule::ScanIdentityAnnotate(
397            scan_identity::ScanIdentityAnnotateRule::new(),
398        ));
399        registry.add(RewriteRule::ScanEdgeIdentityAnnotate(
400            scan_edge_identity::ScanEdgeIdentityAnnotateRule::new(),
401        ));
402        registry.add(RewriteRule::UnnestSimplePatternApply(
403            decorrelation::UnnestSimplePatternApplyRule::new(),
404        ));
405        registry.add(RewriteRule::PushProjectDownJoin(
406            join_optimization::PushProjectDownJoinRule::new(),
407        ));
408        registry.add(RewriteRule::ExtractEquiJoinKeys(
409            join_optimization::ExtractEquiJoinKeysRule::new(),
410        ));
411        registry.add(RewriteRule::LeftJoinToInnerJoin(
412            join_optimization::LeftJoinToInnerJoinRule::new(),
413        ));
414        registry.add(RewriteRule::JoinConditionSimplify(
415            join_optimization::JoinConditionSimplifyRule::new(),
416        ));
417        registry.add(RewriteRule::JoinToExpand(
418            join_optimization::JoinToExpandRule::new(),
419        ));
420        registry.add(RewriteRule::JoinToAppendVertices(
421            join_optimization::JoinToAppendVerticesRule::new(),
422        ));
423        registry.add(RewriteRule::MergeConsecutiveExpand(
424            join_optimization::MergeConsecutiveExpandRule::new(),
425        ));
426        registry.add(RewriteRule::JoinElimination(
427            join_optimization::JoinEliminationRule::new(),
428        ));
429        registry.add(RewriteRule::IndexJoinSelection(
430            join_optimization::IndexJoinSelectionRule::new(),
431        ));
432        registry.add(RewriteRule::JoinReorder(
433            join_optimization::JoinReorderRule::new(),
434        ));
435        registry
436    }
437}
438
439impl RuleRegistry {
440    /// Build a registry filtered by the `[optimizer.rules]` selection.
441    ///
442    /// `enabled` non-empty keeps only those rules (explicit selection wins
443    /// over the default rule set); `disabled` then removes rules. Names match
444    /// rule type names with or without the `Rule` suffix.
445    pub fn filtered(disabled: &[String], enabled: &[String]) -> Self {
446        fn normalise(name: &str) -> &str {
447            name.strip_suffix("Rule").unwrap_or(name)
448        }
449
450        let base = Self::default();
451        if disabled.is_empty() && enabled.is_empty() {
452            return base;
453        }
454        let enabled_set: std::collections::HashSet<&str> =
455            enabled.iter().map(|n| normalise(n)).collect();
456        let disabled_set: std::collections::HashSet<&str> =
457            disabled.iter().map(|n| normalise(n)).collect();
458        let mut registry = Self::new();
459        for rule in base.rules {
460            let name = normalise(rule.name());
461            if !enabled_set.is_empty() && !enabled_set.contains(name) {
462                continue;
463            }
464            if disabled_set.contains(name) {
465                continue;
466            }
467            registry.add(rule);
468        }
469        registry
470    }
471}
472
473#[cfg(test)]
474mod tests {
475    use super::*;
476
477    #[test]
478    fn test_rule_registry_default() {
479        let registry = RuleRegistry::default();
480        assert_eq!(registry.len(), 58);
481    }
482
483    #[test]
484    fn test_rule_names() {
485        let registry = RuleRegistry::default();
486        for rule in registry.iter() {
487            let name = rule.name();
488            assert!(!name.is_empty());
489            assert!(name.ends_with("Rule"));
490        }
491    }
492
493    #[test]
494    fn test_filtered_disabled_removes_rule() {
495        let registry = RuleRegistry::filtered(&["LeftJoinToInnerJoin".to_string()], &[]);
496        assert_eq!(registry.len(), 57);
497        assert!(registry
498            .iter()
499            .all(|rule| rule.name() != "LeftJoinToInnerJoinRule"));
500    }
501
502    #[test]
503    fn test_filtered_enabled_keeps_only_selected() {
504        let registry = RuleRegistry::filtered(
505            &[],
506            &[
507                "EliminateFilterRule".to_string(),
508                "CombineFilter".to_string(),
509            ],
510        );
511        assert_eq!(registry.len(), 2);
512        let names: Vec<&str> = registry.iter().map(|rule| rule.name()).collect();
513        assert!(names.contains(&"EliminateFilterRule"));
514        assert!(names.contains(&"CombineFilterRule"));
515    }
516
517    #[test]
518    fn test_filtered_unknown_names_select_nothing() {
519        let registry = RuleRegistry::filtered(&[], &["NoSuchRule".to_string()]);
520        assert!(registry.is_empty());
521    }
522}