Mermaid.hs

{-# LANGUAGE OverloadedStrings #-}

module Calligraphy.Phases.Render.Mermaid
  ( renderMermaid,
  )
where

import Calligraphy.Phases.Render.Common
import Calligraphy.Prelude hiding (Decl, DeclType, Node)
import Calligraphy.Util.Printer
import Calligraphy.Util.Types
import Control.Monad.State (State, execState, modify)
import Data.List (intercalate)
import Data.List.NonEmpty (NonEmpty)
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Tree (Tree (..))

renderMermaid :: Prints RenderGraph
renderMermaid (RenderGraph roots calls types) = do
  textLn "flowchart TD"
  indent $ do
    case roots of
      Left modules -> mapM_ printModule modules
      Right trees -> mapM_ printTree trees
    forM_ (removeZeroEdges calls) $ \(caller, callee) ->
      strLn $ caller <> " --> " <> callee
    forM_ (removeZeroEdges types) $ \(caller, callee) ->
      strLn $ caller <> " -.-> " <> callee
    strLn "classDef default fill-opacity:0,stroke:#777;"
  where
    printTree :: Prints (Tree RenderNode)
    printTree (Node (RenderNode nodeid typ lbll export) []) = do
      strLn $ nodeid <> nodeShape typ (intercalate "\\n" lbll)
      unless export . strLn $
        "style " <> nodeid <> " stroke-dasharray: 5 5"
    printTree (Node (RenderNode nodeid _typ lbll export) children) = do
      brack ("subgraph " <> nodeid <> "[" <> intercalate "\\n" lbll <> "]") "end" $ do
        unless export . strLn $
          "style " <> nodeid <> " stroke-dasharray: 5 5"
        forM_ children printTree
    -- strLn $ nodeid <> " ~~~ " <> nodeId childNode

    printModule :: Prints RenderModule
    printModule (RenderModule lbl modId decls) =
      brack ("subgraph " <> modId <> " [" <> lbl <> "]") "end" $
        forM_ decls printTree

    -- Because we render hierarchies using subgraphs, there's an odd edge case
    -- when theres an edge between a parent and child; mermaid renders these as
    -- a zero-length edge, i.e. _just_ an arrowhead. So, we remove them.
    removeZeroEdges :: Set (ID, ID) -> Set (ID, ID)
    removeZeroEdges = execState (mapM_ go declRoots)
      where
        declRoots :: NonEmpty (Tree RenderNode)
        declRoots = case roots of
          Left mods -> mods >>= moduleDecls
          Right a -> a
        go :: Tree RenderNode -> State (Set (ID, ID)) ()
        go (Node (RenderNode parentId _ _ _) children) = do
          forM_ children $ \child@(Node (RenderNode childId _ _ _) _) -> do
            modify $ Set.delete (parentId, childId)
            go child

nodeShape :: DeclType -> String -> String
nodeShape DataDecl = wrapBracket "([" "])"
nodeShape ValueDecl = wrapBracket "[" "]"
nodeShape RecDecl = wrapBracket "(" ")"
nodeShape ConDecl = wrapBracket "(" ")"
nodeShape ClassDecl = wrapBracket "[/" "\\]"

wrapBracket :: String -> String -> String -> String
wrapBracket p q inner = p <> inner <> q