fluidattacks-blends-domain 0.2.0

use alloc::collections::{BTreeMap, BTreeSet};
use alloc::vec::Vec;

use crate::query::adj_ast;
use crate::syntax::{SyntaxGraph, SyntaxKind};
use crate::{CodeGraph, NodeId};

#[must_use]
pub const fn nodes_by_type(graph: &SyntaxGraph) -> &BTreeMap<SyntaxKind, Vec<NodeId>> {
    &graph.nodes_by_type
}

#[must_use]
pub fn get_nodes_by_path<G: CodeGraph>(graph: &G, n_id: NodeId, path: &[&str]) -> BTreeSet<NodeId> {
    let Some((first, rest)) = path.split_first() else {
        return BTreeSet::new();
    };
    let matches = adj_ast(graph, n_id, Some(1), &[*first]);
    if rest.is_empty() {
        matches.into_iter().collect()
    } else {
        matches
            .into_iter()
            .flat_map(|child| get_nodes_by_path(graph, child, rest))
            .collect()
    }
}

#[must_use]
pub fn get_node_by_path<G: CodeGraph>(graph: &G, n_id: NodeId, path: &[&str]) -> Option<NodeId> {
    get_nodes_by_path(graph, n_id, path).into_iter().next()
}

#[cfg(test)]
mod tests {
    use super::{get_node_by_path, get_nodes_by_path, nodes_by_type};
    use crate::ast::{AstGraph, AstNode};
    use crate::syntax::{SyntaxGraph, SyntaxKind, SyntaxNode};
    use crate::NodeId;
    use alloc::borrow::ToOwned;
    use alloc::collections::BTreeSet;
    use alloc::vec;

    fn sample() -> AstGraph {
        let mut ast = AstGraph::new();
        ast.add_node(
            NodeId(1),
            AstNode::new(1, 1, "class_declaration".to_owned()),
        );
        ast.add_node(NodeId(2), AstNode::new(1, 1, "base_list".to_owned()));
        ast.add_node(NodeId(3), AstNode::new(1, 1, "identifier".to_owned()));
        ast.add_node(NodeId(4), AstNode::new(1, 1, "modifier".to_owned()));
        ast.add_node(NodeId(5), AstNode::new(1, 2, "identifier".to_owned()));
        ast.add_edge(NodeId(1), NodeId(2), 0);
        ast.add_edge(NodeId(1), NodeId(4), 1);
        ast.add_edge(NodeId(2), NodeId(3), 0);
        ast.add_edge(NodeId(2), NodeId(5), 1);
        ast
    }

    #[test]
    fn nodes_by_type_exposes_the_index_maintained_by_add_node() {
        let mut graph = SyntaxGraph::new();
        graph.add_node(NodeId(9), SyntaxNode::Break);
        graph.add_node(NodeId(1), SyntaxNode::File);
        graph.add_node(NodeId(2), SyntaxNode::Break);

        assert_eq!(
            nodes_by_type(&graph).get(&SyntaxKind::Break),
            Some(&vec![NodeId(9), NodeId(2)])
        );
        assert_eq!(
            nodes_by_type(&graph).get(&SyntaxKind::File),
            Some(&vec![NodeId(1)])
        );
        assert_eq!(nodes_by_type(&graph).get(&SyntaxKind::If), None);
        assert!(nodes_by_type(&SyntaxGraph::new()).is_empty());
    }

    #[test]
    fn nodes_by_path_collect_every_match_of_the_last_step() {
        let ast = sample();
        assert_eq!(
            get_nodes_by_path(&ast, NodeId(1), &["base_list", "identifier"]),
            BTreeSet::from([NodeId(3), NodeId(5)])
        );
        assert_eq!(
            get_nodes_by_path(&ast, NodeId(1), &["base_list"]),
            BTreeSet::from([NodeId(2)])
        );
    }

    #[test]
    fn nodes_by_path_are_empty_when_the_path_breaks_or_is_empty() {
        let ast = sample();
        assert_eq!(
            get_nodes_by_path(&ast, NodeId(1), &["base_list", "block"]),
            BTreeSet::new()
        );
        assert_eq!(get_nodes_by_path(&ast, NodeId(1), &[]), BTreeSet::new());
    }

    #[test]
    fn follows_a_multi_step_label_path() {
        let ast = sample();
        assert_eq!(
            get_node_by_path(&ast, NodeId(1), &["base_list", "identifier"]),
            Some(NodeId(3))
        );
    }

    #[test]
    fn returns_the_direct_child_for_a_single_step_path() {
        let ast = sample();
        assert_eq!(
            get_node_by_path(&ast, NodeId(1), &["base_list"]),
            Some(NodeId(2))
        );
    }

    #[test]
    fn is_none_when_the_path_breaks_or_is_empty() {
        let ast = sample();
        assert_eq!(
            get_node_by_path(&ast, NodeId(1), &["base_list", "block"]),
            None
        );
        assert_eq!(get_node_by_path(&ast, NodeId(1), &["missing"]), None);
        assert_eq!(get_node_by_path(&ast, NodeId(1), &[]), None);
    }
}