smartdb/rust/crates/rustdb-query/src/projection.rs

use bson::{Bson, Document};

use crate::field_path::get_nested_value;

/// Apply a projection to a document.
/// Inclusion mode: only specified fields + _id.
/// Exclusion mode: all fields except specified ones.
/// _id can be explicitly excluded in either mode.
pub fn apply_projection(doc: &Document, projection: &Document) -> Document {
    if projection.is_empty() {
        return doc.clone();
    }

    // Determine mode: inclusion or exclusion
    let mut has_inclusion = false;
    let mut id_explicitly_set = false;

    for (key, value) in projection {
        if key == "_id" {
            id_explicitly_set = true;
            continue;
        }
        match value {
            Bson::Int32(0) | Bson::Int64(0) | Bson::Boolean(false) => {}
            _ => has_inclusion = true,
        }
    }

    if has_inclusion {
        apply_inclusion(doc, projection, id_explicitly_set)
    } else {
        apply_exclusion(doc, projection)
    }
}

fn apply_inclusion(doc: &Document, projection: &Document, id_explicitly_set: bool) -> Document {
    let mut result = Document::new();

    // Include _id by default unless explicitly excluded
    let include_id = if id_explicitly_set {
        is_truthy(projection.get("_id"))
    } else {
        true
    };

    if include_id {
        if let Some(id) = doc.get("_id") {
            result.insert("_id", id.clone());
        }
    }

    for (key, value) in projection {
        if key == "_id" {
            continue;
        }
        if !is_truthy(Some(value)) {
            continue;
        }

        if key.contains('.') {
            if let Some(val) = get_nested_value(doc, key) {
                // Rebuild nested structure
                set_nested_in_result(&mut result, key, val);
            }
        } else if let Some(val) = doc.get(key) {
            result.insert(key.clone(), val.clone());
        }
    }

    result
}

fn apply_exclusion(doc: &Document, projection: &Document) -> Document {
    let mut result = doc.clone();

    for (key, value) in projection {
        if !is_truthy(Some(value)) {
            if key.contains('.') {
                // Remove nested field
                remove_nested_from_result(&mut result, key);
            } else {
                result.remove(key);
            }
        }
    }

    result
}

fn is_truthy(value: Option<&Bson>) -> bool {
    match value {
        None => false,
        Some(Bson::Int32(0)) | Some(Bson::Int64(0)) | Some(Bson::Boolean(false)) => false,
        _ => true,
    }
}

fn set_nested_in_result(doc: &mut Document, path: &str, value: Bson) {
    let parts: Vec<&str> = path.split('.').collect();
    set_nested_recursive(doc, &parts, value);
}

fn set_nested_recursive(doc: &mut Document, parts: &[&str], value: Bson) {
    if parts.len() == 1 {
        doc.insert(parts[0].to_string(), value);
        return;
    }

    let key = parts[0];
    if !doc.contains_key(key) {
        doc.insert(key.to_string(), Bson::Document(Document::new()));
    }
    if let Some(Bson::Document(ref mut nested)) = doc.get_mut(key) {
        set_nested_recursive(nested, &parts[1..], value);
    }
}

fn remove_nested_from_result(doc: &mut Document, path: &str) {
    let parts: Vec<&str> = path.split('.').collect();
    remove_nested_recursive(doc, &parts);
}

fn remove_nested_recursive(doc: &mut Document, parts: &[&str]) {
    if parts.len() == 1 {
        doc.remove(parts[0]);
        return;
    }

    let key = parts[0];
    if let Some(Bson::Document(ref mut nested)) = doc.get_mut(key) {
        remove_nested_recursive(nested, &parts[1..]);
    }
}

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

    #[test]
    fn test_inclusion_projection() {
        let doc = bson::doc! { "_id": 1, "name": "Alice", "age": 30, "email": "a@b.c" };
        let proj = bson::doc! { "name": 1, "age": 1 };
        let result = apply_projection(&doc, &proj);
        assert!(result.contains_key("_id"));
        assert!(result.contains_key("name"));
        assert!(result.contains_key("age"));
        assert!(!result.contains_key("email"));
    }

    #[test]
    fn test_exclusion_projection() {
        let doc = bson::doc! { "_id": 1, "name": "Alice", "age": 30 };
        let proj = bson::doc! { "age": 0 };
        let result = apply_projection(&doc, &proj);
        assert!(result.contains_key("_id"));
        assert!(result.contains_key("name"));
        assert!(!result.contains_key("age"));
    }

    #[test]
    fn test_exclude_id() {
        let doc = bson::doc! { "_id": 1, "name": "Alice" };
        let proj = bson::doc! { "name": 1, "_id": 0 };
        let result = apply_projection(&doc, &proj);
        assert!(!result.contains_key("_id"));
        assert!(result.contains_key("name"));
    }
}