linkrs-query 0.1.0

//! Rewrite rule enumeration – Implementation of static distribution
//!
//! This module uses enumerations to implement static distribution, thereby avoiding the overhead associated with dynamic distribution.
//! All rules are distributed in the form of enumeration variants, using the `match` mechanism.
//!
//! # Advantages
//!
//! No dynamic distribution overhead (no lookup in virtual function tables).
//! No heap allocation (the rules are stored on the stack).
//! Better cache locality
//! Compilers can perform inlining optimizations.
//!
//! # Usage Examples
//!
//! ```rust
//! use linkrs_query::optimizer::heuristic::rule_enum::RuleRegistry;
//!
//! // Create the rule registry
//! let registry = RuleRegistry::default();
//!
//! // Iterate over the registered rules
//! let rule_count = registry.iter().count();
//! assert!(rule_count > 0);
//! ```

use crate::optimizer::heuristic::aggregate;
use crate::optimizer::heuristic::constant_folding;
use crate::optimizer::heuristic::context::RewriteContext;
use crate::optimizer::heuristic::decorrelation;
use crate::optimizer::heuristic::elimination;
use crate::optimizer::heuristic::expand_pushdown;
use crate::optimizer::heuristic::join_optimization;
use crate::optimizer::heuristic::limit_pushdown;
use crate::optimizer::heuristic::merge;
use crate::optimizer::heuristic::pattern::Pattern;
use crate::optimizer::heuristic::predicate_pushdown;
use crate::optimizer::heuristic::projection_pushdown;
use crate::optimizer::heuristic::result::{RewriteResult, TransformResult};
use crate::optimizer::heuristic::rule::RewriteRule as RewriteRuleTrait;
use crate::optimizer::heuristic::scan_edge_identity;
use crate::optimizer::heuristic::scan_identity;
use crate::optimizer::heuristic::slot_coverage;
use crate::planning::plan::PlanNodeEnum;

