wolf-graph-mermaid 0.1.0

use wolf_graph::prelude::*;

use crate::{details::arrow, Arrowhead, EdgeAttributes, EdgeStyle, GraphAttributes, GroupAttributes, LayoutDirection, NodeAttributes, NodeShape};

// Docs: https://mermaid.js.org/intro/
// Editor: https://mermaid.live/

pub trait MermaidEncodable<GData, NData, EData>: VisitableGraph + Sized {
    fn mermaid_graph_attributes(&self) -> Option<GraphAttributes> {
        None
    }

    fn mermaid_node_attributes(&self, _node: &NodeID) -> Option<NodeAttributes> {
        None
    }

    fn mermaid_edge_attributes(&self, _edge: &EdgeID) -> Option<EdgeAttributes> {
        None
    }

    // This should only return `true` if the graph is a compound graph.
    fn mermaid_is_compound(&self) -> bool {
        false
    }

    // This returns the children of a group node.
    fn mermaid_children(&self, _node: Option<&NodeID>) -> Nodes {
        Nodes::new()
    }

    // Default implementation presumes that if a node has children, it is a group.
    // Override if you wish to, for example, support empty groups.
    fn mermaid_is_group(&self, node: &NodeID) -> bool {
        !self.mermaid_children(Some(node)).is_empty()
    }

    fn mermaid_group_attributes(&self, _group: &NodeID) -> Option<GroupAttributes> {
        None
    }
}

pub trait MermaidFormat<GData, NData, EData>: MermaidEncodable<GData, NData, EData> {
    fn mermaid_format(&self) -> String {
        let mut lines = Vec::new();
        format_graph(self, &mut lines);
        format_lines(&lines)
    }
}

// Blanket implementation of `MermaidFormat` for all types that implement `MermaidEncodable`.
impl<GData, NData, EData, T> MermaidFormat<GData, NData, EData> for T where T: MermaidEncodable<GData, NData, EData> {}

fn format_graph<GData, NData, EData>(
    graph: &impl MermaidEncodable<GData, NData, EData>,
    lines: &mut Vec<(usize, String)>
) {
    let indent = 0;

    if let Some(title) = graph
        .mermaid_graph_attributes()
        .and_then(|attrs| attrs.title)
        .map(|title| title.to_string())
    {
        lines.push((0, "---".to_string()));
        lines.push((0, format!("title: {}", title)));
        lines.push((0, "---".to_string()));
    }

    let layout_direction = graph
        .mermaid_graph_attributes()
        .and_then(|attrs| attrs.layout_direction)
        .unwrap_or(LayoutDirection::TopToBottom);
    lines.push((0, format!("graph {}", layout_direction)));

    let mut custom_styles: Vec<String> = vec![];

    if graph.mermaid_is_compound() {
        format_group(graph, lines, &mut custom_styles, None, indent);
    } else {
        format_nodes(graph, lines, &mut custom_styles, &graph.all_nodes(), indent + 1);
    }
    format_edges(graph, lines, &mut custom_styles, &graph.all_edges());

    lines.append(&mut custom_styles.iter().map(|s| (1, s.clone())).collect());
}

fn format_group<GData, NData, EData>(
    graph: &impl MermaidEncodable<GData, NData, EData>,
    lines: &mut Vec<(usize, String)>,
    custom_styles: &mut Vec<String>,
    parent: Option<&NodeID>,
    indent: usize
) {
    let children = graph.mermaid_children(parent);
    format_nodes(graph, lines, custom_styles, &children, indent + 1);
}

fn format_subgraph<GData, NData, EData>(graph: &impl MermaidEncodable<GData, NData, EData>, lines: &mut Vec<(usize, String)>, custom_styles: &mut Vec<String>, parent: &NodeID, indent: usize) {
    let mut components = vec![format!("subgraph {}", parent)];
    if let Some(attributes) = graph.mermaid_group_attributes(parent) {
        if let Some(label) = attributes.label.as_ref() {
            components.push(format!("[{}]", label));
        }
        if let Some(custom_style) = attributes.custom_style() {
            custom_styles.push(format!("style {} {}", parent, custom_style));
        }
    }
    lines.push((indent, components.join(" ")));
    format_group(graph, lines, custom_styles, Some(parent), indent);
    lines.push((indent, "end".to_string()));
}

fn format_nodes<GData, NData, EData>(graph: &impl MermaidEncodable<GData, NData, EData>, lines: &mut Vec<(usize, String)>, custom_styles: &mut Vec<String>, nodes: &Nodes, indent: usize) {
    for node in nodes {
        if graph.mermaid_is_group(node) {
            format_subgraph(graph, lines, custom_styles, node, indent);
        } else {
            format_leaf_node(graph, lines, custom_styles, node, indent);
        }
    }
}

fn format_leaf_node<GData, NData, EData>(graph: &impl MermaidEncodable<GData, NData, EData>, lines: &mut Vec<(usize, String)>, custom_styles: &mut Vec<String>, node: &NodeID, indent: usize) {
    let attributes = graph.mermaid_node_attributes(node);

    let mut line_components = vec![node.to_string()];

    if let Some(label) = attributes.as_ref().and_then(|a| a.label.as_ref()) {
        let delimiters = attributes.as_ref().and_then(|a| a.shape.as_ref()).unwrap_or(&NodeShape::Rectangle).delimiters();
        line_components.push(format!("{}{}{}", delimiters.0, label, delimiters.1));
    }

    lines.push((indent, line_components.join("")));

    if let Some(custom_style) = attributes.and_then(|a| a.custom_style()) {
        custom_styles.push(format!("style {} {}", node, custom_style));
    }
}

fn format_edges<GData, NData, EData>(graph: &impl MermaidEncodable<GData, NData, EData>, lines: &mut Vec<(usize, String)>, custom_styles: &mut Vec<String>, edges: &Edges) {
    for (index, edge) in edges.iter().enumerate() {
        format_edge(graph, lines, custom_styles, index, edge);
    }
}

fn format_edge<GData, NData, EData>(graph: &impl MermaidEncodable<GData, NData, EData>, lines: &mut Vec<(usize, String)>, custom_styles: &mut Vec<String>, index: usize, edge: &EdgeID) {
    let attributes = graph.mermaid_edge_attributes(edge);

    let mut line_components = vec![];

    line_components.push(graph.source(edge).unwrap().to_string());
    line_components.push(" ".to_string());

    line_components.push(arrow(
        attributes.as_ref().and_then(|a| a.length).unwrap_or(1),
        &attributes.as_ref().and_then(|a| a.style).unwrap_or(EdgeStyle::Normal),
        &attributes.as_ref().and_then(|a| a.tail).unwrap_or(Arrowhead::None),
        &attributes.as_ref().and_then(|a| a.head).unwrap_or(Arrowhead::Normal),
    ));

    if let Some(label) = attributes.as_ref().and_then(|a| a.label.as_ref()) {
        line_components.push(format!("|{}|", label));
    }

    line_components.push(" ".to_string());

    line_components.push(graph.target(edge).unwrap().to_string());

    lines.push((1, line_components.join("")));

    if let Some(custom_style) = attributes.and_then(|a| a.custom_style()) {
        custom_styles.push(format!("linkStyle {} {}", index, custom_style));
    }
}

// Translated from above Swift
fn format_lines(lines: &[(usize, String)]) -> String {
    lines
        .iter()
        .map(|(indent, s)| format!("{:indent$}{}", "", s, indent = indent * 4))
        .collect::<Vec<String>>()
        .join("\n")
}