fluidattacks-blends-domain 0.2.0

//! Syntax graph: node table plus AST-mirroring adjacency.

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

use crate::syntax::{SyntaxEdge, SyntaxKind, SyntaxNode};
use crate::{Ast, Cfg, CodeGraph, NodeId};

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct SanitizationEvent {
    pub node_id: NodeId,
    pub kind: String,
    pub subpath: Option<NodeId>,
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum UsageRole {
    Receiver,
    Argument,
    MemberWrite,
}

#[derive(Clone, PartialEq, Eq, Debug)]
pub struct UsageEvent {
    pub node_id: NodeId,
    pub role: UsageRole,
    pub arg_index: Option<i64>,
    pub subpath: Option<NodeId>,
}

#[derive(Clone, PartialEq, Eq, Debug, Default)]
pub struct SyntaxGraph {
    pub nodes: BTreeMap<NodeId, SyntaxNode>,
    pub edges: BTreeMap<NodeId, BTreeMap<NodeId, SyntaxEdge>>,
    pub symbol_index: BTreeMap<String, BTreeMap<NodeId, Vec<NodeId>>>,
    pub scope_parent: BTreeMap<NodeId, NodeId>,
    pub sanitization_index: BTreeMap<String, BTreeMap<NodeId, Vec<SanitizationEvent>>>,
    pub usage_index: BTreeMap<String, BTreeMap<NodeId, Vec<UsageEvent>>>,
    pub subpath_index: BTreeMap<NodeId, Option<NodeId>>,
    pub nodes_by_type: BTreeMap<SyntaxKind, Vec<NodeId>>,
}

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

    pub fn add_node(&mut self, id: NodeId, node: SyntaxNode) {
        self.nodes_by_type
            .entry(SyntaxKind::from(&node))
            .or_default()
            .push(id);
        self.nodes.insert(id, node);
    }

    pub fn register_sanitization(&mut self, symbol: &str, scope: NodeId, event: SanitizationEvent) {
        self.sanitization_index
            .entry(String::from(symbol))
            .or_default()
            .entry(scope)
            .or_default()
            .push(event);
    }

    pub fn register_usage(&mut self, symbol: &str, scope: NodeId, event: UsageEvent) {
        self.usage_index
            .entry(String::from(symbol))
            .or_default()
            .entry(scope)
            .or_default()
            .push(event);
    }

    pub fn add_ast_edge(&mut self, from: NodeId, to: NodeId) {
        self.edges
            .entry(from)
            .or_default()
            .entry(to)
            .or_default()
            .ast = Some(Ast);
    }

    pub fn add_cfg_edge(&mut self, from: NodeId, to: NodeId) {
        self.edges
            .entry(from)
            .or_default()
            .entry(to)
            .or_default()
            .cfg = Some(Cfg);
    }

    pub fn finalize_symbol_index(&mut self) {
        for scope_definitions in self.symbol_index.values_mut() {
            for definition_list in scope_definitions.values_mut() {
                definition_list.reverse();
            }
        }
    }

    pub fn remove_edge(&mut self, from: NodeId, to: NodeId) {
        if let Some(adjacent) = self.edges.get_mut(&from) {
            adjacent.remove(&to);
            if adjacent.is_empty() {
                self.edges.remove(&from);
            }
        }
    }
}

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

    fn ast_children(&self, n_id: NodeId) -> Vec<NodeId> {
        self.edges
            .get(&n_id)
            .map(|adjacent| {
                adjacent
                    .iter()
                    .filter(|(_, edge)| edge.ast.is_some())
                    .map(|(to, _)| *to)
                    .collect()
            })
            .unwrap_or_default()
    }

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

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

    fn cfg_parents(&self, n_id: NodeId) -> Vec<NodeId> {
        self.edges
            .iter()
            .filter(|(_, adjacent)| adjacent.get(&n_id).is_some_and(|edge| edge.cfg.is_some()))
            .map(|(from, _)| *from)
            .collect()
    }

    fn cfg_children(&self, n_id: NodeId) -> Vec<NodeId> {
        self.edges
            .get(&n_id)
            .map(|adjacent| {
                adjacent
                    .iter()
                    .filter(|(_, edge)| edge.cfg.is_some())
                    .map(|(to, _)| *to)
                    .collect()
            })
            .unwrap_or_default()
    }

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

#[cfg(test)]
mod tests {
    use super::SyntaxGraph;
    use crate::syntax::{SyntaxEdge, SyntaxKind, SyntaxNode};
    use crate::{Ast, Cfg, CodeGraph, NodeId};
    use alloc::borrow::ToOwned;

