linkrs-query 0.1.0

//! Rewrite the macro definitions for the rules
//!
//! Provide declarative macros to simplify the definition of rewriting rules and reduce the amount of boilerplate code.

// ==================== Basic Rules Macros ====================

/// Define the basic rules for rule rewriting.
///
/// Automatically generates the rule struct, Default implementation, new() method and RewriteRule trait implementation
///
/// # Example
/// ```rust
/// use linkrs_query::define_rewrite_rule;
/// use linkrs_query::optimizer::heuristic::pattern::Pattern;
///
/// define_rewrite_rule! {
///     name: MyCustomRule,
///     pattern: Pattern::new_with_name("Filter"),
///     apply: |_ctx, _node| {
///         Ok(None)
///     }
/// }
/// ```
#[macro_export]
macro_rules! define_rewrite_rule {
    (
        $(#[$meta:meta])*
        name: $name:ident,
        pattern: $pattern:expr,
        apply: $apply_closure:expr
    ) => {
        $(#[$meta])*
        #[derive(Debug)]
        pub struct $name;

        impl $name {
            /// Create a rule instance
            pub fn new() -> Self {
                Self
            }
        }

        impl Default for $name {
            fn default() -> Self {
                Self::new()
            }
        }

        impl $crate::optimizer::heuristic::rule::RewriteRule for $name {
            fn name(&self) -> &'static str {
                stringify!($name)
            }

            fn pattern(&self) -> $crate::optimizer::heuristic::pattern::Pattern {
                $pattern
            }

            fn apply(
                &self,
                ctx: &mut $crate::optimizer::heuristic::context::RewriteContext,
                node: &$crate::planning::plan::PlanNodeEnum,
            ) -> $crate::optimizer::heuristic::result::RewriteResult<Option<$crate::optimizer::heuristic::result::TransformResult>> {
                let apply_fn: fn(&mut _, &_) -> _ = $apply_closure;
                apply_fn(ctx, node)
            }
        }
    };
}

/// Define rules that match node types.
///
/// Automatic processing of node type matching and unpacking
///
/// # Examples
/// ```rust
/// use linkrs_query::define_typed_rewrite_rule;
/// use linkrs_query::optimizer::heuristic::pattern::Pattern;
///
/// define_typed_rewrite_rule! {
///     name: EliminateFilterRule,
///     pattern: Pattern::new_with_name("Filter"),
///     node_type: Filter,
///     apply: |_ctx, _filter_node| {
///         // `_filter_node` is already unpacked to `&FilterNode` by the macro.
///         Ok(None)
///     }
/// }
/// ```
#[macro_export]
macro_rules! define_typed_rewrite_rule {
    (
        $(#[$meta:meta])*
        name: $name:ident,
        pattern: $pattern:expr,
        node_type: $node_type:ident,
        apply: $apply_closure:expr
    ) => {
        $(#[$meta])*
        #[derive(Debug)]
        pub struct $name;

        impl $name {
            /// Creating rule instances
            pub fn new() -> Self {
                Self
            }
        }

        impl Default for $name {
            fn default() -> Self {
                Self::new()
            }
        }

        impl $crate::optimizer::heuristic::rule::RewriteRule for $name {
            fn name(&self) -> &'static str {
                stringify!($name)
            }

            fn pattern(&self) -> $crate::optimizer::heuristic::pattern::Pattern {
                $pattern
            }

            fn apply(
                &self,
                ctx: &mut $crate::optimizer::heuristic::context::RewriteContext,
                node: &$crate::planning::plan::PlanNodeEnum,
            ) -> $crate::optimizer::heuristic::result::RewriteResult<Option<$crate::optimizer::heuristic::result::TransformResult>> {
                use $crate::planning::plan::PlanNodeEnum;
                use $crate::planning::plan::core::nodes::base::plan_node_traits::SingleInputNode;

                let typed_node = match node {
                    PlanNodeEnum::$node_type(n) => n,
                    _ => return Ok(None),
                };

                let apply_fn: fn(&mut _, &_) -> _ = $apply_closure;
                apply_fn(ctx, typed_node)
            }
        }
    };
}

// ==================== Pushdown Rule Macros ====================

/// Define the push-down rule.
///
/// Automatically generate implementations for the RewriteRule and PushDownRule traits.
///
/// # Examples
/// ```rust
/// use linkrs_query::define_rewrite_pushdown_rule;
/// use linkrs_query::optimizer::heuristic::pattern::Pattern;
///
/// define_rewrite_pushdown_rule! {
///     name: PushFilterDownGetNbrsRule,
///     parent_node: Filter,
///     child_node: GetNeighbors,
///     apply: |_ctx, _filter_node, _get_neighbors_node| {
///         Ok(None)
///     }
/// }
/// ```
#[macro_export]
macro_rules! define_rewrite_pushdown_rule {
    (
        $(#[$meta:meta])*
        name: $name:ident,
        parent_node: $parent_type:ident,
        child_node: $child_type:ident,
        apply: $apply_closure:expr
    ) => {
        $(#[$meta])*
        #[derive(Debug)]
        pub struct $name;

        impl $name {
            /// Creating rule instances
            pub fn new() -> Self {
                Self
            }
        }

        impl Default for $name {
            fn default() -> Self {
                Self::new()
            }
        }

        impl $crate::optimizer::heuristic::rule::RewriteRule for $name {
            fn name(&self) -> &'static str {
                stringify!($name)
            }

            fn pattern(&self) -> $crate::optimizer::heuristic::pattern::Pattern {
                $crate::optimizer::heuristic::pattern::Pattern::new_with_name(stringify!($parent_type))
                    .with_dependency_name(stringify!($child_type))
            }

            fn apply(
                &self,
                ctx: &mut $crate::optimizer::heuristic::context::RewriteContext,
                node: &$crate::planning::plan::PlanNodeEnum,
            ) -> $crate::optimizer::heuristic::result::RewriteResult<Option<$crate::optimizer::heuristic::result::TransformResult>> {
                use $crate::planning::plan::PlanNodeEnum;
                use $crate::planning::plan::core::nodes::base::plan_node_traits::SingleInputNode;

                let parent_node = match node {
                    PlanNodeEnum::$parent_type(n) => n,
                    _ => return Ok(None),
                };

                let input = parent_node.input();
                let child_node = match input {
                    PlanNodeEnum::$child_type(n) => n,
                    _ => return Ok(None),
                };

                let apply_fn: fn(&mut _, &_, &_) -> _ = $apply_closure;
                apply_fn(ctx, parent_node, child_node)
            }
        }

        impl $crate::optimizer::heuristic::rule::PushDownRule for $name {
            fn can_push_down(
                &self,
                node: &$crate::planning::plan::PlanNodeEnum,
                target: &$crate::planning::plan::PlanNodeEnum,
            ) -> bool {
                use $crate::planning::plan::PlanNodeEnum;
                matches!(
                    (node, target),
                    (PlanNodeEnum::$parent_type(_), PlanNodeEnum::$child_type(_))
                )
            }

            fn push_down(
                &self,
                ctx: &mut $crate::optimizer::heuristic::context::RewriteContext,
                node: &$crate::planning::plan::PlanNodeEnum,
                _target: &$crate::planning::plan::PlanNodeEnum,
            ) -> $crate::optimizer::heuristic::result::RewriteResult<Option<$crate::optimizer::heuristic::result::TransformResult>> {
                use $crate::optimizer::heuristic::rule::RewriteRule;
                self.apply(ctx, node)
            }
        }
    };
}

// ==================== Remove Rule Macros ====================

/// Define the elimination rules.
///
/// Automatically generate implementations of the RewriteRule and EliminationRule traits.
///
/// # Examples
/// ```rust
/// use linkrs_query::define_rewrite_elimination_rule;
/// use linkrs_query::optimizer::heuristic::result::TransformResult;
/// use linkrs_query::optimizer::heuristic::rule::RewriteRule;
/// use linkrs_query::planning::plan::core::nodes::base::plan_node_traits::SingleInputNode;
/// use linkrs_query::planning::plan::core::nodes::operation::filter_node::FilterNode;
///
/// define_rewrite_elimination_rule! {
///     name: EliminateFilterRule,
///     node_type: Filter,
///     can_eliminate: |filter_node: &FilterNode| filter_node.condition().is_constant(),
///     eliminate: |_ctx, filter_node: &FilterNode| {
///         let mut result = TransformResult::new();
///         result.erase_curr = true;
///         result.add_new_node(filter_node.input().clone());
///         Ok(Some(result))
///     }
/// }
/// ```
#[macro_export]
macro_rules! define_rewrite_elimination_rule {
    (
        $(#[$meta:meta])*
        name: $name:ident,
        node_type: $node_type:ident,
        can_eliminate: $can_eliminate_closure:expr,
        eliminate: $eliminate_closure:expr
    ) => {
        $(#[$meta])*
        #[derive(Debug)]
        pub struct $name;

        impl $name {
            /// Creating rule instances
            pub fn new() -> Self {
                Self
            }
        }

        impl Default for $name {
            fn default() -> Self {
                Self::new()
            }
        }

        impl $crate::optimizer::heuristic::rule::RewriteRule for $name {
            fn name(&self) -> &'static str {
                stringify!($name)
            }

            fn pattern(&self) -> $crate::optimizer::heuristic::pattern::Pattern {
                $crate::optimizer::heuristic::pattern::Pattern::new_with_name(stringify!($node_type))
            }

            fn apply(
                &self,
                ctx: &mut $crate::optimizer::heuristic::context::RewriteContext,
                node: &$crate::planning::plan::PlanNodeEnum,
            ) -> $crate::optimizer::heuristic::result::RewriteResult<Option<$crate::optimizer::heuristic::result::TransformResult>> {
                use $crate::planning::plan::PlanNodeEnum;

                let typed_node = match node {
                    PlanNodeEnum::$node_type(n) => n,
                    _ => return Ok(None),
                };

                let can_eliminate_fn: fn(&_) -> _ = $can_eliminate_closure;
                if !can_eliminate_fn(typed_node) {
                    return Ok(None);
                }

                let eliminate_fn: fn(&mut _, &_) -> _ = $eliminate_closure;
                eliminate_fn(ctx, typed_node)
            }
        }

        impl $crate::optimizer::heuristic::rule::EliminationRule for $name {
            fn can_eliminate(&self, node: &$crate::planning::plan::PlanNodeEnum) -> bool {
                use $crate::planning::plan::PlanNodeEnum;

                let typed_node = match node {
                    PlanNodeEnum::$node_type(n) => n,
                    _ => return false,
                };

                let can_eliminate_fn: fn(&_) -> _ = $can_eliminate_closure;
                can_eliminate_fn(typed_node)
            }

            fn eliminate(
                &self,
                ctx: &mut $crate::optimizer::heuristic::context::RewriteContext,
                node: &$crate::planning::plan::PlanNodeEnum,
            ) -> $crate::optimizer::heuristic::result::RewriteResult<Option<$crate::optimizer::heuristic::result::TransformResult>> {
                self.apply(ctx, node)
            }
        }
    };
}

/// Define a simple elimination rule (which only deletes the current node).
///
/// # Examples
/// ```rust
/// use linkrs_query::define_simple_rewrite_elimination_rule;
/// use linkrs_query::planning::plan::core::nodes::operation::filter_node::FilterNode;
///
/// define_simple_rewrite_elimination_rule! {
///     name: EliminateTrueFilterRule,
///     node_type: Filter,
///     condition: |filter_node: &FilterNode| filter_node.condition().is_constant()
/// }
/// ```
#[macro_export]
macro_rules! define_simple_rewrite_elimination_rule {
    (
        $(#[$meta:meta])*
        name: $name:ident,
        node_type: $node_type:ident,
        condition: $condition_closure:expr
    ) => {
        $(#[$meta])*
        #[derive(Debug)]
        pub struct $name;

        impl $name {
            /// Creating rule instances
            pub fn new() -> Self {
                Self
            }
        }

        impl Default for $name {
            fn default() -> Self {
                Self::new()
            }
        }

        impl $crate::optimizer::heuristic::rule::RewriteRule for $name {
            fn name(&self) -> &'static str {
                stringify!($name)
            }

            fn pattern(&self) -> $crate::optimizer::heuristic::pattern::Pattern {
                $crate::optimizer::heuristic::pattern::Pattern::new_with_name(stringify!($node_type))
            }

            fn apply(
                &self,
                _ctx: &mut $crate::optimizer::heuristic::context::RewriteContext,
                node: &$crate::planning::plan::PlanNodeEnum,
            ) -> $crate::optimizer::heuristic::result::RewriteResult<Option<$crate::optimizer::heuristic::result::TransformResult>> {
                use $crate::planning::plan::PlanNodeEnum;
                use $crate::planning::plan::core::nodes::base::plan_node_traits::SingleInputNode;

                let typed_node = match node {
                    PlanNodeEnum::$node_type(n) => n,
                    _ => return Ok(None),
                };

                let condition_fn: fn(&_) -> _ = $condition_closure;
                if !condition_fn(typed_node) {
                    return Ok(None);
                }

                let mut result = $crate::optimizer::heuristic::result::TransformResult::new();
                result.erase_curr = true;
                result.add_new_node(typed_node.input().clone());
                Ok(Some(result))
            }
        }

        impl $crate::optimizer::heuristic::rule::EliminationRule for $name {
            fn can_eliminate(&self, node: &$crate::planning::plan::PlanNodeEnum) -> bool {
                use $crate::planning::plan::PlanNodeEnum;

                let typed_node = match node {
                    PlanNodeEnum::$node_type(n) => n,
                    _ => return false,
                };

                let condition_fn: fn(&_) -> _ = $condition_closure;
                condition_fn(typed_node)
            }

            fn eliminate(
                &self,
                ctx: &mut $crate::optimizer::heuristic::context::RewriteContext,
                node: &$crate::planning::plan::PlanNodeEnum,
            ) -> $crate::optimizer::heuristic::result::RewriteResult<Option<$crate::optimizer::heuristic::result::TransformResult>> {
                use $crate::optimizer::heuristic::rule::RewriteRule;
                self.apply(ctx, node)
            }
        }
    };
}

// ==================== Merging Rules Macro ====================

/// Define merge rules
///
/// Automatically generate implementations for the RewriteRule and MergeRule traits.
///
/// # Examples
/// ```rust
/// use linkrs_query::define_rewrite_merge_rule;
/// use linkrs_query::optimizer::heuristic::pattern::Pattern;
///
/// define_rewrite_merge_rule! {
///     name: CombineFilterRule,
///     parent_node: Filter,
///     child_node: Filter,
///     can_merge: |_parent, _child| true,
///     merge: |_ctx, _parent, _child| {
///         Ok(None)
///     }
/// }
/// ```
#[macro_export]
macro_rules! define_rewrite_merge_rule {
    (
        $(#[$meta:meta])*
        name: $name:ident,
        parent_node: $parent_type:ident,
        child_node: $child_type:ident,
        can_merge: $can_merge_closure:expr,
        merge: $merge_closure:expr
    ) => {
        $(#[$meta])*
        #[derive(Debug)]
        pub struct $name;

        impl $name {
            /// Creating rule instances
            pub fn new() -> Self {
                Self
            }
        }

        impl Default for $name {
            fn default() -> Self {
                Self::new()
            }
        }

        impl $crate::optimizer::heuristic::rule::RewriteRule for $name {
            fn name(&self) -> &'static str {
                stringify!($name)
            }

            fn pattern(&self) -> $crate::optimizer::heuristic::pattern::Pattern {
                $crate::optimizer::heuristic::pattern::Pattern::new_with_name(stringify!($parent_type))
                    .with_dependency_name(stringify!($child_type))
            }

            fn apply(
                &self,
                ctx: &mut $crate::optimizer::heuristic::context::RewriteContext,
                node: &$crate::planning::plan::PlanNodeEnum,
            ) -> $crate::optimizer::heuristic::result::RewriteResult<Option<$crate::optimizer::heuristic::result::TransformResult>> {
                use $crate::planning::plan::PlanNodeEnum;
                use $crate::planning::plan::core::nodes::base::plan_node_traits::SingleInputNode;

                let parent_node = match node {
                    PlanNodeEnum::$parent_type(n) => n,
                    _ => return Ok(None),
                };

                let input = parent_node.input();
                let child_node = match input {
                    PlanNodeEnum::$child_type(n) => n,
                    _ => return Ok(None),
                };

                let can_merge_fn: fn(&_, &_) -> _ = $can_merge_closure;
                if !can_merge_fn(parent_node, child_node) {
                    return Ok(None);
                }

                let merge_fn: fn(&mut _, &_, &_) -> _ = $merge_closure;
                merge_fn(ctx, parent_node, child_node)
            }
        }

        impl $crate::optimizer::heuristic::rule::MergeRule for $name {
            fn can_merge(
                &self,
                parent: &$crate::planning::plan::PlanNodeEnum,
                child: &$crate::planning::plan::PlanNodeEnum,
            ) -> bool {
                use $crate::planning::plan::PlanNodeEnum;

                let parent_node = match parent {
                    PlanNodeEnum::$parent_type(n) => n,
                    _ => return false,
                };

                let child_node = match child {
                    PlanNodeEnum::$child_type(n) => n,
                    _ => return false,
                };

                let can_merge_fn: fn(&_, &_) -> _ = $can_merge_closure;
                can_merge_fn(parent_node, child_node)
            }

            fn create_merged_node(
                &self,
                ctx: &mut $crate::optimizer::heuristic::context::RewriteContext,
                parent: &$crate::planning::plan::PlanNodeEnum,
                child: &$crate::planning::plan::PlanNodeEnum,
            ) -> $crate::optimizer::heuristic::result::RewriteResult<Option<$crate::optimizer::heuristic::result::TransformResult>> {
                use $crate::planning::plan::PlanNodeEnum;

                let parent_node = match parent {
                    PlanNodeEnum::$parent_type(n) => n,
                    _ => return Ok(None),
                };

                let child_node = match child {
                    PlanNodeEnum::$child_type(n) => n,
                    _ => return Ok(None),
                };

                let merge_fn: fn(&mut _, &_, &_) -> _ = $merge_closure;
                merge_fn(ctx, parent_node, child_node)
            }
        }
    };
}

// ==================== Rule Registration Macro ====================

/// Definition of rules: Registry
///
/// Automatically generate the default implementation of RuleRegistry, which includes the registration of all rules.
///
/// # Examples
/// ```ignore
/// // The generated code references the crate-internal `RewriteRule` enum,
/// // the `paste` macro and the rule structs, so it can only be expanded
/// // inside the crate itself.
/// define_rewrite_rule_registry! {
///     elimination: [
///         EliminateFilter,
///         RemoveNoopProject,
///     ],
///     merge: [
///         CombineFilter,
///         CollapseProject,
///     ],
/// }
/// ```
#[macro_export]
macro_rules! define_rewrite_rule_registry {
    (
        $(
            $category:ident: [
                $($rule_name:ident),* $(,)?
            ]
        ),* $(,)?
    ) => {
        impl Default for RuleRegistry {
            fn default() -> Self {
                let mut registry = Self::new();
                $(
                    $(
                        registry.add(RewriteRule::$rule_name(
                            $crate::optimizer::heuristic::$category::paste! {[<$rule_name Rule>]}::new()
                        ));
                    )*
                )*
                registry
            }
        }
    };
}

// Re-export all macros
pub use crate::define_rewrite_elimination_rule;
pub use crate::define_rewrite_merge_rule;
pub use crate::define_rewrite_pushdown_rule;
pub use crate::define_rewrite_rule;
pub use crate::define_rewrite_rule_registry;
pub use crate::define_simple_rewrite_elimination_rule;
pub use crate::define_typed_rewrite_rule;