fluidattacks-blends-domain 0.6.0

//! AST graph: node table plus parent → child adjacency.

use alloc::collections::BTreeMap;
use alloc::string::String;
use alloc::vec::Vec;

use crate::ast::AstEdge;
use crate::ast::AstNode;
use crate::{Ast, CodeGraph, NodeId};

#[derive(Clone, PartialEq, Eq, Debug, Default)]
pub struct AstGraph {
    pub nodes: BTreeMap<NodeId, AstNode>,
    pub edges: BTreeMap<NodeId, BTreeMap<NodeId, AstEdge>>,
}

impl AstGraph {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    pub fn add_node(&mut self, id: NodeId, node: AstNode) {
        self.nodes.insert(id, node);
    }

    pub fn add_edge(&mut self, from: NodeId, to: NodeId, index: u32) {
        self.edges
            .entry(from)
            .or_default()
            .insert(to, AstEdge { kind: Ast, index });
    }

    pub fn set_field(&mut self, parent: NodeId, name: String, child: NodeId) {
        if let Some(node) = self.nodes.get_mut(&parent) {
            node.fields.entry(name).or_insert(child);
        }
    }
}

impl CodeGraph for AstGraph {
    fn children(&self, n_id: NodeId) -> Vec<NodeId> {
        self.ast_children(n_id)
    }

    fn ast_children(&self, n_id: NodeId) -> Vec<NodeId> {
        self.edges
            .get(&n_id)
            .map(|adjacent| adjacent.keys().copied().collect())
            .unwrap_or_default()
    }

    fn parents(&self, n_id: NodeId) -> Vec<NodeId> {
        self.ast_parents(n_id)
    }

    fn ast_parents(&self, n_id: NodeId) -> Vec<NodeId> {
        self.edges
            .iter()
            .filter(|(_, adjacent)| adjacent.contains_key(&n_id))
            .map(|(from, _)| *from)
            .collect()
    }

    fn cfg_children(&self, _n_id: NodeId) -> Vec<NodeId> {
        Vec::new()
    }

    fn cfg_parents(&self, _n_id: NodeId) -> Vec<NodeId> {
        Vec::new()
    }

    fn label_type(&self, n_id: NodeId) -> Option<&str> {
        self.nodes.get(&n_id).map(|node| node.kind.as_str())
    }
}

#[cfg(test)]
mod tests {
    use super::AstGraph;
    use crate::ast::AstNode;
    use crate::{Ast, CodeGraph, NodeId};
    use alloc::borrow::ToOwned;
    use alloc::vec::Vec;

    #[test]
    fn builds_parent_child_with_field_and_text() {
        let mut graph = AstGraph::new();

        graph.add_node(
            NodeId(1),
            AstNode::new(1, 1, "method_declaration".to_owned()),
        );
        let mut name = AstNode::new(1, 8, "identifier".to_owned());
        name.text = Some("foo".to_owned());
        graph.add_node(NodeId(2), name);

        graph.add_edge(NodeId(1), NodeId(2), 0);
        graph.set_field(NodeId(1), "name".to_owned(), NodeId(2));

        assert_eq!(graph.nodes.len(), 2);
        assert_eq!(
            graph.nodes.get(&NodeId(2)).and_then(|n| n.text.as_deref()),
            Some("foo")
        );
        assert_eq!(
            graph
                .nodes
                .get(&NodeId(1))
                .and_then(|n| n.fields.get("name")),
            Some(&NodeId(2))
        );
        let edge = graph
            .edges
            .get(&NodeId(1))
            .and_then(|m| m.get(&NodeId(2)))
            .copied();
        assert_eq!(edge.map(|e| e.index), Some(0));
        assert_eq!(edge.map(|e| e.kind), Some(Ast));
    }

    #[test]
    fn set_field_on_unknown_parent_is_noop() {
        let mut graph = AstGraph::new();
        graph.set_field(NodeId(99), "name".to_owned(), NodeId(1));
        assert!(graph.nodes.is_empty());
    }

    #[test]
    fn set_field_keeps_first_write() {
        let mut graph = AstGraph::new();
        graph.add_node(NodeId(1), AstNode::new(1, 1, "x".to_owned()));
        graph.set_field(NodeId(1), "name".to_owned(), NodeId(2));
        graph.set_field(NodeId(1), "name".to_owned(), NodeId(3));
        assert_eq!(
            graph
                .nodes
                .get(&NodeId(1))
                .and_then(|n| n.fields.get("name")),
            Some(&NodeId(2))
        );
    }

    #[test]
    fn iteration_order_is_node_id_order() {
        let mut graph = AstGraph::new();
        graph.add_node(NodeId(2), AstNode::new(1, 1, "b".to_owned()));
        graph.add_node(NodeId(10), AstNode::new(1, 1, "c".to_owned()));
        graph.add_node(NodeId(1), AstNode::new(1, 1, "a".to_owned()));

        let ids: Vec<u64> = graph.nodes.keys().map(|n| n.0).collect();
        assert_eq!(ids, [1, 2, 10]);
    }

    #[test]
    fn ast_children_are_all_edges_in_id_order() {
        let mut graph = AstGraph::new();
        graph.add_node(NodeId(1), AstNode::new(1, 1, "root".to_owned()));
        graph.add_edge(NodeId(1), NodeId(3), 1);
        graph.add_edge(NodeId(1), NodeId(2), 0);

        assert_eq!(graph.ast_children(NodeId(1)), [NodeId(2), NodeId(3)]);
        assert_eq!(graph.ast_children(NodeId(9)), []);
        assert_eq!(graph.label_type(NodeId(1)), Some("root"));
        assert_eq!(graph.label_type(NodeId(9)), None);
    }
}