fluidattacks-blends-domain 0.3.0

//! Counterpart of `blends/query/traversal.py`.

use alloc::collections::BTreeSet;
use alloc::vec::Vec;

use crate::{CodeGraph, NodeId};

/// Which edges a traversal follows.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum EdgeKind {
    Any,
    Ast,
    Cfg,
}

fn children_for<G: CodeGraph>(graph: &G, n_id: NodeId, edges: EdgeKind) -> Vec<NodeId> {
    match edges {
        EdgeKind::Any => graph.children(n_id),
        EdgeKind::Ast => graph.ast_children(n_id),
        EdgeKind::Cfg => graph.cfg_children(n_id),
    }
}

#[derive(Default)]
struct AdjWalk {
    nodes: Vec<NodeId>,
    processed: BTreeSet<NodeId>,
}

fn collect_adj<G: CodeGraph>(
    graph: &G,
    n_id: NodeId,
    depth: i64,
    edges: EdgeKind,
    walk: &mut AdjWalk,
) {
    if depth == 0 || !walk.processed.insert(n_id) {
        return;
    }
    let childs = children_for(graph, n_id, edges);
    walk.nodes.extend(&childs);
    if depth == 1 {
        return;
    }
    for c_id in childs {
        collect_adj(graph, c_id, depth.saturating_sub(1), edges, walk);
    }
}

pub fn adj_lazy<G: CodeGraph>(
    graph: &G,
    n_id: NodeId,
    depth: Option<i64>,
    edges: EdgeKind,
) -> impl Iterator<Item = NodeId> {
    let mut walk = AdjWalk::default();
    collect_adj(graph, n_id, depth.unwrap_or(1), edges, &mut walk);
    walk.nodes.into_iter()
}

#[must_use]
pub fn adj<G: CodeGraph>(
    graph: &G,
    n_id: NodeId,
    depth: Option<i64>,
    edges: EdgeKind,
) -> Vec<NodeId> {
    adj_lazy(graph, n_id, depth, edges).collect()
}

#[must_use]
pub fn adj_ast<G: CodeGraph>(
    graph: &G,
    n_id: NodeId,
    depth: Option<i64>,
    label_types: &[&str],
) -> Vec<NodeId> {
    filter_by_label_type(graph, adj(graph, n_id, depth, EdgeKind::Ast), label_types)
}

#[must_use]
pub fn adj_cfg<G: CodeGraph>(
    graph: &G,
    n_id: NodeId,
    depth: Option<i64>,
    label_types: &[&str],
) -> Vec<NodeId> {
    filter_by_label_type(graph, adj(graph, n_id, depth, EdgeKind::Cfg), label_types)
}

fn parents_for<G: CodeGraph>(graph: &G, n_id: NodeId, edges: EdgeKind) -> Vec<NodeId> {
    match edges {
        EdgeKind::Any => graph.parents(n_id),
        EdgeKind::Ast => graph.ast_parents(n_id),
        EdgeKind::Cfg => graph.cfg_parents(n_id),
    }
}

fn collect_pred<G: CodeGraph>(
    graph: &G,
    n_id: NodeId,
    depth: i64,
    edges: EdgeKind,
    walk: &mut AdjWalk,
) {
    if depth == 0 || !walk.processed.insert(n_id) {
        return;
    }
    let parents = parents_for(graph, n_id, edges);
    walk.nodes.extend(&parents);
    if depth == 1 {
        return;
    }
    for p_id in parents {
        collect_pred(graph, p_id, depth.saturating_sub(1), edges, walk);
    }
}

pub fn pred_lazy<G: CodeGraph>(
    graph: &G,
    n_id: NodeId,
    depth: Option<i64>,
    edges: EdgeKind,
) -> impl Iterator<Item = NodeId> {
    let mut walk = AdjWalk::default();
    collect_pred(graph, n_id, depth.unwrap_or(1), edges, &mut walk);
    walk.nodes.into_iter()
}

#[must_use]
pub fn pred<G: CodeGraph>(
    graph: &G,
    n_id: NodeId,
    depth: Option<i64>,
    edges: EdgeKind,
) -> Vec<NodeId> {
    pred_lazy(graph, n_id, depth, edges).collect()
}

#[must_use]
pub fn pred_ast<G: CodeGraph>(graph: &G, n_id: NodeId, depth: Option<i64>) -> Vec<NodeId> {
    pred(graph, n_id, depth, EdgeKind::Ast)
}

#[must_use]
pub fn pred_cfg<G: CodeGraph>(graph: &G, n_id: NodeId, depth: Option<i64>) -> Vec<NodeId> {
    pred(graph, n_id, depth, EdgeKind::Cfg)
}