macro_rules! define_rewrite_rules {
    (
        $(#[$enum_meta:meta])*
        pub enum $enum_name:ident {
            $(
                $(#[$variant_meta:meta])*
                $variant_name:ident($rule_type:ty)
            ),+ $(,)?
        }
    ) => {
        $(#[$enum_meta])*
        #[derive(Debug)]
        pub enum $enum_name {
            $(
                $(#[$variant_meta])*
                $variant_name($rule_type),
            )+
        }

        impl $enum_name {
            pub fn name(&self) -> &'static str {
                match self {
                    $(
                        $enum_name::$variant_name(_) => {
                            let type_name = stringify!($rule_type);
                            if let Some(pos) = type_name.rfind("::") {
                                &type_name[pos + 2..]
                            } else {
                                type_name
                            }
                        }
                    )+
                }
            }

            pub fn pattern(&self) -> Pattern {
                match self {
                    $(
                        $enum_name::$variant_name(rule) => rule.pattern(),
                    )+
                }
            }

            pub fn apply(
                &self,
                ctx: &mut RewriteContext,
                node: &PlanNodeEnum,
            ) -> RewriteResult<Option<TransformResult>> {
                match self {
                    $(
                        $enum_name::$variant_name(rule) => rule.apply(ctx, node),
                    )+
                }
            }

            pub fn matches(&self, node: &PlanNodeEnum) -> bool {
                self.pattern().matches(node)
            }
        }

        impl RewriteRuleTrait for $enum_name {
            fn name(&self) -> &'static str {
                self.name()
            }

            fn pattern(&self) -> Pattern {
                self.pattern()
            }

            fn apply(
                &self,
                ctx: &mut RewriteContext,
                node: &PlanNodeEnum,
            ) -> RewriteResult<Option<TransformResult>> {
                self.apply(ctx, node)
            }
        }
    };
}

define_rewrite_rules! {
    pub enum RewriteRule {
        // ==================== Remove Rules ====================
        EliminateFilter(elimination::EliminateFilterRule),
        RemoveNoopProject(elimination::RemoveNoopProjectRule),
        EliminateAppendVertices(elimination::EliminateAppendVerticesRule),
        RemoveAppendVerticesBelowJoin(elimination::RemoveAppendVerticesBelowJoinRule),
        EliminateRowCollect(elimination::EliminateRowCollectRule),
        EliminateEmptySetOperation(elimination::EliminateEmptySetOperationRule),
        DedupElimination(elimination::DedupEliminationRule),
        EliminateSort(elimination::EliminateSortRule),

        // ==================== Constant Folding Rule ====================
        FoldConstants(constant_folding::FoldConstantsRule),

        // ==================== Merging Rules ====================
        CombineFilter(merge::CombineFilterRule),
        CollapseProject(merge::CollapseProjectRule),
        CollapseConsecutiveProject(merge::CollapseConsecutiveProjectRule),
        MergeGetVerticesAndProject(merge::MergeGetVerticesAndProjectRule),
        MergeGetVerticesAndDedup(merge::MergeGetVerticesAndDedupRule),
        MergeGetNbrsAndProject(merge::MergeGetNbrsAndProjectRule),
        MergeGetNbrsAndDedup(merge::MergeGetNbrsAndDedupRule),

        // ==================== Predicate Pushdown Rules ====================
        PushFilterDownTraverse(predicate_pushdown::PushFilterDownTraverseRule),
        PushFilterDownExpandAll(predicate_pushdown::PushFilterDownExpandAllRule),
        PushFilterDownNode(predicate_pushdown::PushFilterDownNodeRule),
        PushEFilterDown(predicate_pushdown::PushEFilterDownRule),
        PushVFilterDownScanVertices(predicate_pushdown::PushVFilterDownScanVerticesRule),
        PushFilterDownScanVertices(predicate_pushdown::PushFilterDownScanVerticesRule),
        PushFilterDownScanEdges(predicate_pushdown::PushFilterDownScanEdgesRule),
        EliminateRedundantTagFilter(predicate_pushdown::EliminateRedundantTagFilterRule),
        PushFilterDownInnerJoin(predicate_pushdown::PushFilterDownInnerJoinRule),
        PushFilterDownLeftJoin(predicate_pushdown::PushFilterDownLeftJoinRule),
        PushFilterDownCrossJoin(predicate_pushdown::PushFilterDownCrossJoinRule),
        PushFilterDownGetNbrs(predicate_pushdown::PushFilterDownGetNbrsRule),
        PushFilterDownAllPaths(predicate_pushdown::PushFilterDownAllPathsRule),

        // ==================== Projection Pushdown Rules ====================
        // NOTE: GetVertices / GetNeighbors / GetEdges rules narrow the graph
        // operator's `projected_properties` through typed required-property
        // analysis (optimizer/analysis/required_properties.rs). The Project
        // node is preserved; only the columns read by the graph operator are
        // narrowed, which keeps computed expressions, aliases, and functions
        // sound (bare / opaque variable usage blocks the rewrite).
        PushProjectDownGetVertices(projection_pushdown::PushProjectDownGetVerticesRule),
        PushProjectDownGetNeighbors(projection_pushdown::PushProjectDownGetNeighborsRule),
        PushProjectDownGetEdges(projection_pushdown::PushProjectDownGetEdgesRule),
        PushProjectDownAppendVertices(projection_pushdown::PushProjectDownAppendVerticesRule),
        PushProjectDownScanVertices(projection_pushdown::PushProjectDownScanVerticesRule),
        PushProjectDownScanEdges(projection_pushdown::PushProjectDownScanEdgesRule),
        // Enriches scan `projected_properties` with residual filter
        // predicate columns so the columnar evaluator can serve WHERE
        // clauses directly above scan sources.
        EnrichScanSlotsWithFilterProps(slot_coverage::EnrichScanSlotsWithFilterPropsRule),

        // ==================== Rules for Pushing Limits Down ====================
        PushLimitDownGetVertices(limit_pushdown::PushLimitDownGetVerticesRule),
        PushLimitDownGetEdges(limit_pushdown::PushLimitDownGetEdgesRule),
        PushLimitDownScanVertices(limit_pushdown::PushLimitDownScanVerticesRule),
        PushLimitDownScanEdges(limit_pushdown::PushLimitDownScanEdgesRule),
        PushLimitDownIndexScan(limit_pushdown::PushLimitDownIndexScanRule),
        PushTopNDownIndexScan(limit_pushdown::PushTopNDownIndexScanRule),
        ConvertSortLimitToTopN(limit_pushdown::ConvertSortLimitToTopNRule),

        // ==================== Aggregation Optimization Rules ====================
        PushFilterDownAggregate(aggregate::PushFilterDownAggregateRule),

        // ==================== Expand Pushdown Annotation Rules ====================
        // Runs after PredicatePushdown so filters are already pushed into the
        // scans/expands; annotates ExpandAll hops with id_only/count_only so
        // the executor can skip full vertex/edge materialization.
        ExpandPushdownAnnotate(expand_pushdown::ExpandPushdownAnnotateRule),
        // Runs after ExpandPushdownAnnotate so expand seed tolerance is
        // final; annotates ScanVertices with identity_only so the executor
        // can emit lightweight vertex ids instead of boxed vertices.
        ScanIdentityAnnotate(scan_identity::ScanIdentityAnnotateRule),
        ScanEdgeIdentityAnnotate(scan_edge_identity::ScanEdgeIdentityAnnotateRule),

        // ==================== Decorrelation Rules ====================
        // Stat-free gate: converts simple deterministic PatternApply
        // subqueries to SemiJoin/AntiJoin. Heavy cost-based judgment stays
        // in the CBO phase (cost_based/subquery_unnesting.rs).
        UnnestSimplePatternApply(decorrelation::UnnestSimplePatternApplyRule),

        // ==================== JOIN Optimization Rules ====================
        PushProjectDownJoin(join_optimization::PushProjectDownJoinRule),
        ExtractEquiJoinKeys(join_optimization::ExtractEquiJoinKeysRule),
        LeftJoinToInnerJoin(join_optimization::LeftJoinToInnerJoinRule),
        JoinConditionSimplify(join_optimization::JoinConditionSimplifyRule),
        JoinToExpand(join_optimization::JoinToExpandRule),
        JoinToAppendVertices(join_optimization::JoinToAppendVerticesRule),
        MergeConsecutiveExpand(join_optimization::MergeConsecutiveExpandRule),
        JoinElimination(join_optimization::JoinEliminationRule),
        IndexJoinSelection(join_optimization::IndexJoinSelectionRule),
        JoinReorder(join_optimization::JoinReorderRule),
    }
}

#[derive(Debug)]
pub struct RuleRegistry {
    rules: Vec<RewriteRule>,
}

impl RuleRegistry {
    pub fn new() -> Self {
        Self { rules: Vec::new() }
    }

    pub fn add(&mut self, rule: RewriteRule) {
        self.rules.push(rule);
    }

    pub fn iter(&self) -> impl Iterator<Item = &RewriteRule> {
        self.rules.iter()
    }

    pub fn len(&self) -> usize {
        self.rules.len()
    }

    pub fn is_empty(&self) -> bool {
        self.rules.is_empty()
    }

    pub fn clear(&mut self) {
        self.rules.clear();
    }

    pub fn into_vec(self) -> Vec<RewriteRule> {
        self.rules
    }
}

impl Default for RuleRegistry {
    fn default() -> Self {
        let mut registry = Self::new();
        registry.add(RewriteRule::EliminateFilter(
            elimination::EliminateFilterRule::new(),
        ));
        registry.add(RewriteRule::FoldConstants(
            constant_folding::FoldConstantsRule::new(),
        ));
        registry.add(RewriteRule::RemoveNoopProject(
            elimination::RemoveNoopProjectRule::new(),
        ));
        registry.add(RewriteRule::EliminateAppendVertices(
            elimination::EliminateAppendVerticesRule::new(),
        ));
        registry.add(RewriteRule::RemoveAppendVerticesBelowJoin(
            elimination::RemoveAppendVerticesBelowJoinRule::new(),
        ));
        registry.add(RewriteRule::EliminateRowCollect(
            elimination::EliminateRowCollectRule::new(),
        ));
        registry.add(RewriteRule::EliminateEmptySetOperation(
            elimination::EliminateEmptySetOperationRule::new(),
        ));
        registry.add(RewriteRule::DedupElimination(
            elimination::DedupEliminationRule::new(),
        ));
        registry.add(RewriteRule::EliminateSort(
            elimination::EliminateSortRule::new(),
        ));
        registry.add(RewriteRule::CombineFilter(merge::CombineFilterRule::new()));
        registry.add(RewriteRule::CollapseProject(
            merge::CollapseProjectRule::new(),
        ));
        registry.add(RewriteRule::CollapseConsecutiveProject(
            merge::CollapseConsecutiveProjectRule::new(),
        ));
        registry.add(RewriteRule::MergeGetVerticesAndProject(
            merge::MergeGetVerticesAndProjectRule::new(),
        ));
        registry.add(RewriteRule::MergeGetVerticesAndDedup(
            merge::MergeGetVerticesAndDedupRule::new(),
        ));
        registry.add(RewriteRule::MergeGetNbrsAndProject(
            merge::MergeGetNbrsAndProjectRule::new(),
        ));
        registry.add(RewriteRule::MergeGetNbrsAndDedup(
            merge::MergeGetNbrsAndDedupRule::new(),
        ));
        registry.add(RewriteRule::PushFilterDownTraverse(
            predicate_pushdown::PushFilterDownTraverseRule::new(),
        ));
        registry.add(RewriteRule::PushFilterDownExpandAll(
            predicate_pushdown::PushFilterDownExpandAllRule::new(),
        ));
        registry.add(RewriteRule::PushFilterDownNode(
            predicate_pushdown::PushFilterDownNodeRule::new(),
        ));
        registry.add(RewriteRule::PushEFilterDown(
            predicate_pushdown::PushEFilterDownRule::new(),
        ));
        registry.add(RewriteRule::PushVFilterDownScanVertices(
            predicate_pushdown::PushVFilterDownScanVerticesRule::new(),
        ));
        registry.add(RewriteRule::EliminateRedundantTagFilter(
            predicate_pushdown::EliminateRedundantTagFilterRule::new(),
        ));
        registry.add(RewriteRule::PushFilterDownScanVertices(
            predicate_pushdown::PushFilterDownScanVerticesRule::new(),
        ));
        registry.add(RewriteRule::PushFilterDownScanEdges(
            predicate_pushdown::PushFilterDownScanEdgesRule::new(),
        ));
        registry.add(RewriteRule::PushFilterDownInnerJoin(
            predicate_pushdown::PushFilterDownInnerJoinRule::new(),
        ));
        registry.add(RewriteRule::PushFilterDownLeftJoin(
            predicate_pushdown::PushFilterDownLeftJoinRule::new(),
        ));
        registry.add(RewriteRule::PushFilterDownCrossJoin(
            predicate_pushdown::PushFilterDownCrossJoinRule::new(),
        ));
        registry.add(RewriteRule::PushFilterDownGetNbrs(
            predicate_pushdown::PushFilterDownGetNbrsRule::new(),
        ));
        registry.add(RewriteRule::PushFilterDownAllPaths(
            predicate_pushdown::PushFilterDownAllPathsRule::new(),
        ));
        registry.add(RewriteRule::PushProjectDownScanVertices(
            projection_pushdown::PushProjectDownScanVerticesRule::new(),
        ));
        registry.add(RewriteRule::PushProjectDownScanEdges(
            projection_pushdown::PushProjectDownScanEdgesRule::new(),
        ));
        registry.add(RewriteRule::PushProjectDownGetVertices(
            projection_pushdown::PushProjectDownGetVerticesRule::new(),
        ));
        registry.add(RewriteRule::PushProjectDownGetNeighbors(
            projection_pushdown::PushProjectDownGetNeighborsRule::new(),
        ));
        registry.add(RewriteRule::PushProjectDownGetEdges(
            projection_pushdown::PushProjectDownGetEdgesRule::new(),
        ));
        registry.add(RewriteRule::PushProjectDownAppendVertices(
            projection_pushdown::PushProjectDownAppendVerticesRule::new(),
        ));
        registry.add(RewriteRule::EnrichScanSlotsWithFilterProps(
            slot_coverage::EnrichScanSlotsWithFilterPropsRule::new(),
        ));
        registry.add(RewriteRule::PushLimitDownGetVertices(
            limit_pushdown::PushLimitDownGetVerticesRule::new(),
        ));
        registry.add(RewriteRule::PushLimitDownGetEdges(
            limit_pushdown::PushLimitDownGetEdgesRule::new(),
        ));
        registry.add(RewriteRule::PushLimitDownScanVertices(
            limit_pushdown::PushLimitDownScanVerticesRule::new(),
        ));
        registry.add(RewriteRule::PushLimitDownScanEdges(
            limit_pushdown::PushLimitDownScanEdgesRule::new(),
        ));
        registry.add(RewriteRule::PushLimitDownIndexScan(
            limit_pushdown::PushLimitDownIndexScanRule::new(),
        ));
        registry.add(RewriteRule::PushTopNDownIndexScan(
            limit_pushdown::PushTopNDownIndexScanRule::new(),
        ));
        registry.add(RewriteRule::ConvertSortLimitToTopN(
            limit_pushdown::ConvertSortLimitToTopNRule::new(),
        ));
        registry.add(RewriteRule::PushFilterDownAggregate(
            aggregate::PushFilterDownAggregateRule::new(),
        ));
        registry.add(RewriteRule::ExpandPushdownAnnotate(
            expand_pushdown::ExpandPushdownAnnotateRule::new(),
        ));
        registry.add(RewriteRule::ScanIdentityAnnotate(
            scan_identity::ScanIdentityAnnotateRule::new(),
        ));
        registry.add(RewriteRule::ScanEdgeIdentityAnnotate(
            scan_edge_identity::ScanEdgeIdentityAnnotateRule::new(),
        ));
        registry.add(RewriteRule::UnnestSimplePatternApply(
            decorrelation::UnnestSimplePatternApplyRule::new(),
        ));
        registry.add(RewriteRule::PushProjectDownJoin(
            join_optimization::PushProjectDownJoinRule::new(),
        ));
        registry.add(RewriteRule::ExtractEquiJoinKeys(
            join_optimization::ExtractEquiJoinKeysRule::new(),
        ));
        registry.add(RewriteRule::LeftJoinToInnerJoin(
            join_optimization::LeftJoinToInnerJoinRule::new(),
        ));
        registry.add(RewriteRule::JoinConditionSimplify(
            join_optimization::JoinConditionSimplifyRule::new(),
        ));
        registry.add(RewriteRule::JoinToExpand(
            join_optimization::JoinToExpandRule::new(),
        ));
        registry.add(RewriteRule::JoinToAppendVertices(
            join_optimization::JoinToAppendVerticesRule::new(),
        ));
        registry.add(RewriteRule::MergeConsecutiveExpand(
            join_optimization::MergeConsecutiveExpandRule::new(),
        ));
        registry.add(RewriteRule::JoinElimination(
            join_optimization::JoinEliminationRule::new(),
        ));
        registry.add(RewriteRule::IndexJoinSelection(
            join_optimization::IndexJoinSelectionRule::new(),
        ));
        registry.add(RewriteRule::JoinReorder(
            join_optimization::JoinReorderRule::new(),
        ));
        registry
    }
}

impl RuleRegistry {
    /// Build a registry filtered by the `[optimizer.rules]` selection.
    ///
    /// `enabled` non-empty keeps only those rules (explicit selection wins
    /// over the default rule set); `disabled` then removes rules. Names match
    /// rule type names with or without the `Rule` suffix.
    pub fn filtered(disabled: &[String], enabled: &[String]) -> Self {
        fn normalise(name: &str) -> &str {
            name.strip_suffix("Rule").unwrap_or(name)
        }

        let base = Self::default();
        if disabled.is_empty() && enabled.is_empty() {
            return base;
        }
        let enabled_set: std::collections::HashSet<&str> =
            enabled.iter().map(|n| normalise(n)).collect();
        let disabled_set: std::collections::HashSet<&str> =
            disabled.iter().map(|n| normalise(n)).collect();
        let mut registry = Self::new();
        for rule in base.rules {
            let name = normalise(rule.name());
            if !enabled_set.is_empty() && !enabled_set.contains(name) {
                continue;
            }
            if disabled_set.contains(name) {
                continue;
            }
            registry.add(rule);
        }
        registry
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rule_registry_default() {
        let registry = RuleRegistry::default();
        assert_eq!(registry.len(), 58);
    }

    #[test]
    fn test_rule_names() {
        let registry = RuleRegistry::default();
        for rule in registry.iter() {
            let name = rule.name();
            assert!(!name.is_empty());
            assert!(name.ends_with("Rule"));
        }
    }

    #[test]
    fn test_filtered_disabled_removes_rule() {
        let registry = RuleRegistry::filtered(&["LeftJoinToInnerJoin".to_string()], &[]);
        assert_eq!(registry.len(), 57);
        assert!(registry
            .iter()
            .all(|rule| rule.name() != "LeftJoinToInnerJoinRule"));
    }

    #[test]
    fn test_filtered_enabled_keeps_only_selected() {
        let registry = RuleRegistry::filtered(
            &[],
            &[
                "EliminateFilterRule".to_string(),
                "CombineFilter".to_string(),
            ],
        );
        assert_eq!(registry.len(), 2);
        let names: Vec<&str> = registry.iter().map(|rule| rule.name()).collect();
        assert!(names.contains(&"EliminateFilterRule"));
        assert!(names.contains(&"CombineFilterRule"));
    }

    #[test]
    fn test_filtered_unknown_names_select_nothing() {
        let registry = RuleRegistry::filtered(&[], &["NoSuchRule".to_string()]);
        assert!(registry.is_empty());
    }
}