{-# LANGUAGE DeriveFunctor       #-}
{-# LANGUAGE DerivingStrategies  #-}
{-# LANGUAGE Rank2Types          #-}
{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Test.StateMachine.DotDrawing
  ( GraphOptions (..)
  , GraphvizOutput (..)
  , Rose (..)
  , printDotGraph
  ) where

import           Control.Exception
import           Control.Monad
import           Data.GraphViz.Attributes.Complete
import           Data.GraphViz.Commands
import           Data.GraphViz.Exception
import           Data.GraphViz.Types.Canonical
import           Data.List
                   (uncons)
import           Data.List.Split
import           Data.Map                          hiding
                   (null)
import           Data.Maybe
import           Data.Text.Lazy
                   (pack)
import           Prelude
import           Test.StateMachine.Types.History

------------------------------------------------------------------------

data GraphOptions = GraphOptions {
      GraphOptions -> String
filePath       :: FilePath       --  Where to store the graph
                                       --  (note: file extensions are not checked)
    , GraphOptions -> GraphvizOutput
graphvizOutput :: GraphvizOutput --  output formats (like Jpeg, Png ..)
    }

data Rose a = Rose a (Map Pid a)
    deriving stock ((forall a b. (a -> b) -> Rose a -> Rose b)
-> (forall a b. a -> Rose b -> Rose a) -> Functor Rose
forall a b. a -> Rose b -> Rose a
forall a b. (a -> b) -> Rose a -> Rose b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Rose a -> Rose b
fmap :: forall a b. (a -> b) -> Rose a -> Rose b
$c<$ :: forall a b. a -> Rose b -> Rose a
<$ :: forall a b. a -> Rose b -> Rose a
Functor, Int -> Rose a -> ShowS
[Rose a] -> ShowS
Rose a -> String
(Int -> Rose a -> ShowS)
-> (Rose a -> String) -> ([Rose a] -> ShowS) -> Show (Rose a)
forall a. Show a => Int -> Rose a -> ShowS
forall a. Show a => [Rose a] -> ShowS
forall a. Show a => Rose a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Rose a -> ShowS
showsPrec :: Int -> Rose a -> ShowS
$cshow :: forall a. Show a => Rose a -> String
show :: Rose a -> String
$cshowList :: forall a. Show a => [Rose a] -> ShowS
showList :: [Rose a] -> ShowS
Show)

printDotGraph :: GraphOptions -> Rose [String] -> IO ()
printDotGraph :: GraphOptions -> Rose [String] -> IO ()
printDotGraph GraphOptions{String
GraphvizOutput
filePath :: GraphOptions -> String
graphvizOutput :: GraphOptions -> GraphvizOutput
filePath :: String
graphvizOutput :: GraphvizOutput
..} (Rose [String]
pref Map Pid [String]
sfx) = do
    let

        -- create barrier nodes

        nThreads :: Int
nThreads = Map Pid [String] -> Int
forall k a. Map k a -> Int
size Map Pid [String]
sfx
        barrierRecord :: [RecordField]
barrierRecord = (\Int
n -> PortName -> RecordField
PortName (PN {portName :: Text
portName = String -> Text
pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
n})) (Int -> RecordField) -> [Int] -> [RecordField]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
1..Int
nThreads]
        barrierNode :: DotNode String
barrierNode = DotNode {
                        nodeID :: String
nodeID = String
"barrier"
                        , nodeAttributes :: Attributes
nodeAttributes =
                            [Shape -> Attribute
Shape Shape
Record,NodeSize -> Attribute
FixedSize NodeSize
SetNodeSize,Double -> Attribute
Width Double
4.0,
                                Double -> Attribute
Height Double
0.0,
                                Label -> Attribute
Label ([RecordField] -> Label
RecordLabel [RecordField]
barrierRecord)]
                        }

        -- create preffix

        prefixWithResp :: [(Int, (String, String))]
prefixWithResp = [Int] -> [(String, String)] -> [(Int, (String, String))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..] ([(String, String)] -> [(Int, (String, String))])
-> [(String, String)] -> [(Int, (String, String))]
forall a b. (a -> b) -> a -> b
$ String -> [String] -> [(String, String)]
forall a. String -> [a] -> [(a, a)]
byTwoUnsafe String
"prefix" [String]
pref
        prefixNodes :: [DotNode String]
prefixNodes = String -> (Int, (String, String)) -> DotNode String
toDotNode String
"prefix" ((Int, (String, String)) -> DotNode String)
-> [(Int, (String, String))] -> [DotNode String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, (String, String))]
prefixWithResp
        prefixEdges :: [DotEdge String]
prefixEdges = [DotNode String] -> [DotEdge String]
forall a. [DotNode a] -> [DotEdge a]
connectNodes [DotNode String]
prefixNodes

        -- create suffixes

        nodesAndEdges :: [(Int, [DotNode String], [DotEdge String])]
nodesAndEdges = (((Pid, [String]) -> (Int, [DotNode String], [DotEdge String]))
 -> [(Pid, [String])]
 -> [(Int, [DotNode String], [DotEdge String])])
-> [(Pid, [String])]
-> ((Pid, [String]) -> (Int, [DotNode String], [DotEdge String]))
-> [(Int, [DotNode String], [DotEdge String])]
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Pid, [String]) -> (Int, [DotNode String], [DotEdge String]))
-> [(Pid, [String])] -> [(Int, [DotNode String], [DotEdge String])]
forall a b. (a -> b) -> [a] -> [b]
Prelude.map (Map Pid [String] -> [(Pid, [String])]
forall k a. Map k a -> [(k, a)]
toList Map Pid [String]
sfx) (((Pid, [String]) -> (Int, [DotNode String], [DotEdge String]))
 -> [(Int, [DotNode String], [DotEdge String])])
