byre/mac/Sources/Byre/Diagram.swift

import AppKit
import CoreText
import SwiftUI

enum NodeShape {
    case box, rounded, diamond, circle
}

enum EdgeStyle {
    case solid, dotted, thick
}

enum FlowDirection {
    case lr, rl, td, bt
}

struct NodeColors {
    var fill: Color?
    var stroke: Color?
    var text: Color?

    mutating func merge(_ other: NodeColors) {
        fill = other.fill ?? fill
        stroke = other.stroke ?? stroke
        text = other.text ?? text
    }
}

struct FlowNode {
    let id: String
    var label: String
    var shape: NodeShape = .box
    var classes: [String] = []
    var style = NodeColors()
}

struct FlowEdge {
    let from: String
    let to: String
    let label: String?
    let style: EdgeStyle
    let arrow: Bool
}

struct Flowchart {
    var direction: FlowDirection = .td
    var order: [String] = []
    var nodes: [String: FlowNode] = [:]
    var edges: [FlowEdge] = []
    var classDefs: [String: NodeColors] = [:]

    mutating func node(_ id: String, label: String? = nil, shape: NodeShape? = nil) {
        if nodes[id] == nil {
            order.append(id)
            nodes[id] = FlowNode(id: id, label: id)
        }
        if let label {
            nodes[id]!.label = label
        }
        if let shape {
            nodes[id]!.shape = shape
        }
    }

    func colors(_ id: String) -> NodeColors {
        let n = nodes[id]!
        var c = NodeColors()
        for cls in n.classes {
            if let d = classDefs[cls] {
                c.merge(d)
            }
        }
        c.merge(n.style)
        return c
    }
}

struct DiagramError: Error {
    let reason: String
}

