polydat_grammar/

viz.rs

// Copyright 2024-2026 Jonathan Shook
// SPDX-License-Identifier: Apache-2.0

//! DAG visualization for Polydat Kernels.
//!
//! Renders a Polydat Kernel's node graph as DOT (with record nodes and
//! port-based edge routing), Mermaid, or self-contained SVG.
//!
//! The DOT output uses graphviz record syntax:
//! - Each node has named input ports (top) and output ports (bottom)
//! - Edges connect from output ports to input ports
//! - Dark theme colors via graph/node/edge attributes
//!
//! The program's inputs are drawn as port nodes at the top of the graph,
//! distinct from the binding nodes that read them:
//! - Coordinates (`input name: type`) share one `INPUTS` register with one
//!   output port per coordinate.
//! - External ports (`extern name: type [= default]`, the kernel's
//!   `InputKind::ExternalWrite` slots the host writes between cycles) each
//!   get their own register labeled with the port's kind, name, type, and
//!   default, with one output port wired to every node that reads it.

use std::collections::{HashMap, HashSet};

use crate::ast::*;
use crate::{lexer, parser};

/// How a visualization node is drawn.
#[derive(Clone, Copy, PartialEq, Eq)]
enum VizKind {
    /// A binding: input ports | label | output ports.
    Func,
    /// The `INPUTS` or `OUTPUTS` register: coordinates or terminal wires
    /// as ports, with the blue or green accent.
    Register,
    /// One external port (`extern name: type [= default]`): a register
    /// with a single output port and the amber accent.
    Extern,
}

/// A node in the visualization graph.
struct VizNode {
    /// Unique ID for DOT (e.g., "n0", "n1", "x0")
    id: String,
    /// Display label (function name or binding expression)
    label: String,
    /// Input wire names (from upstream nodes/coords)
    inputs: Vec<String>,
    /// Output wire names (what this node produces)
    outputs: Vec<String>,
    /// How the node is drawn.
    kind: VizKind,
}

/// An edge connecting an output to an input.
struct VizEdge {
    from_node: String,
    from_port: String,
    to_node: String,
    to_port: String,
}

