{-# 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 -> FilePath
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 (a -> Rose b -> Rose a
(a -> b) -> Rose a -> Rose b
(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
<$ :: a -> Rose b -> Rose a
$c<$ :: forall a b. a -> Rose b -> Rose a
fmap :: (a -> b) -> Rose a -> Rose b
$cfmap :: forall a b. (a -> b) -> Rose a -> Rose b
Functor, Int -> Rose a -> ShowS
[Rose a] -> ShowS
Rose a -> FilePath
(Int -> Rose a -> ShowS)
-> (Rose a -> FilePath) -> ([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 -> FilePath
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
showList :: [Rose a] -> ShowS
$cshowList :: forall a. Show a => [Rose a] -> ShowS
show :: Rose a -> FilePath
$cshow :: forall a. Show a => Rose a -> FilePath
showsPrec :: Int -> Rose a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Rose a -> ShowS
Show)

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

        -- create barrier nodes

        nThreads :: Int
nThreads = Map Pid [FilePath] -> Int
forall k a. Map k a -> Int
size Map Pid [FilePath]
sfx
        barrierRecord :: [RecordField]
barrierRecord = (\Int
n -> PortName -> RecordField
PortName (PN :: Text -> PortName
PN {portName :: Text
portName = FilePath -> Text
pack (FilePath -> Text) -> FilePath -> Text
forall a b. (a -> b) -> a -> b
$ Int -> FilePath
forall a. Show a => a -> FilePath
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 FilePath
barrierNode = DotNode :: forall n. n -> Attributes -> DotNode n
DotNode {
                        nodeID :: FilePath
nodeID = FilePath
"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, (FilePath, FilePath))]
prefixWithResp = [Int] -> [(FilePath, FilePath)] -> [(Int, (FilePath, FilePath))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
1..] ([(FilePath, FilePath)] -> [(Int, (FilePath, FilePath))])
-> [(FilePath, FilePath)] -> [(Int, (FilePath, FilePath))]
forall a b. (a -> b) -> a -> b
$ FilePath -> [FilePath] -> [(FilePath, FilePath)]
forall a. FilePath -> [a] -> [(a, a)]
byTwoUnsafe FilePath
"prefix" [FilePath]
pref
        prefixNodes :: [DotNode FilePath]
prefixNodes = FilePath -> (Int, (FilePath, FilePath)) -> DotNode FilePath
toDotNode FilePath
"prefix" ((Int, (FilePath, FilePath)) -> DotNode FilePath)
-> [(Int, (FilePath, FilePath))] -> [DotNode FilePath]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, (FilePath, FilePath))]
prefixWithResp
        prefixEdges :: [DotEdge FilePath]
prefixEdges = [DotNode FilePath] -> [DotEdge FilePath]
forall a. [DotNode a] -> [DotEdge a]
connectNodes [DotNode FilePath]
prefixNodes

        -- create suffixes

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

        -- create barrier edges

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

        -- create dot graph

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

        dg :: DotGraph FilePath
dg = DotGraph :: forall n.
Bool -> Bool -> Maybe GraphID -> DotStatements n -> DotGraph n
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
$ FilePath -> Text
pack FilePath
"G")
            , graphStatements :: DotStatements FilePath
graphStatements = DotStatements FilePath
dotStmts
        }

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

toDotNode :: String -> (Int, (String,String)) -> DotNode String
toDotNode :: FilePath -> (Int, (FilePath, FilePath)) -> DotNode FilePath
toDotNode FilePath
nodeIdGroup (Int
n, (FilePath
invocation, FilePath
resp)) =
    DotNode :: forall n. n -> Attributes -> DotNode n
DotNode {
          nodeID :: FilePath
nodeID = (FilePath
nodeIdGroup FilePath -> ShowS
forall a. [a] -> [a] -> [a]
++ FilePath
"-" FilePath -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> FilePath
forall a. Show a => a -> FilePath
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
$ FilePath -> Text
pack (FilePath -> Text) -> FilePath -> Text
forall a b. (a -> b) -> a -> b
$
                (FilePath -> Int -> ShowS
newLinesAfter FilePath
"\\l" Int
60 FilePath
invocation)
                FilePath -> ShowS
forall a. [a] -> [a] -> [a]
++ FilePath
"\\n"
                FilePath -> ShowS
forall a. [a] -> [a] -> [a]
++ (FilePath -> Int -> ShowS
newLinesAfter FilePath
"\\r" Int
60 FilePath
resp)]
    }

byTwoUnsafe :: String -> [a] -> [(a,a)]
byTwoUnsafe :: FilePath -> [a] -> [(a, a)]
byTwoUnsafe FilePath
str [a]
ls = [(a, a)] -> Maybe [(a, a)] -> [(a, a)]
forall a. a -> Maybe a -> a
fromMaybe (FilePath -> [(a, a)]
forall a. HasCallStack => FilePath -> a
error (FilePath -> [(a, a)]) -> FilePath -> [(a, a)]
forall a b. (a -> b) -> a -> b
$ FilePath
"couldn't split " FilePath -> ShowS
forall a. [a] -> [a] -> [a]
++ if FilePath -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null FilePath
str then FilePath
" " else FilePath
str FilePath -> ShowS
forall a. [a] -> [a] -> [a]
++ FilePath
" 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 :: [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 :: [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 :: FilePath -> Int -> ShowS
newLinesAfter FilePath
esc Int
n FilePath
str = ShowS -> [FilePath] -> FilePath
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (FilePath -> ShowS
forall a. [a] -> [a] -> [a]
++ FilePath
esc) (Int -> FilePath -> [FilePath]
forall e. Int -> [e] -> [[e]]
chunksOf Int
n FilePath
str)