wowlab_common/

mermaid.rs

// #t(file: rust_alloc_in_loop) graph builder produces owned strings per element, unavoidable

use crate::markdown;

/// Direction for a mermaid flowchart.
#[derive(Clone, Copy, Debug)]
// #t(rust_non_exhaustive_on_public) LR and TB are the only mermaid flowchart directions
pub enum Direction {
    LR,
    TB,
}

impl Direction {
    const fn into_str(self) -> &'static str {
        match self {
            Self::LR => "LR",
            Self::TB => "TB",
        }
    }
}

/// Builder for mermaid flowchart diagrams.
#[derive(Debug)]
pub struct Graph {
    direction: Direction,
    nodes: Vec<Box<str>>,
    edges: Vec<(Box<str>, Box<str>)>,
}

impl Graph {
    /// Create a new graph with the given direction.
    #[must_use]
    pub fn new(direction: Direction) -> Self {
        Self {
            direction,
            nodes: Vec::new(),
            edges: Vec::new(),
        }
    }

    /// Add a standalone node (rendered only if it has no edges).
    #[must_use]
    pub fn node(mut self, name: impl Into<String>) -> Self {
        self.nodes.push(name.into().into_boxed_str());

        self
    }

    /// Add multiple standalone nodes.
    #[must_use]
    pub fn nodes<I, T>(mut self, names: I) -> Self
    where
        I: IntoIterator<Item = T>,
        T: Into<String>,
    {
        self.nodes
            .extend(names.into_iter().map(|name| name.into().into_boxed_str()));

        self
    }

    /// Add a directed edge from `from` to `to`.
    #[must_use]
    pub fn edge(mut self, from: impl Into<String>, to: impl Into<String>) -> Self {
        self.edges
            .push((from.into().into_boxed_str(), to.into().into_boxed_str()));

        self
    }

    /// Add multiple edges.
    #[must_use]
    pub fn edges<I, F, T>(mut self, edges: I) -> Self
    where
        I: IntoIterator<Item = (F, T)>,
        F: Into<String>,
        T: Into<String>,
    {
        self.edges.extend(
            edges
                .into_iter()
                .map(|(from, to)| (from.into().into_boxed_str(), to.into().into_boxed_str())),
        );

        self
    }

    /// Render as a mermaid code block, or empty string if there are no edges.
    #[must_use]
    pub fn build(mut self) -> String {
        self.edges.sort();
        self.edges.dedup();

        if self.edges.is_empty() {
            return String::new();
        }

        let mut connected: std::collections::BTreeSet<&str> = std::collections::BTreeSet::new();

        for (from, to) in &self.edges {
            connected.insert(from.as_ref());
            connected.insert(to.as_ref());
        }

        let mut lines = vec![format!("graph {}", self.direction.into_str())];

        self.nodes.sort();
        self.nodes.dedup();

        for node in &self.nodes {
            if !connected.contains(node.as_ref()) {
                lines.push(format!("    {node}"));
            }
        }

        for (from, to) in &self.edges {
            lines.push(format!("    {from} --> {to}"));
        }

        markdown::code_block("mermaid", lines.join("\n"))
    }
}

impl Default for Graph {
    fn default() -> Self {
        Self::new(Direction::LR)
    }
}

#[cfg(test)]
mod tests {
    use googletest::prelude::*;

    use super::*;

    #[gtest]
    fn empty_graph_returns_empty_string() -> Result<()> {
        let result = Graph::new(Direction::LR).build();

        verify_that!(result, eq(""))?;

        Ok(())
    }

    #[gtest]
    fn single_edge() -> Result<()> {
        let result = Graph::new(Direction::LR).edge("a", "b").build();

        verify_that!(result, contains_substring("graph LR"))?;
        verify_that!(result, contains_substring("    a --> b"))?;

        Ok(())
    }

    #[gtest]
    fn standalone_nodes_without_edges() -> Result<()> {
        let result = Graph::new(Direction::LR).node("orphan").build();

        verify_that!(result, eq(""))?;

        Ok(())
    }

    #[gtest]
    fn standalone_nodes_with_edges() -> Result<()> {
        let result = Graph::new(Direction::LR)
            .node("orphan")
            .node("a")
            .edge("a", "b")
            .build();

        verify_that!(result, contains_substring("    orphan"))?;
        verify_that!(result, contains_substring("    a --> b"))?;
        let standalone_lines: Vec<&str> = result.lines().filter(|l| l.trim() == "a").collect();

        verify_that!(standalone_lines, is_empty())?;

        Ok(())
    }

    #[gtest]
    fn edges_are_sorted_and_deduped() -> Result<()> {
        let result = Graph::new(Direction::LR)
            .edge("b", "c")
            .edge("a", "b")
            .edge("b", "c")
            .build();
        let edge_lines: Vec<&str> = result.lines().filter(|l| l.contains("-->")).collect();

        verify_that!(
            edge_lines,
            elements_are![contains_substring("a --> b"), contains_substring("b --> c")]
        )?;

        Ok(())
    }

    #[gtest]
    fn top_to_bottom_direction() -> Result<()> {
        let result = Graph::new(Direction::TB).edge("x", "y").build();

        verify_that!(result, contains_substring("graph TB"))?;

        Ok(())
    }

    #[gtest]
    fn bulk_edges() -> Result<()> {
        let result = Graph::new(Direction::LR)
            .edges([("a", "b"), ("c", "d")])
            .build();

        verify_that!(result, contains_substring("a --> b"))?;
        verify_that!(result, contains_substring("c --> d"))?;

        Ok(())
    }

    #[gtest]
    fn wrapped_in_mermaid_code_block() -> Result<()> {
        let result = Graph::new(Direction::LR).edge("x", "y").build();

        verify_that!(result, starts_with("```mermaid\n"))?;
        verify_that!(result, ends_with("\n```"))?;

        Ok(())
    }
}