/// Render a Polydat source string as DOT with record nodes and ports.
pub fn polydat_to_dot(source: &str) -> Result<String, String> {
    let (nodes, edges) = build_graph(source)?;
    let mut dot = String::new();

    dot.push_str("digraph Polydat {\n");
    dot.push_str("    rankdir=TB;\n");
    dot.push_str("    bgcolor=\"#1a1a2e\";\n");
    dot.push_str("    node [shape=record, style=filled, fontname=\"monospace\", fontsize=11];\n");
    dot.push_str("    edge [color=\"#4da6ff\", fontcolor=\"#8888a0\", fontname=\"monospace\", fontsize=9];\n");
    dot.push('\n');

    for node in &nodes {
        if node.kind == VizKind::Extern {
            // External port: output port only (top of graph), amber accent
            let ports: Vec<String> = node
                .outputs
                .iter()
                .map(|name| format!("<o_{name}> {name}"))
                .collect();
            dot.push_str(&format!(
                "    {} [label=\"{{ {} | {{ {} }} }}\", fillcolor=\"#0f3460\", \
                 fontcolor=\"#ffb454\", color=\"#ffb454\", penwidth=2];\n",
                node.id,
                dot_escape(&node.label),
                ports.join(" | "),
            ));
        } else if node.kind == VizKind::Register {
            // Register nodes (INPUTS / OUTPUTS): record with labeled ports
            if node.outputs.is_empty() && !node.inputs.is_empty() {
                // OUTPUTS register: input ports only (bottom of graph)
                let ports: Vec<String> = node
                    .inputs
                    .iter()
                    .map(|name| format!("<i_{name}> {name}"))
                    .collect();
                dot.push_str(&format!(
                    "    {} [label=\"{{ {{ {} }} | {} }}\", fillcolor=\"#0f3460\", \
                     fontcolor=\"#4ecca3\", color=\"#4ecca3\", penwidth=2];\n",
                    node.id,
                    ports.join(" | "),
                    dot_escape(&node.label),
                ));
            } else if !node.outputs.is_empty() && node.inputs.is_empty() {
                // INPUTS register: output ports only (top of graph)
                let ports: Vec<String> = node
                    .outputs
                    .iter()
                    .map(|name| format!("<o_{name}> {name}"))
                    .collect();
                dot.push_str(&format!(
                    "    {} [label=\"{{ {} | {{ {} }} }}\", fillcolor=\"#0f3460\", \
                     fontcolor=\"#4da6ff\", color=\"#4da6ff\", penwidth=2];\n",
                    node.id,
                    dot_escape(&node.label),
                    ports.join(" | "),
                ));
            } else {
                // Fallback
                dot.push_str(&format!(
                    "    {} [label=\"{}\", shape=oval, fillcolor=\"#16213e\", \
                     fontcolor=\"#4da6ff\", color=\"#4da6ff\"];\n",
                    node.id,
                    dot_escape(&node.label)
                ));
            }
        } else {
            // Function nodes: record with input ports | label | output ports
            let input_ports = if node.inputs.is_empty() {
                String::new()
            } else {
                let ports: Vec<String> = node
                    .inputs
                    .iter()
                    .map(|name| format!("<i_{name}> {name}"))
                    .collect();
                format!("{{ {} }} | ", ports.join(" | "))
            };

            let output_ports = if node.outputs.is_empty() {
                String::new()
            } else {
                let ports: Vec<String> = node
                    .outputs
                    .iter()
                    .map(|name| format!("<o_{name}> {name}"))
                    .collect();
                format!(" | {{ {} }}", ports.join(" | "))
            };

            dot.push_str(&format!(
                "    {} [label=\"{}{}{}\", fillcolor=\"#16213e\", \
                 fontcolor=\"#e0e0e0\", color=\"#0f3460\"];\n",
                node.id,
                input_ports,
                dot_escape(&node.label),
                output_ports,
            ));
        }
    }

    dot.push('\n');

    for edge in &edges {
        let from = if edge.from_port.is_empty() {
            edge.from_node.clone()
        } else {
            format!("{}:o_{}", edge.from_node, edge.from_port)
        };
        let to = if edge.to_port.is_empty() {
            edge.to_node.clone()
        } else {
            format!("{}:i_{}", edge.to_node, edge.to_port)
        };
        dot.push_str(&format!("    {} -> {};\n", from, to));
    }

    dot.push_str("}\n");
    Ok(dot)
}

/// Render a Polydat source string as a Mermaid flowchart.
pub fn polydat_to_mermaid(source: &str) -> Result<String, String> {
    let (nodes, edges) = build_graph(source)?;
    let mut lines = vec!["flowchart TD".to_string()];

    for node in &nodes {
        let escaped = node.label.replace('"', "'");
        if node.kind == VizKind::Func {
            lines.push(format!("    {}[\"{}\"]", node.id, escaped));
        } else {
            lines.push(format!("    {}([\"{}\"])", node.id, escaped));
        }
    }

    for edge in &edges {
        let label = if edge.from_port.is_empty() && edge.to_port.is_empty() {
            String::new()
        } else {
            let port_name = if !edge.from_port.is_empty() {
                &edge.from_port
            } else {
                &edge.to_port
            };
            format!("|{}|", port_name)
        };
        lines.push(format!(
            "    {} -->{} {}",
            edge.from_node, label, edge.to_node
        ));
    }

    lines.push("    classDef coord fill:#16213e,stroke:#4da6ff,color:#4da6ff".into());
    lines.push("    classDef func fill:#16213e,stroke:#0f3460,color:#e0e0e0".into());
    if nodes.iter().any(|n| n.kind == VizKind::Extern) {
        lines.push("    classDef extern fill:#16213e,stroke:#ffb454,color:#ffb454".into());
    }
    for node in &nodes {
        let class = match node.kind {
            VizKind::Register => "input",
            VizKind::Extern => "extern",
            VizKind::Func => "func",
        };
        lines.push(format!("    class {} {class}", node.id));
    }

    Ok(lines.join("\n"))
}

