fluidattacks-blends 0.6.0

//! Graphs as Graphviz DOT source.
//!
//! Every node carries its whole attribute map as a multi-line label, so the
//! rendered graph shows the same information the JSON export does. Edge colors
//! separate the layers: blue for AST, red for CFG, purple for both.

use std::collections::BTreeMap;

use blends_domain::ast::AstGraph;
use blends_domain::syntax::{SyntaxGraph, SyntaxNode};
use blends_domain::NodeId;
use serde_json::Value;

use crate::attrs::{ast_edge_attrs, ast_node_attrs, syntax_edge_attrs, syntax_node_attrs};

const ARROWHEAD: &str = "open";
const COLOR_AST: &str = "blue";
const COLOR_AST_CFG: &str = "purple";
const COLOR_CFG: &str = "red";

/// A graph that can describe its nodes and edges as attribute maps.
///
/// Private on purpose: it carries `serde_json::Value`, and leaking that into
/// the public surface would make a future `serde_json` major a breaking change
/// here. [`DotGraph`] is the public face, and its blanket impl over this trait
/// keeps it sealed.
trait GraphAttrs {
    fn attr_nodes(&self) -> Vec<(NodeId, BTreeMap<String, Value>)>;
    fn attr_edges(&self) -> Vec<(NodeId, NodeId, BTreeMap<String, Value>)>;
}

/// A graph that can render itself as Graphviz DOT source.
pub trait DotGraph {
    #[must_use]
    fn to_dot(&self, name: &str) -> String;
}

impl<G: GraphAttrs> DotGraph for G {
    fn to_dot(&self, name: &str) -> String {
        let mut statements = vec![format!("digraph \"{}\" {{", escape(name))];

        for (n_id, mut attrs) in self.attr_nodes() {
            attrs.insert("id".to_owned(), Value::from(n_id.0.to_string()));
            statements.push(node_statement(n_id, &create_label(&attrs)));
        }

        for (from, to, attrs) in self.attr_edges() {
            statements.push(edge_statement(
                from,
                to,
                &create_label(&attrs),
                edge_color(&attrs),
            ));
        }

        statements.push("}".to_owned());
        statements.push(String::new());

        statements.join("\n")
    }
}

impl GraphAttrs for AstGraph {
    fn attr_nodes(&self) -> Vec<(NodeId, BTreeMap<String, Value>)> {
        self.nodes
            .iter()
            .map(|(id, node)| (*id, ast_node_attrs(node)))
            .collect()
    }

    fn attr_edges(&self) -> Vec<(NodeId, NodeId, BTreeMap<String, Value>)> {
        self.edges
            .iter()
            .flat_map(|(from, targets)| {
                targets
                    .iter()
                    .map(move |(to, edge)| (*from, *to, ast_edge_attrs(*edge)))
            })
            .collect()
    }
}

impl GraphAttrs for SyntaxGraph {
    fn attr_nodes(&self) -> Vec<(NodeId, BTreeMap<String, Value>)> {
        self.nodes
            .iter()
            .map(|(id, node)| {
                (
                    *id,
                    syntax_node_attrs(node).unwrap_or_else(|| fallback(node)),
                )
            })
            .collect()
    }

    fn attr_edges(&self) -> Vec<(NodeId, NodeId, BTreeMap<String, Value>)> {
        self.edges
            .iter()
            .flat_map(|(from, targets)| {
                targets
                    .iter()
                    .map(move |(to, edge)| (*from, *to, syntax_edge_attrs(*edge)))
            })
            .collect()
    }
}

/// A node type with no attribute export yet still renders, by its type alone.
/// Dropping it would leave a dangling edge in the picture.
fn fallback(node: &SyntaxNode) -> BTreeMap<String, Value> {
    let mut attrs = BTreeMap::new();
    attrs.insert("label_type".to_owned(), Value::from(node.label_type()));
    attrs
}

fn escape(value: &str) -> String {
    value
        .replace('\\', "\\\\")
        .replace('"', "\\\"")
        .replace('\r', "")
        .replace('\n', "\\n")
}

/// `serde_json` renders strings with their quotes; a label wants the text.
fn display(value: &Value) -> String {
    match value {
        Value::String(text) => text.clone(),
        other => other.to_string(),
    }
}

fn create_label(attrs: &BTreeMap<String, Value>) -> String {
    attrs
        .iter()
        .map(|(key, value)| format!("{key}: {}", display(value)))
        .collect::<Vec<_>>()
        .join("\n")
}

fn edge_color(attrs: &BTreeMap<String, Value>) -> Option<&'static str> {
    match (
        attrs.contains_key("label_ast"),
        attrs.contains_key("label_cfg"),
    ) {
        (true, true) => Some(COLOR_AST_CFG),
        (true, false) => Some(COLOR_AST),
        (false, true) => Some(COLOR_CFG),
        (false, false) => None,
    }
}

fn render_attrs(attrs: &[(&str, &str)]) -> String {
    attrs
        .iter()
        .map(|(key, value)| format!("{key}=\"{}\"", escape(value)))
        .collect::<Vec<_>>()
        .join(", ")
}

fn node_statement(n_id: NodeId, label: &str) -> String {
    format!("  \"{}\" [{}];", n_id.0, render_attrs(&[("label", label)]))
}

