{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}

module Network.TypedProtocol.Documentation.GraphViz
where

import Data.Maybe
import qualified Data.GraphViz as Dot
import qualified Data.GraphViz.Attributes.Complete as Dot
import Data.Graph.Inductive.Graph (mkGraph, Node)
import Data.Graph.Inductive.PatriciaTree (Gr)
import qualified Data.Text.Lazy as LText
import qualified Data.Map.Strict as Map
import System.IO (Handle)
import qualified Data.ByteString as BS

import Network.TypedProtocol.Documentation.Types

protocolToDot :: ProtocolDescription codec -> Dot.DotGraph Node
protocolToDot :: forall codec. ProtocolDescription codec -> DotGraph Node
protocolToDot ProtocolDescription codec
proto =
  GraphvizParams
  Node (StateRef, AgencyID) String () (StateRef, AgencyID)
-> Gr (StateRef, AgencyID) String -> DotGraph Node
forall cl (gr :: * -> * -> *) nl el l.
(Ord cl, Graph gr) =>
GraphvizParams Node nl el cl l -> gr nl el -> DotGraph Node
Dot.graphToDot
          GraphvizParams
  Node (StateRef, AgencyID) Any () (StateRef, AgencyID)
forall n nl el. GraphvizParams n nl el () nl
Dot.nonClusteredParams
            { Dot.fmtNode = \ (Node
_, (StateRef
name, AgencyID
agency)) ->
                StateRef -> AgencyID -> Attributes
stateToNode StateRef
name AgencyID
agency
            , Dot.fmtEdge = \ (Node
_, Node
_, String
name) ->
                [ Label -> Attribute
Dot.Label (Text -> Label
Dot.StrLabel (Text -> Label) -> Text -> Label
forall a b. (a -> b) -> a -> b
$ String -> Text
LText.pack String
name)
                ]
            , Dot.globalAttributes =
                [ Dot.GraphAttrs
                    [ Dot.FontSize 10.0
                    , Dot.LabelFontSize 10.0
                    ]
                , Dot.NodeAttrs
                    [ Dot.FontName "Noto Sans"
                    , Dot.LabelFontName "Noto Sans"
                    ]
                , Dot.EdgeAttrs
                    [ Dot.FontName "Noto Sans"
                    , Dot.LabelFontName "Noto Sans"
                    , Dot.FontSize 8.0
                    , Dot.LabelFontSize 8.0
                    ]
                ]
            }
          ([(Node, (StateRef, AgencyID))]
-> [(Node, Node, String)] -> Gr (StateRef, AgencyID) String
forall a b. [LNode a] -> [LEdge b] -> Gr a b
forall (gr :: * -> * -> *) a b.
Graph gr =>
[LNode a] -> [LEdge b] -> gr a b
mkGraph [(Node, (StateRef, AgencyID))]
numberedStateInfos [(Node, Node, String)]
edges :: Gr (StateRef, AgencyID) String)
  where
    stateInfos :: [(StateRef, AgencyID)]
stateInfos = [ (StateRef
name, AgencyID
agency)  | (StateRef
name, [Description]
_, AgencyID
agency) <- ProtocolDescription codec -> [(StateRef, [Description], AgencyID)]
forall codec.
ProtocolDescription codec -> [(StateRef, [Description], AgencyID)]
protocolStates ProtocolDescription codec
proto ]
    stateNames :: [StateRef]
stateNames = ((StateRef, AgencyID) -> StateRef)
-> [(StateRef, AgencyID)] -> [StateRef]
forall a b. (a -> b) -> [a] -> [b]
map (StateRef, AgencyID) -> StateRef
forall a b. (a, b) -> a
fst [(StateRef, AgencyID)]
stateInfos
    numberedStateInfos :: [(Node, (StateRef, AgencyID))]
numberedStateInfos = [Node] -> [(StateRef, AgencyID)] -> [(Node, (StateRef, AgencyID))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Node
0::Int,Node
1..] [(StateRef, AgencyID)]
stateInfos
    numberedStateNames :: [(Node, StateRef)]
numberedStateNames = [Node] -> [StateRef] -> [(Node, StateRef)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Node
0::Int,Node
1..] [StateRef]
stateNames
    stateDict :: Map StateRef Node
stateDict =
      [(StateRef, Node)] -> Map StateRef Node
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(StateRef, Node)] -> Map StateRef Node)
-> [(StateRef, Node)] -> Map StateRef Node
forall a b. (a -> b) -> a -> b
$ ((Node, StateRef) -> (StateRef, Node))
-> [(Node, StateRef)] -> [(StateRef, Node)]
forall a b. (a -> b) -> [a] -> [b]
map (Node, StateRef) -> (StateRef, Node)
forall a b. (a, b) -> (b, a)
flipPair [(Node, StateRef)]
numberedStateNames

    agencyColor :: AgencyID -> ColorList
agencyColor AgencyID
NobodyAgencyID = [Color] -> ColorList
Dot.toColorList [X11Color -> Color
Dot.X11Color X11Color
Dot.Black]
    agencyColor AgencyID
ServerAgencyID = [Color] -> ColorList
Dot.toColorList [X11Color -> Color
Dot.X11Color X11Color
Dot.Blue]
    agencyColor AgencyID
ClientAgencyID = [Color] -> ColorList
Dot.toColorList [X11Color -> Color
Dot.X11Color X11Color
Dot.Brown]

    edges :: [(Node, Node, String)]