/// Render a Polydat source string as self-contained SVG.
///
/// Generates DOT with record nodes and port syntax, then renders
/// through layout-rs (pure Rust, no external graphviz needed).
/// Layout-rs supports record shapes and port-based edge routing.
pub fn polydat_to_svg(source: &str) -> Result<String, String> {
    let dot_source = polydat_to_dot(source)?;

    // Parse DOT through layout-rs
    let mut parser = layout::gv::DotParser::new(&dot_source);
    let graph = parser
        .process()
        .map_err(|e| format!("DOT parse error: {e}"))?;

    // Build visual graph from parsed DOT
    let mut builder = layout::gv::GraphBuilder::new();
    builder.visit_graph(&graph);
    let mut visual = builder.get();

    // Layout and render to SVG
    let mut svg_writer = layout::backends::svg::SVGWriter::new();
    visual.do_it(false, false, false, &mut svg_writer);

    let raw = svg_writer.finalize();
    // Inject dark background
    let styled = raw.replacen("<svg ", "<svg style=\"background:#1a1a2e\" ", 1);
    Ok(styled)
}

// ─── Graph building ─────────────────────────────────────────

fn build_graph(source: &str) -> Result<(Vec<VizNode>, Vec<VizEdge>), String> {
    let tokens = lexer::lex(source)?;
    let ast = parser::parse(tokens)?;

    let mut nodes: Vec<VizNode> = Vec::new();
    let mut edges: Vec<VizEdge> = Vec::new();
    let mut name_to_node_id: HashMap<String, String> = HashMap::new();
    let mut node_counter = 0usize;

    // Collect coordinates, external ports, and defined names
    let mut input_names: Vec<String> = Vec::new();
    let mut externs: Vec<&ExternPort> = Vec::new();
    let mut defined_names: HashSet<String> = HashSet::new();
    let mut all_output_names: Vec<String> = Vec::new();

    for stmt in &ast.statements {
        match stmt {
            Statement::InputDecl(d) => input_names.push(d.name.clone()),
            Statement::Binding(b) => {
                for t in &b.targets {
                    defined_names.insert(t.clone());
                    all_output_names.push(t.clone());
                }
            }
            Statement::ExternPort(e) => externs.push(e),
            Statement::ModuleDef(_) => {}
            Statement::Cursor(_) => {}
            Statement::Pragma { .. } => {}
            Statement::For(_) => {}
            Statement::Tile(_) => {}
        }
    }

    // Infer coordinates
    if input_names.is_empty() {
        let mut refs: HashSet<String> = HashSet::new();
        for stmt in &ast.statements {
            let expr = match stmt {
                Statement::InputDecl(_)
                | Statement::ModuleDef(_)
                | Statement::ExternPort(_)
                | Statement::Cursor(_)
                | Statement::Pragma { .. }
                | Statement::For(_)
                | Statement::Tile(_) => continue,
                Statement::Binding(b) => &b.value,
            };
            collect_expr_idents(expr, &mut refs);
        }
        let extern_names: HashSet<&str> = externs.iter().map(|e| e.name.as_str()).collect();
        for name in refs {
            if !defined_names.contains(&name) && !extern_names.contains(name.as_str()) {
                input_names.push(name);
            }
        }
        input_names.sort();
    }

    // Determine which outputs are terminal (not consumed by other nodes)
    let mut consumed: HashSet<String> = HashSet::new();
    for stmt in &ast.statements {
        let expr = match stmt {
            Statement::InputDecl(_)
            | Statement::ModuleDef(_)
            | Statement::ExternPort(_)
            | Statement::Cursor(_)
            | Statement::Pragma { .. }
            | Statement::For(_)
            | Statement::Tile(_) => continue,
            Statement::Binding(b) => &b.value,
        };
        collect_expr_idents(expr, &mut consumed);
    }
    let terminal_outputs: Vec<String> = all_output_names
        .iter()
        .filter(|name| !consumed.contains(*name))
        .cloned()
        .collect();

    // ─── INPUTS register (top) ──────────────────────────
    // Single record node with all coordinates as output ports
    let inputs_id = "inputs".to_string();
    {
        let mut input_ports: Vec<String> = Vec::new();
        for name in &input_names {
            input_ports.push(name.clone());
        }
        nodes.push(VizNode {
            id: inputs_id.clone(),
            label: "INPUTS".into(),
            inputs: vec![],
            outputs: input_ports,
            kind: VizKind::Register,
        });
        for name in &input_names {
            name_to_node_id.insert(name.clone(), inputs_id.clone());
        }
    }

    // ─── External ports (top) ───────────────────────────
    // One register per `extern name: type [= default]`. These are the
    // kernel's `InputKind::ExternalWrite` slots: written by the host
    // between cycles and read by nodes like any other input.
    // The label carries kind, name, type, and default so a reader can
    // tell a port with a declared default from one that is unset until
    // the host writes it.
    for (idx, port) in externs.iter().enumerate() {
        let id = format!("x{idx}");
        let label = match &port.default {
            Some(default) => format!(
                "extern {}: {} = {}",
                port.name,
                port.typ,
                format_expr_short(default)
            ),
            None => format!("extern {}: {} (unset)", port.name, port.typ),
        };
        name_to_node_id.insert(port.name.clone(), id.clone());
        nodes.push(VizNode {
            id,
            label,
            inputs: vec![],
            outputs: vec![port.name.clone()],
            kind: VizKind::Extern,
        });
    }

    // ─── Function nodes (middle) ────────────────────────
    for stmt in &ast.statements {
        match stmt {
            Statement::InputDecl(_)
            | Statement::ModuleDef(_)
            | Statement::ExternPort(_)
            | Statement::Cursor(_)
            | Statement::Pragma { .. }
            | Statement::For(_)
            | Statement::Tile(_) => continue,
            Statement::Binding(b) => {
                let id = format!("n{node_counter}");
                node_counter += 1;

                let target_label = if b.targets.len() == 1 {
                    b.targets[0].clone()
                } else {
                    format!("({})", b.targets.join(", "))
                };

                let mut input_refs: Vec<String> = Vec::new();
                collect_expr_idents_ordered(&b.value, &mut input_refs);

                let label = format_node_label(&b.value, &target_label);

                for ref_name in &input_refs {
                    if let Some(src_id) = name_to_node_id.get(ref_name) {
                        edges.push(VizEdge {
                            from_node: src_id.clone(),
                            from_port: ref_name.clone(),
                            to_node: id.clone(),
                            to_port: ref_name.clone(),
                        });
                    }
                }

                for t in &b.targets {
                    name_to_node_id.insert(t.clone(), id.clone());
                }
                nodes.push(VizNode {
                    id,
                    label,
                    inputs: input_refs,
                    outputs: b.targets.clone(),
                    kind: VizKind::Func,
                });
            }
        }
    }

    // ─── OUTPUTS register (bottom) ──────────────────────
    // Single record node with all terminal outputs as input ports
    if !terminal_outputs.is_empty() {
        let outputs_id = "outputs".to_string();
        for name in &terminal_outputs {
            if let Some(src_id) = name_to_node_id.get(name) {
                edges.push(VizEdge {
                    from_node: src_id.clone(),
                    from_port: name.clone(),
                    to_node: outputs_id.clone(),
                    to_port: name.clone(),
                });
            }
        }
        nodes.push(VizNode {
            id: outputs_id,
            label: "OUTPUTS".into(),
            inputs: terminal_outputs,
            outputs: vec![],
            kind: VizKind::Register,
        });
    }

    Ok((nodes, edges))
}