fn edge_statement(from: NodeId, to: NodeId, label: &str, color: Option<&str>) -> String {
    let mut attrs = vec![("label", label), ("arrowhead", ARROWHEAD)];
    if let Some(color) = color {
        attrs.push(("color", color));
    }

    format!(
        "  \"{}\" -> \"{}\" [{}];",
        from.0,
        to.0,
        render_attrs(&attrs)
    )
}

#[cfg(test)]
mod tests {
    use super::{DotGraph, GraphAttrs};
    use blends_domain::ast::{AstGraph, AstNode};
    use blends_domain::syntax::{SyntaxEdge, SyntaxGraph, SyntaxNode};
    use blends_domain::NodeId;

    fn ast_graph() -> AstGraph {
        let mut graph = AstGraph::new();
        graph.add_node(NodeId(1), AstNode::new(1, 0, "module".to_owned()));
        let mut child = AstNode::new(2, 4, "identifier".to_owned());
        child.text = Some("say \"hi\"\nand bye".to_owned());
        graph.add_node(NodeId(2), child);
        graph.add_edge(NodeId(1), NodeId(2), 0);

        graph
    }

    /// Nodes 1..=3 with an AST-only, a CFG-only, and a both-layers edge.
    fn syntax_graph() -> SyntaxGraph {
        let mut graph = SyntaxGraph::new();
        graph.add_node(NodeId(1), SyntaxNode::File);
        graph.add_node(NodeId(2), SyntaxNode::ArgumentList);
        graph.add_node(NodeId(3), SyntaxNode::ArgumentList);

        graph.add_ast_edge(NodeId(1), NodeId(2));
        graph.add_cfg_edge(NodeId(2), NodeId(3));
        graph.add_ast_edge(NodeId(1), NodeId(3));
        graph.add_cfg_edge(NodeId(1), NodeId(3));

        graph
    }

    #[test]
    fn emits_a_labeled_digraph() {
        let source = ast_graph().to_dot("python_py");

        assert!(source.starts_with("digraph \"python_py\" {\n"));
        assert!(source.ends_with("}\n"));
        assert!(source.contains("label_type: module"));
    }

    #[test]
    fn labels_carry_the_node_id() {
        assert!(ast_graph().to_dot("g").contains("id: 1"));
    }

    #[test]
    fn escapes_quotes_and_newlines() {
        let source = ast_graph().to_dot("g");

        assert!(source.contains("say \\\"hi\\\"\\nand bye"));
        assert!(!source.contains("say \"hi\"\nand bye"));
    }

    #[test]
    fn colors_ast_edges_blue() {
        let source = ast_graph().to_dot("g");

        assert!(source.contains("arrowhead=\"open\", color=\"blue\""));
        assert!(source.contains("\"1\" -> \"2\""));
    }

    #[test]
    fn every_label_line_reaches_the_source() {
        let graph = ast_graph();
        let source = graph.to_dot("g");

        for (n_id, attrs) in graph.attr_nodes() {
            for key in attrs.keys() {
                assert!(source.contains(&format!("{key}: ")), "{key} missing");
            }
            assert!(source.contains(&format!("\"{}\" [", n_id.0)));
        }
    }

    #[test]
    fn renders_numeric_attributes_without_quotes() {
        let mut graph = ast_graph();
        graph.set_field(NodeId(1), "body".to_owned(), NodeId(2));

        assert!(graph.to_dot("g").contains("body: 2"));
    }

    #[test]
    fn renders_the_syntax_graph_by_node_type() {
        let source = syntax_graph().to_dot("python_py");

        assert!(source.starts_with("digraph \"python_py\" {\n"));
        assert!(source.contains("label_type: File"));
        assert!(source.contains("label_type: ArgumentList"));
    }

    #[test]
    fn colors_syntax_edges_by_layer() {
        let source = syntax_graph().to_dot("g");

        assert!(source.contains("label_ast: AST\\nlabel_cfg: CFG"));
        assert!(source.contains("color=\"purple\""));
        assert!(source.contains("color=\"blue\""));
        assert!(source.contains("color=\"red\""));
    }

    #[test]
    fn leaves_a_layerless_edge_uncolored() {
        let mut graph = SyntaxGraph::new();
        graph.add_node(NodeId(1), SyntaxNode::File);
        graph.add_node(NodeId(2), SyntaxNode::ArgumentList);
        graph.edges.entry(NodeId(1)).or_default().insert(
            NodeId(2),
            SyntaxEdge {
                ast: None,
                cfg: None,
            },
        );

        let source = graph.to_dot("g");

        assert!(source.contains("\"1\" -> \"2\" [label=\"\", arrowhead=\"open\"];"));
        assert!(!source.contains("color="));
    }

    /// `SyntaxNode::Default` is one of the variants with no attribute export, so
    /// it exercises the fallback rather than `syntax_node_attrs`.
    #[test]
    fn renders_an_unexported_node_type_by_type_alone() {
        let mut graph = SyntaxGraph::new();
        graph.add_node(NodeId(1), SyntaxNode::Default);

        let source = graph.to_dot("g");

        assert!(source.contains("label_type: Default"));
        assert!(source.contains("id: 1"));
    }
}