edges = [Maybe (Node, Node, String)] -> [(Node, Node, String)]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (Node, Node, String)] -> [(Node, Node, String)])
-> [Maybe (Node, Node, String)] -> [(Node, Node, String)]
forall a b. (a -> b) -> a -> b
$ ((MessageDescription codec -> Maybe (Node, Node, String))
 -> [MessageDescription codec] -> [Maybe (Node, Node, String)])
-> [MessageDescription codec]
-> (MessageDescription codec -> Maybe (Node, Node, String))
-> [Maybe (Node, Node, String)]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (MessageDescription codec -> Maybe (Node, Node, String))
-> [MessageDescription codec] -> [Maybe (Node, Node, String)]
forall a b. (a -> b) -> [a] -> [b]
map (ProtocolDescription codec -> [MessageDescription codec]
forall codec.
ProtocolDescription codec -> [MessageDescription codec]
protocolMessages ProtocolDescription codec
proto) ((MessageDescription codec -> Maybe (Node, Node, String))
 -> [Maybe (Node, Node, String)])
-> (MessageDescription codec -> Maybe (Node, Node, String))
-> [Maybe (Node, Node, String)]
forall a b. (a -> b) -> a -> b
$ \MessageDescription codec
msg -> do
              fromIndex <- StateRef -> Map StateRef Node -> Maybe Node
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (MessageDescription codec -> StateRef
forall codec. MessageDescription codec -> StateRef
messageFromState MessageDescription codec
msg) Map StateRef Node
stateDict
              toIndex <- Map.lookup (messageToState msg) stateDict
              return (fromIndex, toIndex, messageName msg)

    stateToNode :: StateRef -> AgencyID -> Attributes
stateToNode StateRef
AnyState AgencyID
_ =
      [ Label -> Attribute
Dot.Label (Text -> Label
Dot.StrLabel Text
"any state")
      , ColorList -> Attribute
Dot.Color ([Color] -> ColorList
Dot.toColorList [X11Color -> Color
Dot.X11Color X11Color
Dot.Gray])
      ]
    stateToNode (State String
name) AgencyID
agency =
      [ Label -> Attribute
Dot.Label (Text -> Label
Dot.StrLabel (Text -> Label) -> Text -> Label
forall a b. (a -> b) -> a -> b
$ String -> Text
LText.pack String
name)
      , ColorList -> Attribute
Dot.Color (AgencyID -> ColorList
agencyColor AgencyID
agency)
      ]

flipPair :: (a, b) -> (b, a)
flipPair :: forall a b. (a, b) -> (b, a)
flipPair (a
x, b
y) = (b
y, a
x)

protocolToSVGFile :: ProtocolDescription codec -> FilePath -> IO FilePath
protocolToSVGFile :: forall codec. ProtocolDescription codec -> String -> IO String
protocolToSVGFile ProtocolDescription codec
proto =
  GraphvizCommand
-> DotGraph Node -> GraphvizOutput -> String -> IO String
forall (dg :: * -> *) n.
PrintDotRepr dg n =>
GraphvizCommand -> dg n -> GraphvizOutput -> String -> IO String
Dot.runGraphvizCommand GraphvizCommand
Dot.Dot (ProtocolDescription codec -> DotGraph Node
forall codec. ProtocolDescription codec -> DotGraph Node
protocolToDot ProtocolDescription codec
proto) GraphvizOutput
Dot.Svg

protocolToDotFile :: ProtocolDescription codec -> FilePath -> IO FilePath
protocolToDotFile :: forall codec. ProtocolDescription codec -> String -> IO String
protocolToDotFile ProtocolDescription codec
proto =
  GraphvizCommand
-> DotGraph Node -> GraphvizOutput -> String -> IO String
forall (dg :: * -> *) n.
PrintDotRepr dg n =>
GraphvizCommand -> dg n -> GraphvizOutput -> String -> IO String
Dot.runGraphvizCommand GraphvizCommand
Dot.Dot (ProtocolDescription codec -> DotGraph Node
forall codec. ProtocolDescription codec -> DotGraph Node
protocolToDot ProtocolDescription codec
proto) GraphvizOutput
Dot.Canon

hProtocolToSVG :: ProtocolDescription codec -> Handle -> IO ()
hProtocolToSVG :: forall codec. ProtocolDescription codec -> Handle -> IO ()
hProtocolToSVG ProtocolDescription codec
proto Handle
dst =
  GraphvizCommand
-> DotGraph Node -> GraphvizOutput -> (Handle -> IO ()) -> IO ()
forall (dg :: * -> *) n a.
PrintDotRepr dg n =>
GraphvizCommand
-> dg n -> GraphvizOutput -> (Handle -> IO a) -> IO a
Dot.graphvizWithHandle GraphvizCommand
Dot.Dot (ProtocolDescription codec -> DotGraph Node
forall codec. ProtocolDescription codec -> DotGraph Node
protocolToDot ProtocolDescription codec
proto) GraphvizOutput
Dot.Svg Handle -> IO ()
consume
  where
    consume :: Handle -> IO ()
consume Handle
src = Handle -> IO ByteString
BS.hGetContents Handle
src IO ByteString -> (ByteString -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Handle -> ByteString -> IO ()
BS.hPut Handle
dst