    #[test]
    fn stores_nodes_and_edges_by_id() {
        let mut graph = SyntaxGraph::new();
        graph.add_node(
            NodeId(1),
            SyntaxNode::MissingNode {
                node_type: "stream".to_owned(),
            },
        );
        graph.add_ast_edge(NodeId(1), NodeId(2));

        assert_eq!(
            graph.nodes.get(&NodeId(1)),
            Some(&SyntaxNode::MissingNode {
                node_type: "stream".to_owned()
            })
        );
        assert_eq!(
            graph
                .edges
                .get(&NodeId(1))
                .and_then(|adjacent| adjacent.get(&NodeId(2))),
            Some(&SyntaxEdge {
                ast: Some(Ast),
                cfg: None
            })
        );
    }

    #[test]
    fn add_node_indexes_ids_by_syntax_type_in_insertion_order() {
        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!(
            graph.nodes_by_type.get(&SyntaxKind::Break),
            Some(&alloc::vec![NodeId(9), NodeId(2)])
        );
        assert_eq!(
            graph.nodes_by_type.get(&SyntaxKind::File),
            Some(&alloc::vec![NodeId(1)])
        );
        assert_eq!(graph.nodes_by_type.get(&SyntaxKind::If), None);
    }

    #[test]
    fn ast_edge_merges_into_existing_edge() {
        let mut graph = SyntaxGraph::new();
        graph.add_ast_edge(NodeId(1), NodeId(2));
        graph.add_ast_edge(NodeId(1), NodeId(2));

        assert_eq!(
            graph
                .edges
                .get(&NodeId(1))
                .map(alloc::collections::BTreeMap::len),
            Some(1)
        );
    }

    #[test]
    fn ast_children_skip_edges_without_the_ast_mark() {
        let mut graph = SyntaxGraph::new();
        graph.add_ast_edge(NodeId(1), NodeId(2));
        graph.edges.entry(NodeId(1)).or_default().insert(
            NodeId(3),
            SyntaxEdge {
                ast: None,
                cfg: None,
            },
        );

        assert_eq!(graph.ast_children(NodeId(1)), [NodeId(2)]);
    }

    #[test]
    fn cfg_edge_merges_onto_the_ast_edge() {
        let mut graph = SyntaxGraph::new();
        graph.add_ast_edge(NodeId(1), NodeId(2));
        graph.add_cfg_edge(NodeId(1), NodeId(2));

        assert_eq!(
            graph
                .edges
                .get(&NodeId(1))
                .and_then(|adjacent| adjacent.get(&NodeId(2))),
            Some(&SyntaxEdge {
                ast: Some(Ast),
                cfg: Some(Cfg)
            })
        );
    }

    #[test]
    fn remove_edge_drops_the_edge_and_prunes_the_empty_entry() {
        let mut graph = SyntaxGraph::new();
        graph.add_ast_edge(NodeId(1), NodeId(2));
        graph.add_ast_edge(NodeId(1), NodeId(3));
        graph.add_cfg_edge(NodeId(1), NodeId(2));

        graph.remove_edge(NodeId(1), NodeId(2));

        assert_eq!(graph.children(NodeId(1)), [NodeId(3)]);
        assert_eq!(graph.parents(NodeId(2)), []);

        graph.remove_edge(NodeId(1), NodeId(3));

        assert!(!graph.edges.contains_key(&NodeId(1)));
    }

    #[test]
    fn finalize_symbol_index_puts_the_most_recent_definition_first() {
        let mut graph = SyntaxGraph::new();
        graph
            .symbol_index
            .entry("x".to_owned())
            .or_default()
            .entry(NodeId(1))
            .or_default()
            .extend([NodeId(2), NodeId(5), NodeId(9)]);

        graph.finalize_symbol_index();

        assert_eq!(
            graph
                .symbol_index
                .get("x")
                .and_then(|scopes| scopes.get(&NodeId(1))),
            Some(&alloc::vec![NodeId(9), NodeId(5), NodeId(2)])
        );
    }

    #[test]
    fn remove_edge_ignores_a_missing_edge() {
        let mut graph = SyntaxGraph::new();
        graph.add_ast_edge(NodeId(1), NodeId(2));

        graph.remove_edge(NodeId(7), NodeId(8));
        graph.remove_edge(NodeId(1), NodeId(9));

        assert_eq!(graph.children(NodeId(1)), [NodeId(2)]);
    }

    #[test]
    fn cfg_children_skip_edges_without_the_cfg_mark() {
        let mut graph = SyntaxGraph::new();
        graph.add_ast_edge(NodeId(1), NodeId(2));
        graph.add_ast_edge(NodeId(1), NodeId(3));
        graph.add_cfg_edge(NodeId(1), NodeId(3));

        assert_eq!(graph.cfg_children(NodeId(1)), [NodeId(3)]);
        assert_eq!(graph.cfg_children(NodeId(9)), []);
    }
}