smartdb/rust/crates/rustdb-index/src/planner.rs

use std::collections::HashSet;

use bson::{Bson, Document};
use tracing::debug;

use crate::engine::IndexEngine;

/// The execution plan for a query.
#[derive(Debug, Clone)]
pub enum QueryPlan {
    /// Full collection scan - no suitable index found.
    CollScan,
    /// Index scan with exact/equality matches.
    IxScan {
        /// Name of the index used.
        index_name: String,
        /// Candidate document IDs from the index.
        candidate_ids: HashSet<String>,
    },
    /// Index scan with range-based matches.
    IxScanRange {
        /// Name of the index used.
        index_name: String,
        /// Candidate document IDs from the range scan.
        candidate_ids: HashSet<String>,
    },
}

/// Plans query execution by selecting the best available index.
pub struct QueryPlanner;

impl QueryPlanner {
    /// Analyze a filter and the available indexes to produce a query plan.
    pub fn plan(filter: &Document, engine: &IndexEngine) -> QueryPlan {
        if filter.is_empty() {
            debug!("Empty filter -> CollScan");
            return QueryPlan::CollScan;
        }

        let indexes = engine.list_indexes();
        let mut best_plan: Option<QueryPlan> = None;
        let mut best_score: f64 = 0.0;

        for idx_info in &indexes {
            let index_fields: Vec<String> = idx_info.key.keys().map(|k| k.to_string()).collect();

            let mut matched = false;
            let mut score: f64 = 0.0;
            let mut is_range = false;

            for field in &index_fields {
                if let Some(condition) = filter.get(field) {
                    matched = true;
                    let field_score = Self::score_condition(condition);
                    score += field_score;

                    if Self::is_range_condition(condition) {
                        is_range = true;
                    }
                }
            }

            if !matched {
                continue;
            }

            // Unique index bonus
            if idx_info.unique {
                score += 0.5;
            }

            if score > best_score {
                best_score = score;

                // Try to get candidates from the engine
                // We build a sub-filter with only the fields this index covers
                let mut sub_filter = Document::new();
                for field in &index_fields {
                    if let Some(val) = filter.get(field) {
                        sub_filter.insert(field.clone(), val.clone());
                    }
                }

                if let Some(candidates) = engine.find_candidate_ids(&sub_filter) {
                    if is_range {
                        best_plan = Some(QueryPlan::IxScanRange {
                            index_name: idx_info.name.clone(),
                            candidate_ids: candidates,
                        });
                    } else {
                        best_plan = Some(QueryPlan::IxScan {
                            index_name: idx_info.name.clone(),
                            candidate_ids: candidates,
                        });
                    }
                }
            }
        }

        match best_plan {
            Some(plan) => {
                debug!(score = best_score, "Selected index plan");
                plan
            }
            None => {
                debug!("No suitable index found -> CollScan");
                QueryPlan::CollScan
            }
        }
    }

    /// Score a filter condition for index selectivity.
    /// Higher scores indicate more selective (better) index usage.
    fn score_condition(condition: &Bson) -> f64 {
        match condition {
            Bson::Document(doc) if Self::has_operators(doc) => {
                let mut score: f64 = 0.0;
                for (op, _) in doc {
                    score += match op.as_str() {
                        "$eq" => 2.0,
                        "$in" => 1.5,
                        "$gt" | "$gte" | "$lt" | "$lte" => 1.0,
                        _ => 0.0,
                    };
                }
                score
            }
            // Direct equality
            _ => 2.0,
        }
    }

    /// Check if a condition involves range operators.
    fn is_range_condition(condition: &Bson) -> bool {
        match condition {
            Bson::Document(doc) => {
                doc.keys().any(|k| matches!(k.as_str(), "$gt" | "$gte" | "$lt" | "$lte"))
            }
            _ => false,
        }
    }

    fn has_operators(doc: &Document) -> bool {
        doc.keys().any(|k| k.starts_with('$'))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::engine::IndexOptions;
    use bson::oid::ObjectId;

    #[test]
    fn test_empty_filter_collscan() {
        let engine = IndexEngine::new();
        let plan = QueryPlanner::plan(&bson::doc! {}, &engine);
        assert!(matches!(plan, QueryPlan::CollScan));
    }

    #[test]
    fn test_id_equality_ixscan() {
        let mut engine = IndexEngine::new();
        let oid = ObjectId::new();
        let doc = bson::doc! { "_id": oid.clone(), "name": "Alice" };
        engine.on_insert(&doc).unwrap();

        let filter = bson::doc! { "_id": oid };
        let plan = QueryPlanner::plan(&filter, &engine);
        assert!(matches!(plan, QueryPlan::IxScan { .. }));
    }

    #[test]
    fn test_indexed_field_ixscan() {
        let mut engine = IndexEngine::new();
        engine.create_index(
            bson::doc! { "status": 1 },
            IndexOptions::default(),
        ).unwrap();

        let doc = bson::doc! { "_id": ObjectId::new(), "status": "active" };
        engine.on_insert(&doc).unwrap();

        let filter = bson::doc! { "status": "active" };
        let plan = QueryPlanner::plan(&filter, &engine);
        assert!(matches!(plan, QueryPlan::IxScan { .. }));
    }

    #[test]
    fn test_unindexed_field_collscan() {
        let engine = IndexEngine::new();
        let filter = bson::doc! { "unindexed_field": "value" };
        let plan = QueryPlanner::plan(&filter, &engine);
        assert!(matches!(plan, QueryPlan::CollScan));
    }

    #[test]
    fn test_range_query_ixscan_range() {
        let mut engine = IndexEngine::new();
        engine.create_index(
            bson::doc! { "age": 1 },
            IndexOptions::default(),
        ).unwrap();

        let doc = bson::doc! { "_id": ObjectId::new(), "age": 30 };
        engine.on_insert(&doc).unwrap();

        let filter = bson::doc! { "age": { "$gte": 25, "$lt": 35 } };
        let plan = QueryPlanner::plan(&filter, &engine);
        assert!(matches!(plan, QueryPlan::IxScanRange { .. }));
    }

    #[test]
    fn test_unique_index_preferred() {
        let mut engine = IndexEngine::new();
        engine.create_index(
            bson::doc! { "email": 1 },
            IndexOptions { unique: true, ..Default::default() },
        ).unwrap();
        engine.create_index(
            bson::doc! { "email": 1, "name": 1 },
            IndexOptions { name: Some("email_name".to_string()), ..Default::default() },
        ).unwrap();

        let doc = bson::doc! { "_id": ObjectId::new(), "email": "a@b.com", "name": "Alice" };
        engine.on_insert(&doc).unwrap();

        let filter = bson::doc! { "email": "a@b.com" };
        let plan = QueryPlanner::plan(&filter, &engine);

        // The unique index on email should be preferred (higher score)
        match plan {
            QueryPlan::IxScan { index_name, .. } => {
                assert_eq!(index_name, "email_1");
            }
            _ => panic!("Expected IxScan"),
        }
    }
}