fn filter_by_label_type<G: CodeGraph>(
    graph: &G,
    childs: Vec<NodeId>,
    label_types: &[&str],
) -> Vec<NodeId> {
    if label_types.is_empty() {
        return childs;
    }
    childs
        .into_iter()
        .filter(|c_id| {
            graph
                .label_type(*c_id)
                .is_some_and(|kind| label_types.contains(&kind))
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::{adj, adj_ast, adj_cfg, adj_lazy, pred, pred_ast, pred_cfg, pred_lazy, EdgeKind};
    use crate::ast::{AstGraph, AstNode};
    use crate::syntax::{SyntaxGraph, SyntaxNode};
    use crate::NodeId;
    use alloc::borrow::ToOwned;
    use alloc::vec::Vec;

    fn sample_ast() -> AstGraph {
        let mut ast = AstGraph::new();
        ast.add_node(NodeId(1), AstNode::new(1, 1, "stream".to_owned()));
        ast.add_node(NodeId(2), AstNode::new(1, 1, "document".to_owned()));
        ast.add_node(NodeId(3), AstNode::new(2, 1, "document".to_owned()));
        ast.add_node(NodeId(4), AstNode::new(2, 1, "block_node".to_owned()));
        ast.add_edge(NodeId(1), NodeId(2), 0);
        ast.add_edge(NodeId(1), NodeId(3), 1);
        ast.add_edge(NodeId(3), NodeId(4), 0);
        ast
    }

    #[test]
    fn adj_defaults_to_direct_children_in_id_order() {
        let ast = sample_ast();
        assert_eq!(
            adj(&ast, NodeId(1), None, EdgeKind::Any),
            [NodeId(2), NodeId(3)]
        );
        assert_eq!(adj(&ast, NodeId(4), None, EdgeKind::Any), []);
        assert_eq!(adj(&ast, NodeId(1), Some(0), EdgeKind::Any), []);
    }

    #[test]
    fn adj_depth_covers_descendants_and_minus_one_is_infinite() {
        let ast = sample_ast();
        assert_eq!(
            adj(&ast, NodeId(1), Some(2), EdgeKind::Any),
            [NodeId(2), NodeId(3), NodeId(4)]
        );
        assert_eq!(
            adj(&ast, NodeId(1), Some(-1), EdgeKind::Any),
            [NodeId(2), NodeId(3), NodeId(4)]
        );
    }

    #[test]
    fn pred_depths_mirror_the_python_comprehensive_suite() {
        let mut ast = AstGraph::new();
        for (n_id, kind) in [(1, "a"), (2, "b"), (3, "c"), (4, "d"), (5, "e"), (6, "f")] {
            ast.add_node(NodeId(n_id), AstNode::new(1, 1, kind.to_owned()));
        }
        ast.add_edge(NodeId(1), NodeId(2), 0);
        ast.add_edge(NodeId(1), NodeId(3), 1);
        ast.add_edge(NodeId(2), NodeId(6), 0);
        ast.add_edge(NodeId(3), NodeId(4), 0);
        ast.add_edge(NodeId(4), NodeId(5), 0);

        assert_eq!(
            pred(&ast, NodeId(5), Some(-1), EdgeKind::Any),
            [NodeId(4), NodeId(3), NodeId(1)]
        );
        assert_eq!(pred(&ast, NodeId(5), Some(0), EdgeKind::Any), []);
        assert_eq!(pred(&ast, NodeId(5), None, EdgeKind::Any), [NodeId(4)]);
        assert_eq!(
            pred(&ast, NodeId(5), Some(2), EdgeKind::Any),
            [NodeId(4), NodeId(3)]
        );
        assert_eq!(pred(&ast, NodeId(1), Some(-1), EdgeKind::Any), []);
    }

    #[test]
    fn pred_lazy_yields_parents_in_order_expanding_each_node_once() {
        let mut ast = AstGraph::new();
        ast.add_node(NodeId(1), AstNode::new(1, 1, "a".to_owned()));
        ast.add_node(NodeId(2), AstNode::new(1, 1, "b".to_owned()));
        ast.add_node(NodeId(3), AstNode::new(1, 1, "c".to_owned()));
        ast.add_edge(NodeId(1), NodeId(2), 0);
        ast.add_edge(NodeId(2), NodeId(3), 0);
        ast.add_edge(NodeId(1), NodeId(3), 1);

        let direct: Vec<NodeId> = pred_lazy(&ast, NodeId(3), Some(1), EdgeKind::Any).collect();
        assert_eq!(direct, [NodeId(1), NodeId(2)]);

        let deep: Vec<NodeId> = pred_lazy(&ast, NodeId(3), Some(-1), EdgeKind::Any).collect();
        assert_eq!(deep, [NodeId(1), NodeId(2), NodeId(1)]);

        let zero: Vec<NodeId> = pred_lazy(&ast, NodeId(3), Some(0), EdgeKind::Any).collect();
        assert_eq!(zero, []);
    }

    #[test]
    fn pred_terminates_on_an_edge_cycle() {
        let mut cyclic = AstGraph::new();
        cyclic.add_node(NodeId(1), AstNode::new(1, 1, "a".to_owned()));
        cyclic.add_node(NodeId(2), AstNode::new(1, 1, "b".to_owned()));
        cyclic.add_edge(NodeId(1), NodeId(2), 0);
        cyclic.add_edge(NodeId(2), NodeId(1), 0);

        assert_eq!(
            pred_ast(&cyclic, NodeId(1), Some(-1)),
            [NodeId(2), NodeId(1)]
        );
        assert_eq!(pred_ast(&cyclic, NodeId(1), None), [NodeId(2)]);
    }

    #[test]
    fn pred_follows_only_the_requested_edge_kind() {
        let mut syntax = SyntaxGraph::new();
        syntax.add_node(NodeId(1), SyntaxNode::File);
        syntax.add_node(NodeId(2), SyntaxNode::ExecutionBlock);
        syntax.add_node(NodeId(3), SyntaxNode::ExecutionBlock);
        syntax.add_ast_edge(NodeId(1), NodeId(3));
        syntax.add_cfg_edge(NodeId(2), NodeId(3));

        assert_eq!(pred_ast(&syntax, NodeId(3), None), [NodeId(1)]);
        assert_eq!(pred_cfg(&syntax, NodeId(3), None), [NodeId(2)]);
        assert_eq!(
            pred(&syntax, NodeId(3), None, EdgeKind::Any),
            [NodeId(1), NodeId(2)]
        );
    }

    #[test]
    fn adj_lazy_repeats_a_shared_child_across_parents_expanding_it_once() {
        let mut ast = AstGraph::new();
        ast.add_node(NodeId(1), AstNode::new(1, 1, "a".to_owned()));
        ast.add_node(NodeId(2), AstNode::new(1, 1, "b".to_owned()));
        ast.add_node(NodeId(3), AstNode::new(1, 1, "c".to_owned()));
        ast.add_edge(NodeId(1), NodeId(2), 0);
        ast.add_edge(NodeId(1), NodeId(3), 1);
        ast.add_edge(NodeId(2), NodeId(3), 0);

        let yielded: Vec<NodeId> = adj_lazy(&ast, NodeId(1), Some(-1), EdgeKind::Any).collect();
        assert_eq!(yielded, [NodeId(2), NodeId(3), NodeId(3)]);
    }

    #[test]
    fn adj_terminates_on_an_edge_cycle() {
        let mut ast = AstGraph::new();
        ast.add_node(NodeId(1), AstNode::new(1, 1, "a".to_owned()));
        ast.add_node(NodeId(2), AstNode::new(1, 1, "b".to_owned()));
        ast.add_edge(NodeId(1), NodeId(2), 0);
        ast.add_edge(NodeId(2), NodeId(1), 0);

        assert_eq!(
            adj(&ast, NodeId(1), Some(-1), EdgeKind::Any),
            [NodeId(2), NodeId(1)]
        );
        assert_eq!(
            adj(&ast, NodeId(2), Some(-1), EdgeKind::Any),
            [NodeId(1), NodeId(2)]
        );
    }

    #[test]
    fn adj_ast_delegates_to_adj_and_filters_by_label_type() {
        let ast = sample_ast();
        assert_eq!(adj_ast(&ast, NodeId(1), None, &[]), [NodeId(2), NodeId(3)]);
        assert_eq!(
            adj_ast(&ast, NodeId(1), None, &["document"]),
            [NodeId(2), NodeId(3)]
        );
        assert_eq!(adj_ast(&ast, NodeId(1), None, &["block_node"]), []);
        assert_eq!(
            adj_ast(&ast, NodeId(1), Some(-1), &["block_node"]),
            [NodeId(4)]
        );
    }

    fn sample_syntax() -> SyntaxGraph {
        let mut graph = SyntaxGraph::new();
        graph.add_node(NodeId(1), SyntaxNode::File);
        graph.add_node(NodeId(2), SyntaxNode::ArrayInitializer);
        graph.add_node(
            NodeId(3),
            SyntaxNode::Object {
                name: None,
                tf_reference: None,
            },
        );
        graph.add_node(
            NodeId(4),
            SyntaxNode::Literal {
                value: "doe".to_owned(),
                value_type: "string".to_owned(),
            },
        );
        graph.add_ast_edge(NodeId(1), NodeId(2));
        graph.add_ast_edge(NodeId(1), NodeId(3));
        graph.add_cfg_edge(NodeId(1), NodeId(3));
        graph.add_cfg_edge(NodeId(3), NodeId(4));
        graph
    }

    #[test]
    fn adj_cfg_follows_only_edges_with_the_cfg_mark() {
        let syntax = sample_syntax();
        assert_eq!(adj_cfg(&syntax, NodeId(1), None, &[]), [NodeId(3)]);
        assert_eq!(
            adj_cfg(&syntax, NodeId(1), Some(-1), &[]),
            [NodeId(3), NodeId(4)]
        );
        assert_eq!(adj_cfg(&syntax, NodeId(2), None, &[]), []);
    }

    #[test]
    fn adj_cfg_filters_by_label_type() {
        let syntax = sample_syntax();
        assert_eq!(
            adj_cfg(&syntax, NodeId(1), Some(-1), &["Literal"]),
            [NodeId(4)]
        );
        assert_eq!(adj_cfg(&syntax, NodeId(1), None, &["Literal"]), []);
    }
}