streamd/src/query/find.rs

use crate::models::LocalizedShard;

/// Find all shards matching a predicate, recursively searching through children.
///
/// The search is depth-first, with the parent tested before its children.
pub fn find_shard<F>(shards: &[LocalizedShard], predicate: F) -> Vec<LocalizedShard>
where
    F: Fn(&LocalizedShard) -> bool + Copy,
{
    let mut found_shards = Vec::new();

    for shard in shards {
        if predicate(shard) {
            found_shards.push(shard.clone());
        }
        found_shards.extend(find_shard(&shard.children, predicate));
    }

    found_shards
}

/// Find all shards where a specific dimension has a specific value.
pub fn find_shard_by_position(
    shards: &[LocalizedShard],
    dimension: &str,
    value: &str,
) -> Vec<LocalizedShard> {
    find_shard(shards, |shard| {
        shard
            .location
            .get(dimension)
            .map(|v| v == value)
            .unwrap_or(false)
    })
}

/// Find all shards where a specific dimension is set (regardless of value).
pub fn find_shard_by_set_dimension(
    shards: &[LocalizedShard],
    dimension: &str,
) -> Vec<LocalizedShard> {
    find_shard(shards, |shard| shard.location.contains_key(dimension))
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::{TimeZone, Utc};
    use indexmap::IndexMap;

    fn generate_localized_shard(
        location: Option<IndexMap<String, String>>,
        children: Option<Vec<LocalizedShard>>,
    ) -> LocalizedShard {
        LocalizedShard {
            start_line: 1,
            end_line: 1,
            moment: Utc.with_ymd_and_hms(2020, 1, 1, 0, 0, 0).unwrap(),
            location: location.unwrap_or_default(),
            children: children.unwrap_or_default(),
            markers: vec![],
            tags: vec![],
        }
    }

    #[test]
    fn test_returns_empty_when_no_match() {
        let mut loc = IndexMap::new();
        loc.insert("file".to_string(), "a.md".to_string());
        let root = generate_localized_shard(Some(loc), None);
        let shards = vec![root];

        let result = find_shard(&shards, |s| s.location.contains_key("missing"));

        assert!(result.is_empty());
    }

    #[test]
    fn test_finds_matches_depth_first_and_preserves_order() {
        let mut loc1 = IndexMap::new();
        loc1.insert("k".to_string(), "match".to_string());
        let grandchild = generate_localized_shard(Some(loc1.clone()), None);

        let child1 = generate_localized_shard(Some(loc1), Some(vec![grandchild.clone()]));

        let mut loc2 = IndexMap::new();
        loc2.insert("k".to_string(), "nope".to_string());
        let child2 = generate_localized_shard(Some(loc2.clone()), None);

        let root = generate_localized_shard(Some(loc2), Some(vec![child1.clone(), child2]));

        let result = find_shard(&[root], |s| {
            s.location.get("k") == Some(&"match".to_string())
        });

        assert_eq!(result.len(), 2);
        assert_eq!(result[0], child1);
        assert_eq!(result[1], grandchild);
    }

    #[test]
    fn test_includes_root_if_it_matches() {
        let mut loc = IndexMap::new();
        loc.insert("k".to_string(), "match".to_string());

        let child = generate_localized_shard(Some(loc.clone()), None);
        let root = generate_localized_shard(Some(loc), Some(vec![child]));

        let result = find_shard(std::slice::from_ref(&root), |s| {
            s.location.get("k") == Some(&"match".to_string())
        });

        assert_eq!(result[0], root);
        assert_eq!(result.len(), 2);
    }

    #[test]
    fn test_multiple_roots_keeps_left_to_right_order() {
        let mut loc_match = IndexMap::new();
        loc_match.insert("k".to_string(), "match".to_string());

        let mut loc_nope = IndexMap::new();
        loc_nope.insert("k".to_string(), "nope".to_string());

        let a = generate_localized_shard(Some(loc_match.clone()), None);
        let b = generate_localized_shard(Some(loc_match), None);
        let c = generate_localized_shard(Some(loc_nope), None);

        let result = find_shard(&[a.clone(), b.clone(), c], |s| {
            s.location.get("k") == Some(&"match".to_string())
        });

        assert_eq!(result, vec![a, b]);
    }

    #[test]
    fn test_query_function_can_use_arbitrary_logic() {
        let mut loc1 = IndexMap::new();
        loc1.insert("x".to_string(), "1".to_string());

        let mut loc2 = IndexMap::new();
        loc2.insert("x".to_string(), "2".to_string());

        let mut loc3 = IndexMap::new();
        loc3.insert("x".to_string(), "3".to_string());

        let a = generate_localized_shard(Some(loc1), None);
        let b = generate_localized_shard(Some(loc2), None);
        let c = generate_localized_shard(Some(loc3), None);
        let root = generate_localized_shard(None, Some(vec![a, b.clone(), c]));

        let result = find_shard(&[root], |shard| {
            shard
                .location
                .get("x")
                .and_then(|x| x.parse::<i32>().ok())
                .map(|x| x % 2 == 0)
                .unwrap_or(false)
        });

        assert_eq!(result, vec![b]);
    }

    #[test]
    fn test_matches_only_when_dimension_present_and_equal() {
        let mut loc_match = IndexMap::new();
        loc_match.insert("file".to_string(), "a.md".to_string());
        loc_match.insert("line".to_string(), "10".to_string());

        let mut loc_wrong = IndexMap::new();
        loc_wrong.insert("file".to_string(), "a.md".to_string());
        loc_wrong.insert("line".to_string(), "11".to_string());

        let mut loc_missing = IndexMap::new();
        loc_missing.insert("file".to_string(), "a.md".to_string());

        let match_shard = generate_localized_shard(Some(loc_match), None);
        let wrong_value = generate_localized_shard(Some(loc_wrong), None);
        let missing_dim = generate_localized_shard(Some(loc_missing), None);

        let mut root_loc = IndexMap::new();
        root_loc.insert("root".to_string(), "x".to_string());
        let root = generate_localized_shard(
            Some(root_loc),
            Some(vec![match_shard.clone(), wrong_value, missing_dim]),
        );

        let result = find_shard_by_position(&[root], "line", "10");

        assert_eq!(result, vec![match_shard]);
    }

    #[test]
    fn test_recurses_through_children() {
        let mut loc_deep = IndexMap::new();
        loc_deep.insert("section".to_string(), "s1".to_string());
        let deep = generate_localized_shard(Some(loc_deep), None);

        let mut loc_mid = IndexMap::new();
        loc_mid.insert("section".to_string(), "s0".to_string());
        let mid = generate_localized_shard(Some(loc_mid), Some(vec![deep.clone()]));

        let root = generate_localized_shard(None, Some(vec![mid]));

        let result = find_shard_by_position(&[root], "section", "s1");

        assert_eq!(result, vec![deep]);
    }
}