-> ((Pid, [String]) -> (Int, [DotNode String], [DotEdge String]))
-> [(Int, [DotNode String], [DotEdge String])]
forall a b. (a -> b) -> a -> b
$ \(Pid
pid, [String]
str) ->
            let p :: Int
p = Pid -> Int
unPid Pid
pid
                s :: [(Int, (String, String))]
s = [Int] -> [(String, String)] -> [(Int, (String, String))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..] ([(String, String)] -> [(Int, (String, String))])
-> [(String, String)] -> [(Int, (String, String))]
forall a b. (a -> b) -> a -> b
$ String -> [String] -> [(String, String)]
forall a. String -> [a] -> [(a, a)]
byTwoUnsafe (Int -> String
forall a. Show a => a -> String
show Int
p) [String]
str
                n :: [DotNode String]
n = String -> (Int, (String, String)) -> DotNode String
toDotNode (Int -> String
forall a. Show a => a -> String
show Int
p) ((Int, (String, String)) -> DotNode String)
-> [(Int, (String, String))] -> [DotNode String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, (String, String))]
s
                e :: [DotEdge String]
e = [DotNode String] -> [DotEdge String]
forall a. [DotNode a] -> [DotEdge a]
connectNodes [DotNode String]
n
            in (Int
p, [DotNode String]
n, [DotEdge String]
e)
        nodes :: [DotNode String]
nodes = ((Int, [DotNode String], [DotEdge String]) -> [DotNode String])
-> [(Int, [DotNode String], [DotEdge String])] -> [DotNode String]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(Int
_,[DotNode String]
n,[DotEdge String]
_) -> [DotNode String]
n) [(Int, [DotNode String], [DotEdge String])]
nodesAndEdges
        edges :: [DotEdge String]
edges = ((Int, [DotNode String], [DotEdge String]) -> [DotEdge String])
-> [(Int, [DotNode String], [DotEdge String])] -> [DotEdge String]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(Int
_,[DotNode String]
_,[DotEdge String]
e) -> [DotEdge String]
e) [(Int, [DotNode String], [DotEdge String])]
nodesAndEdges
        firstOfEachPid :: [(Int, Maybe (DotNode String))]
firstOfEachPid = (\(Int
p, [DotNode String]
n, [DotEdge String]
_) -> (Int
p, (DotNode String, [DotNode String]) -> DotNode String
forall a b. (a, b) -> a
fst ((DotNode String, [DotNode String]) -> DotNode String)
-> Maybe (DotNode String, [DotNode String])
-> Maybe (DotNode String)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DotNode String] -> Maybe (DotNode String, [DotNode String])
forall a. [a] -> Maybe (a, [a])
uncons [DotNode String]
n)) ((Int, [DotNode String], [DotEdge String])
 -> (Int, Maybe (DotNode String)))
-> [(Int, [DotNode String], [DotEdge String])]
-> [(Int, Maybe (DotNode String))]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, [DotNode String], [DotEdge String])]
nodesAndEdges

        -- create barrier edges

        edgesFromBarrier :: [DotEdge String]
edgesFromBarrier = ((Int, Maybe (DotNode String)) -> [DotEdge String])
-> [(Int, Maybe (DotNode String))] -> [DotEdge String]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(Int
p, Maybe (DotNode String)
mn) -> case Maybe (DotNode String)
mn of
                            Maybe (DotNode String)
Nothing -> []
                            Just DotNode String
