//! Hamelin-local optimizer rules.
//!
//! `extract_leaf_expressions` is a fork of the DataFusion 53 optimizer rules
//! of the same names, with one fix: refuse to push extraction projections
//! past `LogicalPlan::Unnest`. See the module docstring for the rationale.

mod extract_leaf_expressions;

pub use extract_leaf_expressions::{HamelinExtractLeafExpressions, HamelinPushDownLeafProjections};

use std::sync::Arc;

use datafusion::execution::session_state::SessionStateBuilder;
use datafusion_optimizer::{Optimizer, OptimizerRule};

use crate::json::JsonRewrite;

/// Build the Hamelin optimizer rule list: DataFusion's default rules with
/// the upstream `extract_leaf_expressions`/`push_down_leaf_projections`
/// pair swapped out for the Hamelin forks, plus `JsonRewrite` appended at
/// the end so it runs after `OptimizeProjections` has pruned dead
/// passthroughs.
pub fn hamelin_optimizer_rules() -> Vec<Arc<dyn OptimizerRule + Send + Sync>> {
    let default_rules = Optimizer::new().rules;
    let mut rules: Vec<Arc<dyn OptimizerRule + Send + Sync>> = default_rules
        .into_iter()
        .flat_map(|rule| {
            let name = rule.name();
            if name == "extract_leaf_expressions" {
                vec![Arc::new(HamelinExtractLeafExpressions::new())
                    as Arc<dyn OptimizerRule + Send + Sync>]
            } else if name == "push_down_leaf_projections" {
                vec![Arc::new(HamelinPushDownLeafProjections::new())
                    as Arc<dyn OptimizerRule + Send + Sync>]
            } else {
                vec![rule]
            }
        })
        .collect();
    rules.push(Arc::new(JsonRewrite));
    rules
}

/// Apply Hamelin's optimizer rule overrides to a `SessionStateBuilder`.
pub fn with_hamelin_optimizer_rules(builder: SessionStateBuilder) -> SessionStateBuilder {
    builder.with_optimizer_rules(hamelin_optimizer_rules())
}