private let header = Re(#"^(flowchart|graph)(?:\s+(LR|RL|TD|TB|BT))?$"#, caseless: true)
private let nodeId = Re(#"^[A-Za-z0-9_]+"#)
private let plainEdge = Re(#"^(-{2,}>|-{3,}|-\.+->|-\.+-|={2,}>|={3,})(?:\|([^|]*)\|)?"#)
private let dashLabel = Re(#"^--\s*([^->].*?)\s*(-{2,}>|-{3,})"#)
private let dotLabel = Re(#"^-\.\s*(.*?)\s*(\.+->|\.+-)"#)
private let thickLabel = Re(#"^==\s*([^=>].*?)\s*(={2,}>|={3,})"#)
private let classDefLine = Re(#"^classDef\s+([A-Za-z0-9_,]+)\s+(.+)$"#)
private let classLine = Re(#"^class\s+([A-Za-z0-9_,\s]+?)\s+([A-Za-z0-9_]+)$"#)
private let styleLine = Re(#"^style\s+([A-Za-z0-9_]+)\s+(.+)$"#)
private let hexColor = Re(#"^#([0-9a-fA-F]{3}|[0-9a-fA-F]{6})$"#)

private let shapes: [(open: String, close: String, shape: NodeShape)] = [
    ("((", "))", .circle), ("([", "])", .rounded), ("[[", "]]", .box), ("[(", ")]", .rounded), ("{{", "}}", .diamond),
    ("[", "]", .box), ("(", ")", .rounded), ("{", "}", .diamond),
]

func parseFlowchart(_ src: String) throws(DiagramError) -> Flowchart {
    var lines: [String] = []
    for raw in src.components(separatedBy: "\n") {
        let line = trimmed(raw)
        if line.isEmpty || line.hasPrefix("%%") {
            continue
        }
        for part in line.split(separator: ";") {
            let s = trimmed(String(part))
            if !s.isEmpty {
                lines.append(s)
            }
        }
    }
    guard let first = lines.first else {
        throw DiagramError(reason: "empty diagram")
    }
    guard let h = header.exec(first) else {
        let kind = first.split(separator: " ").first.map(String.init) ?? first
        throw DiagramError(reason: "\(kind) diagrams aren't drawn here, only flowcharts")
    }
    var chart = Flowchart()
    switch h[2]?.uppercased() {
    case "LR": chart.direction = .lr
    case "RL": chart.direction = .rl
    case "BT": chart.direction = .bt
    default: chart.direction = .td
    }
    for line in lines.dropFirst() {
        if line.hasPrefix("linkStyle ") {
            continue
        }
        if let m = classDefLine.exec(line) {
            let colors = parseStyle(m[2]!)
            for name in m[1]!.split(separator: ",") {
                chart.classDefs[String(name)] = colors
            }
            continue
        }
        if let m = classLine.exec(line) {
            for id in m[1]!.split(whereSeparator: { $0 == "," || $0 == " " }) {
                chart.node(String(id))
                chart.nodes[String(id)]!.classes.append(m[2]!)
            }
            continue
        }
        if let m = styleLine.exec(line) {
            chart.node(m[1]!)
            chart.nodes[m[1]!]!.style.merge(parseStyle(m[2]!))
            continue
        }
        try parseChain(line, &chart)
    }
    return chart
}

private func parseStyle(_ text: String) -> NodeColors {
    var c = NodeColors()
    for prop in text.split(separator: ",") {
        let kv = prop.split(separator: ":", maxSplits: 1).map { trimmed(String($0)) }
        guard kv.count == 2, let color = parseColor(kv[1]) else {
            continue
        }
        switch kv[0] {
        case "fill": c.fill = color
        case "stroke": c.stroke = color
        case "color": c.text = color
        default: break
        }
    }
    return c
}

private func parseColor(_ text: String) -> Color? {
    guard let m = hexColor.exec(text) else {
        return nil
    }
    var hex = m[1]!
    if hex.count == 3 {
        hex = hex.map { "\($0)\($0)" }.joined()
    }
    let v = UInt32(hex, radix: 16)!
    return Color(.sRGB, red: Double(v >> 16 & 0xff) / 255, green: Double(v >> 8 & 0xff) / 255, blue: Double(v & 0xff) / 255)
}

private func unquote(_ s: String) -> String {
    var t = trimmed(s)
    if t.count >= 2 && t.hasPrefix("\"") && t.hasSuffix("\"") {
        t = String(t.dropFirst().dropLast())
    }
    return t.replacingOccurrences(of: "<br/>", with: "\n").replacingOccurrences(of: "<br>", with: "\n").replacingOccurrences(of: "<br />", with: "\n")
}

private func parseChain(_ line: String, _ chart: inout Flowchart) throws(DiagramError) {
    var rest = Substring(line)
    let fail = DiagramError(reason: "can't parse \"\(line)\"")
    func skipSpaces() {
        rest = rest.drop { $0 == " " || $0 == "\t" }
    }
    func node() throws(DiagramError) -> String {
        skipSpaces()
        guard let m = nodeId.match(String(rest)) else {
            throw fail
        }
        let id = m.groups[0]!
        rest = rest.dropFirst(id.count)
        var label: String?
        var shape: NodeShape?
        for s in shapes where rest.hasPrefix(s.open) {
            let inner = rest.dropFirst(s.open.count)
            var after: Substring?
            if inner.first == "\"", let q = inner.dropFirst().firstIndex(of: "\"") {
                let tail = inner[inner.index(after: q)...]
                if tail.hasPrefix(s.close) {
                    label = unquote(String(inner[...q]))
                    after = tail.dropFirst(s.close.count)
                }
            }
            if after == nil, let r = inner.range(of: s.close) {
                label = unquote(String(inner[..<r.lowerBound]))
                after = inner[r.upperBound...]
            }
            guard let after else {
                throw fail
            }
            rest = after
            shape = s.shape
            break
        }
        chart.node(id, label: label, shape: shape)
        if rest.hasPrefix(":::") {
            rest = rest.dropFirst(3)
            guard let m = nodeId.match(String(rest)) else {
                throw fail
            }
            chart.nodes[id]!.classes.append(m.groups[0]!)
            rest = rest.dropFirst(m.groups[0]!.count)
        }
        return id
    }
    func edge() throws(DiagramError) -> (label: String?, style: EdgeStyle, arrow: Bool) {
        skipSpaces()
        let s = String(rest)
        if let m = plainEdge.match(s) {
            rest = rest.dropFirst(s.distance(from: s.startIndex, to: m.range.upperBound))
            let arrow = m.groups[1]!
            return (m.groups[2].map(unquote), arrow.hasPrefix("=") ? .thick : arrow.contains(".") ? .dotted : .solid, arrow.hasSuffix(">"))
        }
        for (re, style) in [(dashLabel, EdgeStyle.solid), (dotLabel, .dotted), (thickLabel, .thick)] {
            if let m = re.match(s) {
                rest = rest.dropFirst(s.distance(from: s.startIndex, to: m.range.upperBound))
                return (unquote(m.groups[1]!), style, m.groups[2]!.hasSuffix(">"))
            }
        }
        throw fail
    }
    var from = try node()
    while true {
        skipSpaces()
        if rest.isEmpty {
            return
        }
        let e = try edge()
        let to = try node()
        chart.edges.append(FlowEdge(from: from, to: to, label: e.label?.isEmpty == true ? nil : e.label, style: e.style, arrow: e.arrow))
        from = to
    }
}

struct PlacedNode {
    let node: FlowNode
    let center: CGPoint
    let size: CGSize
    let colors: NodeColors
}

struct DiagramLayout {
    let nodes: [String: PlacedNode]
    let edges: [FlowEdge]
    let size: CGSize
}

nonisolated(unsafe) private let labelFont = Fonts.ct(13)
nonisolated(unsafe) private let edgeFont = Fonts.ct(12)

func textSize(_ text: String, _ font: CTFont) -> CGSize {
    let lines = text.components(separatedBy: "\n")
    var width: CGFloat = 0
    for l in lines {
        let line = CTLineCreateWithAttributedString(NSAttributedString(string: l, attributes: [.font: font as NSFont]))
        width = max(width, CGFloat(CTLineGetTypographicBounds(line, nil, nil, nil)))
    }
    let lineHeight = CTFontGetAscent(font) + CTFontGetDescent(font) + CTFontGetLeading(font) + 3
    return CGSize(width: ceil(width), height: ceil(lineHeight * CGFloat(lines.count)))
}

func layoutFlowchart(_ chart: Flowchart) -> DiagramLayout {
    let ids = chart.order
    var succ: [String: [String]] = [:]
    for e in chart.edges where e.from != e.to {
        succ[e.from, default: []].append(e.to)
    }

    var state: [String: Int] = [:]
    var dag: [String: [String]] = [:]
    func visit(_ n: String) {
        state[n] = 1
        for m in succ[n] ?? [] {
            if state[m] == 1 {
                dag[m, default: []].append(n)
                continue
            }
            dag[n, default: []].append(m)
            if state[m] == nil {
                visit(m)
            }
        }
        state[n] = 2
    }
    for n in ids where state[n] == nil {
        visit(n)
    }

    var indegree: [String: Int] = [:]
    for (_, ms) in dag {
        for m in ms {
            indegree[m, default: 0] += 1
        }
    }
    var layer: [String: Int] = [:]
    var queue = ids.filter { (indegree[$0] ?? 0) == 0 }
    for n in queue {
        layer[n] = 0
    }
    var qi = 0
    while qi < queue.count {
        let n = queue[qi]
        qi += 1
        for m in dag[n] ?? [] {
            layer[m] = max(layer[m] ?? 0, layer[n]! + 1)
            indegree[m]! -= 1
            if indegree[m] == 0 {
                queue.append(m)
            }
        }
    }
    let depth = (layer.values.max() ?? 0) + 1
    var layers = Array(repeating: [String](), count: depth)
    for n in ids {
        layers[layer[n] ?? 0].append(n)
    }

    var preds: [String: [String]] = [:]
    for (n, ms) in dag {
        for m in ms {
            preds[m, default: []].append(n)
        }
    }
    func position(_ l: Int) -> [String: Double] {
        Dictionary(uniqueKeysWithValues: layers[l].enumerated().map { ($0.element, Double($0.offset)) })
    }
    func reorder(_ l: Int, _ ref: Int, _ links: [String: [String]]) {
        let pos = position(ref)
        let current = position(l)
        layers[l].sort { a, b in
            func bary(_ n: String) -> Double {
                let ps = (links[n] ?? []).compactMap { pos[$0] }
                return ps.isEmpty ? current[n]! : ps.reduce(0, +) / Double(ps.count)
            }
            let (x, y) = (bary(a), bary(b))
            return x == y ? current[a]! < current[b]! : x < y
        }
    }
    for _ in 0..<4 {
        for l in layers.indices.dropFirst() {
            reorder(l, l - 1, preds)
        }
        for l in layers.indices.dropLast().reversed() {
            reorder(l, l + 1, dag)
        }
    }

    var sizes: [String: CGSize] = [:]
    for n in ids {
        let node = chart.nodes[n]!
        let t = textSize(node.label, labelFont)
        switch node.shape {
        case .box, .rounded:
            sizes[n] = CGSize(width: t.width + 30, height: t.height + 18)
        case .diamond:
            sizes[n] = CGSize(width: (t.width * 0.75 + 14) * 2, height: (t.height * 1.5 + 8) * 2)
        case .circle:
            let d = (t.width * t.width + t.height * t.height).squareRoot() + 16
            sizes[n] = CGSize(width: d, height: d)
        }
    }

    let horizontal = chart.direction == .lr || chart.direction == .rl
    let rankGap: CGFloat = horizontal ? 70 : 56
    let nodeGap: CGFloat = horizontal ? 24 : 32
    func along(_ s: CGSize) -> CGFloat {
        horizontal ? s.width : s.height
    }
    func across(_ s: CGSize) -> CGFloat {
        horizontal ? s.height : s.width
    }
    let spans = layers.map { l in l.map { across(sizes[$0]!) }.reduce(0, +) + nodeGap * CGFloat(max(0, l.count - 1)) }
    let widest = spans.max() ?? 0
    var centers: [String: CGPoint] = [:]
    var rank: CGFloat = 0
    for (l, nodes) in layers.enumerated() {
        let thick = nodes.map { along(sizes[$0]!) }.max() ?? 0
        var offset = (widest - spans[l]) / 2
        for n in nodes {
            let a = across(sizes[n]!)
            let c1 = rank + thick / 2
            let c2 = offset + a / 2
            centers[n] = horizontal ? CGPoint(x: c1, y: c2) : CGPoint(x: c2, y: c1)
            offset += a + nodeGap
        }
        rank += thick + rankGap
    }
    let length = max(0, rank - rankGap)
    let margin: CGFloat = 8
    let total = horizontal ? CGSize(width: length, height: widest) : CGSize(width: widest, height: length)
    var placed: [String: PlacedNode] = [:]
    for n in ids {
        var c = centers[n]!
        if chart.direction == .rl {
            c.x = total.width - c.x
        }
        if chart.direction == .bt {
            c.y = total.height - c.y
        }
        placed[n] = PlacedNode(node: chart.nodes[n]!, center: CGPoint(x: c.x + margin, y: c.y + margin), size: sizes[n]!, colors: chart.colors(n))
    }
    return DiagramLayout(nodes: placed, edges: chart.edges, size: CGSize(width: total.width + margin * 2, height: total.height + margin * 2))
}

private func boundary(_ n: PlacedNode, toward p: CGPoint) -> CGPoint {
    let dx = p.x - n.center.x
    let dy = p.y - n.center.y
    if dx == 0 && dy == 0 {
        return n.center
    }
    let a = n.size.width / 2
    let b = n.size.height / 2
    let t: CGFloat
    switch n.node.shape {
    case .box, .rounded:
        t = min(dx == 0 ? .infinity : a / abs(dx), dy == 0 ? .infinity : b / abs(dy))
    case .diamond:
        t = 1 / (abs(dx) / a + abs(dy) / b)
    case .circle:
        t = a / (dx * dx + dy * dy).squareRoot()
    }
    return CGPoint(x: n.center.x + dx * t, y: n.center.y + dy * t)
}

private func shapePath(_ n: PlacedNode) -> Path {
    let r = CGRect(x: n.center.x - n.size.width / 2, y: n.center.y - n.size.height / 2, width: n.size.width, height: n.size.height)
    switch n.node.shape {
    case .box:
        return Path(roundedRect: r, cornerRadius: 3)
    case .rounded:
        return Path(roundedRect: r, cornerRadius: min(r.height / 2, 14))
    case .circle:
        return Path(ellipseIn: r)
    case .diamond:
        var p = Path()
        p.move(to: CGPoint(x: r.midX, y: r.minY))
        p.addLine(to: CGPoint(x: r.maxX, y: r.midY))
        p.addLine(to: CGPoint(x: r.midX, y: r.maxY))
        p.addLine(to: CGPoint(x: r.minX, y: r.midY))
        p.closeSubpath()
        return p
    }
}

struct FlowchartCanvas: View {
    let layout: DiagramLayout

    var body: some View {
        FitWidth(natural: layout.size) {
            Canvas { context, size in
                let scale = layout.size.width > 0 ? size.width / layout.size.width : 1
                context.scaleBy(x: scale, y: scale)
                draw(&context)
            }
        }
    }

    private func draw(_ context: inout GraphicsContext) {
        for e in layout.edges {
            guard let a = layout.nodes[e.from], let b = layout.nodes[e.to] else {
                continue
            }
            let stroke = StrokeStyle(lineWidth: e.style == .thick ? 3 : 1.5, lineCap: .round, dash: e.style == .dotted ? [3, 4] : [])
            var mid: CGPoint
            if e.from == e.to {
                let top = CGPoint(x: a.center.x, y: a.center.y - a.size.height / 2)
                var loop = Path()
                loop.move(to: CGPoint(x: top.x - 8, y: top.y))
                loop.addCurve(to: CGPoint(x: top.x + 8, y: top.y), control1: CGPoint(x: top.x - 24, y: top.y - 30), control2: CGPoint(x: top.x + 24, y: top.y - 30))
                context.stroke(loop, with: .color(Palette.muted), style: stroke)
                if e.arrow {
                    arrowHead(&context, from: CGPoint(x: top.x + 20, y: top.y - 20), to: CGPoint(x: top.x + 8, y: top.y))
                }
                mid = CGPoint(x: top.x, y: top.y - 26)
            } else {
                let start = boundary(a, toward: b.center)
                let end = boundary(b, toward: a.center)
                var line = Path()
                line.move(to: start)
                line.addLine(to: end)
                context.stroke(line, with: .color(Palette.muted), style: stroke)
                if e.arrow {
                    arrowHead(&context, from: start, to: end)
                }
                mid = CGPoint(x: (start.x + end.x) / 2, y: (start.y + end.y) / 2)
            }
            if let label = e.label {
                let size = textSize(label, edgeFont)
                let box = CGRect(x: mid.x - size.width / 2 - 4, y: mid.y - size.height / 2 - 2, width: size.width + 8, height: size.height + 4)
                context.fill(Path(roundedRect: box, cornerRadius: 3), with: .color(Palette.bg))
                context.draw(Text(label).font(Font(edgeFont)).foregroundColor(Palette.text), at: mid, anchor: .center)
            }
        }
        for (_, n) in layout.nodes {
            let path = shapePath(n)
            context.fill(path, with: .color(n.colors.fill ?? Palette.surface))
            context.stroke(path, with: .color(n.colors.stroke ?? Palette.border), lineWidth: 1)
            for (i, line) in n.node.label.components(separatedBy: "\n").enumerated() {
                let count = CGFloat(n.node.label.components(separatedBy: "\n").count)
                let lineHeight = textSize("X", labelFont).height
                let y = n.center.y + (CGFloat(i) - (count - 1) / 2) * lineHeight
                context.draw(Text(line).font(Font(labelFont)).foregroundColor(n.colors.text ?? Palette.text), at: CGPoint(x: n.center.x, y: y), anchor: .center)
            }
        }
    }

    private func arrowHead(_ context: inout GraphicsContext, from: CGPoint, to: CGPoint) {
        let angle = atan2(to.y - from.y, to.x - from.x)
        let size: CGFloat = 9
        var head = Path()
        head.move(to: to)
        head.addLine(to: CGPoint(x: to.x - size * cos(angle - .pi / 7), y: to.y - size * sin(angle - .pi / 7)))
        head.addLine(to: CGPoint(x: to.x - size * cos(angle + .pi / 7), y: to.y - size * sin(angle + .pi / 7)))
        head.closeSubpath()
        context.fill(head, with: .color(Palette.muted))
    }
}

struct DiagramView: View {
    let source: String

    var body: some View {
        switch Result(catching: { () throws(DiagramError) -> Flowchart in try parseFlowchart(source) }) {
        case .success(let chart):
            FlowchartCanvas(layout: layoutFlowchart(chart))
                .frame(maxWidth: .infinity, alignment: .leading)
        case .failure(let error):
            VStack(alignment: .leading, spacing: 4) {
                Text(verbatim: source)
                    .textSelection(.enabled)
                    .padding(.horizontal, 8)
                    .padding(.vertical, 6)
                    .frame(maxWidth: .infinity, alignment: .leading)
                    .background(Palette.surface, in: RoundedRectangle(cornerRadius: 6))
                Text(verbatim: "diagram not rendered: \(error.reason)")
                    .foregroundStyle(Palette.muted)
            }
        }
    }
}