n -> [DotEdge {
                                      fromNode :: String
fromNode = DotNode String -> String
forall n. DotNode n -> n
nodeID DotNode String
barrierNode
                                    , toNode :: String
toNode = DotNode String -> String
forall n. DotNode n -> n
nodeID DotNode String
n
                                    , edgeAttributes :: Attributes
edgeAttributes = [PortPos -> Attribute
TailPort (PortName -> Maybe CompassPoint -> PortPos
LabelledPort (PN {portName :: Text
portName = String -> Text
pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
p}) Maybe CompassPoint
forall a. Maybe a
Nothing)]
                                    }]) [(Int, Maybe (DotNode String))]
firstOfEachPid
        prefixToBarrier :: [DotEdge String]
prefixToBarrier = case [DotNode String]
prefixNodes of
            [] -> []
            [DotNode String]
_  -> [DotEdge {
                        fromNode :: String
fromNode = DotNode String -> String
forall n. DotNode n -> n
nodeID ([DotNode String] -> DotNode String
forall a. HasCallStack => [a] -> a
last [DotNode String]
prefixNodes)
                    , toNode :: String
toNode = DotNode String -> String
forall n. DotNode n -> n
nodeID DotNode String
barrierNode
                    , edgeAttributes :: Attributes
edgeAttributes = [] -- [HeadPort (LabelledPort (PN {portName = "1"}) Nothing)]]
                    }]

        -- create dot graph

        dotStmts :: DotStatements String
dotStmts = DotStmts {
                    attrStmts :: [GlobalAttributes]
attrStmts = [NodeAttrs {attrs :: Attributes
attrs = [Shape -> Attribute
Shape Shape
BoxShape,Double -> Attribute
Width Double
4.0]}]
                    , subGraphs :: [DotSubGraph String]
subGraphs = [] -- do we want to put commands with same pid on the same group?
                    , nodeStmts :: [DotNode String]
nodeStmts = DotNode String
barrierNode DotNode String -> [DotNode String] -> [DotNode String]
forall a. a -> [a] -> [a]
: ([DotNode String]
prefixNodes [DotNode String] -> [DotNode String] -> [DotNode String]
forall a. [a] -> [a] -> [a]
++ [DotNode String]
nodes)
                    , edgeStmts :: [DotEdge String]
edgeStmts = [DotEdge String]
prefixToBarrier [DotEdge String] -> [DotEdge String] -> [DotEdge String]
forall a. [a] -> [a] -> [a]
++ [DotEdge String]
prefixEdges [DotEdge String] -> [DotEdge String] -> [DotEdge String]
forall a. [a] -> [a] -> [a]
++ [DotEdge String]
edges [DotEdge String] -> [DotEdge String] -> [DotEdge String]
forall a. [a] -> [a] -> [a]
++ [DotEdge String]
edgesFromBarrier
                    }

        dg :: DotGraph String
dg = DotGraph {
            strictGraph :: Bool
strictGraph = Bool
False
            , directedGraph :: Bool
directedGraph = Bool
True
            , graphID :: Maybe GraphID
graphID = GraphID -> Maybe GraphID
forall a. a -> Maybe a
Just (Text -> GraphID
Str (Text -> GraphID) -> Text -> GraphID
forall a b. (a -> b) -> a -> b
$ String -> Text
pack String
"G")
            , graphStatements :: DotStatements String
graphStatements = DotStatements String
dotStmts
        }

    Either GraphvizException (Either IOException String)
err <- IO (Either IOException String)
-> IO (Either GraphvizException (Either IOException String))
forall e a. Exception e => IO a -> IO (Either e a)
try (IO (Either IOException String)
 -> IO (Either GraphvizException (Either IOException String)))
-> IO (Either IOException String)
-> IO (Either GraphvizException (Either IOException String))
forall a b. (a -> b) -> a -> b
$ IO String -> IO (Either IOException String)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO String -> IO (Either IOException String))
-> IO String -> IO (Either IOException String)
forall a b. (a -> b) -> a -> b
$ DotGraph String -> GraphvizOutput -> String -> IO String
forall (dg :: * -> *) n.
PrintDotRepr dg n =>
dg n -> GraphvizOutput -> String -> IO String
runGraphviz DotGraph String
dg GraphvizOutput
graphvizOutput String
filePath
    case Either GraphvizException (Either IOException String)
err of
        Left (GraphvizException
e :: GraphvizException) ->
            String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ GraphvizException -> String
forall e. Exception e => e -> String
displayException GraphvizException
e
        Right (Left (IOException
e :: IOException)) ->
            String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ IOException -> String
forall e. Exception e => e -> String
displayException IOException
e
        Right (Right String
_) ->
            () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