fn format_node_label(expr: &Expr, target: &str) -> String {
    match expr {
        Expr::Call(call) => {
            let args: Vec<String> = call
                .args
                .iter()
                .map(|a| match a {
                    Arg::Positional(e) => format_expr_short(e),
                    Arg::Named(n, e) => format!("{}: {}", n, format_expr_short(e)),
                })
                .collect();
            format!("{} := {}({})", target, call.func, args.join(", "))
        }
        Expr::Ident(id, _) => format!("{} := {}", target, id),
        Expr::IntLit(v, _) => format!("{} = {}", target, v),
        Expr::FloatLit(v, _) => format!("{} = {}", target, v),
        Expr::StringLit(s, _) => {
            let trunc = if s.len() > 20 {
                format!("{}...", &s[..20])
            } else {
                s.clone()
            };
            format!("{} = \"{}\"", target, trunc)
        }
        _ => target.to_string(),
    }
}

fn format_expr_short(expr: &Expr) -> String {
    match expr {
        Expr::Ident(id, _) => id.clone(),
        Expr::IntLit(v, _) => v.to_string(),
        Expr::FloatLit(v, _) => format!("{v}"),
        Expr::StringLit(s, _) => format!("\"{s}\""),
        Expr::Call(call) => format!("{}(..)", call.func),
        _ => "..".into(),
    }
}