toDotNode :: String -> (Int, (String,String)) -> DotNode String
toDotNode :: String -> (Int, (String, String)) -> DotNode String
toDotNode String
nodeIdGroup (Int
n, (String
invocation, String
resp)) =
    DotNode {
          nodeID :: String
nodeID = String
nodeIdGroup String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"-" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n
        , nodeAttributes :: Attributes
nodeAttributes = [Label -> Attribute
Label (Label -> Attribute) -> Label -> Attribute
forall a b. (a -> b) -> a -> b
$ Text -> Label
StrLabel (Text -> Label) -> Text -> Label
forall a b. (a -> b) -> a -> b
$ String -> Text
pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$
                String -> Int -> ShowS
newLinesAfter String
"\\l" Int
60 String
invocation
                String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"\\n"
                String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> Int -> ShowS
newLinesAfter String
"\\r" Int
60 String
resp]
    }

byTwoUnsafe :: String -> [a] -> [(a,a)]
byTwoUnsafe :: forall a. String -> [a] -> [(a, a)]
byTwoUnsafe String
str [a]
ls = [(a, a)] -> Maybe [(a, a)] -> [(a, a)]
forall a. a -> Maybe a -> a
fromMaybe (String -> [(a, a)]
forall a. HasCallStack => String -> a
error (String -> [(a, a)]) -> String -> [(a, a)]
forall a b. (a -> b) -> a -> b
$ String
"couldn't split " String -> ShowS
forall a. [a] -> [a] -> [a]
++ if String -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
str then String
" " else String
str String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" in pairs") (Maybe [(a, a)] -> [(a, a)]) -> Maybe [(a, a)] -> [(a, a)]
forall a b. (a -> b) -> a -> b
$ [a] -> Maybe [(a, a)]
forall a. [a] -> Maybe [(a, a)]
byTwo [a]
ls

byTwo :: [a] -> Maybe [(a,a)]
byTwo :: forall a. [a] -> Maybe [(a, a)]
byTwo = [(a, a)] -> [a] -> Maybe [(a, a)]
forall {b}. [(b, b)] -> [b] -> Maybe [(b, b)]
go []
  where
    go :: [(b, b)] -> [b] -> Maybe [(b, b)]
go [(b, b)]
acc []            = [(b, b)] -> Maybe [(b, b)]
forall a. a -> Maybe a
Just ([(b, b)] -> Maybe [(b, b)]) -> [(b, b)] -> Maybe [(b, b)]
forall a b. (a -> b) -> a -> b
$ [(b, b)] -> [(b, b)]
forall a. [a] -> [a]
reverse [(b, b)]
acc
    go [(b, b)]
_acc [b
_]          = Maybe [(b, b)]
forall a. Maybe a
Nothing
    go [(b, b)]
acc (b
a: b
b : [b]
rest) = [(b, b)] -> [b] -> Maybe [(b, b)]
go ((b
a,b
b) (b, b) -> [(b, b)] -> [(b, b)]
forall a. a -> [a] -> [a]
: [(b, b)]
acc) [b]
rest

connectNodes :: [DotNode a] -> [DotEdge a]
connectNodes :: forall a. [DotNode a] -> [DotEdge a]
connectNodes = [DotEdge a] -> [DotNode a] -> [DotEdge a]
forall {n}. [DotEdge n] -> [DotNode n] -> [DotEdge n]
go []
  where
    go :: [DotEdge n] -> [DotNode n] -> [DotEdge n]
go [DotEdge n]
acc [] = [DotEdge n] -> [DotEdge n]
forall a. [a] -> [a]
reverse [DotEdge n]
acc
    go [DotEdge n]
acc [DotNode n
_] = [DotEdge n] -> [DotEdge n]
forall a. [a] -> [a]
reverse [DotEdge n]
acc
    go [DotEdge n]
acc (DotNode n
a:DotNode n
b:[DotNode n]
rest) = [DotEdge n] -> [DotNode n] -> [DotEdge n]
go (n -> n -> Attributes -> DotEdge n
forall n. n -> n -> Attributes -> DotEdge n
DotEdge (DotNode n -> n
forall n. DotNode n -> n
nodeID DotNode n
a) (DotNode n -> n
forall n. DotNode n -> n
nodeID DotNode n
b) [] DotEdge n -> [DotEdge n] -> [DotEdge n]
forall a. a -> [a] -> [a]
: [DotEdge n]
acc) (DotNode n
bDotNode n -> [DotNode n] -> [DotNode n]
forall a. a -> [a] -> [a]
:[DotNode n]
rest)

newLinesAfter :: String -> Int -> String -> String
newLinesAfter :: String -> Int -> ShowS
newLinesAfter String
esc Int
n String
str = ShowS -> [String] -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
esc) (Int -> String -> [String]
forall e. Int -> [e] -> [[e]]
chunksOf Int
n String
str)