fn collect_expr_idents(expr: &Expr, out: &mut HashSet<String>) {
    match expr {
        Expr::Ident(name, _) => {
            out.insert(name.clone());
        }
        Expr::Call(call) => {
            for arg in &call.args {
                let inner = match arg {
                    Arg::Positional(e) | Arg::Named(_, e) => e,
                };
                collect_expr_idents(inner, out);
            }
        }
        Expr::ArrayLit(elems, _) => {
            for e in elems {
                collect_expr_idents(e, out);
            }
        }
        _ => {}
    }
}

/// Like collect_expr_idents but preserves order and avoids duplicates.
fn collect_expr_idents_ordered(expr: &Expr, out: &mut Vec<String>) {
    match expr {
        Expr::Ident(name, _) => {
            if !out.contains(name) {
                out.push(name.clone());
            }
        }
        Expr::Call(call) => {
            for arg in &call.args {
                let inner = match arg {
                    Arg::Positional(e) | Arg::Named(_, e) => e,
                };
                collect_expr_idents_ordered(inner, out);
            }
        }
        Expr::ArrayLit(elems, _) => {
            for e in elems {
                collect_expr_idents_ordered(e, out);
            }
        }
        _ => {}
    }
}

fn dot_escape(s: &str) -> String {
    s.replace('\\', "\\\\")
        .replace('"', "\\\"")
        .replace('{', "\\{")
        .replace('}', "\\}")
        .replace('<', "\\<")
        .replace('>', "\\>")
        .replace('|', "\\|")
}

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

    const SIMPLE_POLYDAT: &str = "input cycle: u64\nh := hash(cycle)\nuser_id := mod(h, 1000000)";

    #[test]
    fn dot_has_ports() {
        let dot = polydat_to_dot(SIMPLE_POLYDAT).unwrap();
        assert!(dot.contains("shape=record"));
        assert!(dot.contains("bgcolor"));
        assert!(dot.contains(":o_")); // output port syntax
        assert!(dot.contains(":i_")); // input port syntax
    }

    #[test]
    fn dot_dark_theme() {
        let dot = polydat_to_dot(SIMPLE_POLYDAT).unwrap();
        assert!(dot.contains("#1a1a2e")); // dark bg
        assert!(dot.contains("#16213e")); // node fill
        assert!(dot.contains("#e0e0e0")); // light text
    }

    #[test]
    fn mermaid_output() {
        let mermaid = polydat_to_mermaid(SIMPLE_POLYDAT).unwrap();
        assert!(mermaid.contains("flowchart TD"));
        assert!(mermaid.contains("-->"));
    }

    #[test]
    fn svg_dark_background() {
        let svg = polydat_to_svg(SIMPLE_POLYDAT).unwrap();
        assert!(svg.contains("<svg"));
        assert!(svg.contains("#1a1a2e"));
    }

    #[test]
    fn inferred_coords() {
        let src = "h := hash(cycle)\nid := mod(h, 100)";
        let dot = polydat_to_dot(src).unwrap();
        assert!(dot.contains("cycle"));
    }

    #[test]
    fn multi_output() {
        let src = "input cycle: u64\n(x, y) := mixed_radix(cycle, 100, 0)\nhx := hash(x)";
        let dot = polydat_to_dot(src).unwrap();
        assert!(dot.contains("mixed_radix"));
        assert!(dot.contains("hash"));
    }

    /// One coordinate, one extern with a default, one extern without.
    const EXTERN_POLYDAT: &str = "input cycle: u64\n\
        extern balance: f64 = 0.5\n\
        extern session_id: u64\n\
        h := hash(cycle)\n\
        scaled := f64_mul(balance, 2.0)\n\
        token := u64_add(h, session_id)";

    #[test]
    fn extern_ports_are_drawn_with_edges() {
        let dot = polydat_to_dot(EXTERN_POLYDAT).unwrap();
        // The coordinate register is unchanged and holds only coordinates.
        assert!(dot.contains("inputs [label=\"{ INPUTS | { <o_cycle> cycle } }\""));
        // Each extern is its own port node, labeled kind / name / type / default.
        assert!(
            dot.contains("x0 [label=\"{ extern balance: f64 = 0.5 | { <o_balance> balance } }\"")
        );
        assert!(dot.contains(
            "x1 [label=\"{ extern session_id: u64 (unset) | { <o_session_id> session_id } }\""
        ));
        // Externs carry their own accent, distinct from coordinates and outputs.
        assert!(dot.contains("#ffb454"));
        // Edges run from each input port to the node that reads it.
        assert!(dot.contains("inputs:o_cycle -> n0:i_cycle;"));
        assert!(dot.contains("x0:o_balance -> n1:i_balance;"));
        assert!(dot.contains("x1:o_session_id -> n2:i_session_id;"));
        assert!(dot.contains("n0:o_h -> n2:i_h;"));
    }

    #[test]
    fn extern_ports_are_not_inferred_as_coordinates() {
        // No `input` declaration: `cycle` is inferred, `k` is an extern.
        let src = "extern k: u64 = 7\nh := hash(cycle)\nz := u64_add(h, k)";
        let dot = polydat_to_dot(src).unwrap();
        assert!(dot.contains("inputs [label=\"{ INPUTS | { <o_cycle> cycle } }\""));
        assert!(dot.contains("x0 [label=\"{ extern k: u64 = 7 | { <o_k> k } }\""));
        assert!(dot.contains("x0:o_k -> n1:i_k;"));
    }

    #[test]
    fn extern_ports_in_mermaid_and_svg() {
        let mermaid = polydat_to_mermaid(EXTERN_POLYDAT).unwrap();
        assert!(mermaid.contains("x0([\"extern balance: f64 = 0.5\"])"));
        assert!(mermaid.contains("x1([\"extern session_id: u64 (unset)\"])"));
        assert!(mermaid.contains("x0 -->|balance| n1"));
        assert!(mermaid.contains("classDef extern"));
        assert!(mermaid.contains("class x0 extern"));

        let svg = polydat_to_svg(EXTERN_POLYDAT).unwrap();
        assert!(svg.contains("<svg"));
        assert!(svg.contains("balance"));
        assert!(svg.contains("session_id"));
    }

    #[test]
    fn programs_without_externs_draw_no_extern_nodes() {
        let dot = polydat_to_dot(SIMPLE_POLYDAT).unwrap();
        assert!(!dot.contains("extern"));
        assert!(!dot.contains("#ffb454"));
        let mermaid = polydat_to_mermaid(SIMPLE_POLYDAT).unwrap();
        assert!(!mermaid.contains("extern"));
    }
}