{-# LANGUAGE CPP                  #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE PatternGuards        #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE RecordWildCards      #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TupleSections        #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE ViewPatterns         #-}
-- |
-- Module      : Data.Array.Accelerate.Pretty.Graphviz
-- Copyright   : [2015..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
module Data.Array.Accelerate.Pretty.Graphviz (

  Graph,
  PrettyGraph(..), Detail(..),

  graphDelayedAcc, graphDelayedAfun,

) where

import Data.Array.Accelerate.AST
import Data.Array.Accelerate.AST.Idx
import Data.Array.Accelerate.AST.LeftHandSide
import Data.Array.Accelerate.AST.Var
import Data.Array.Accelerate.Analysis.Match
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Pretty.Graphviz.Monad
import Data.Array.Accelerate.Pretty.Graphviz.Type
import Data.Array.Accelerate.Pretty.Print               hiding ( Keyword(..) )
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Stencil
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Sugar.Foreign
import Data.Array.Accelerate.Trafo.Delayed
import Data.Array.Accelerate.Trafo.Substitution

import Control.Applicative                              hiding ( Const, empty )
import Control.Arrow                                    ( (&&&) )
import Control.Monad.State                              ( modify, gets, state )
import Data.HashSet                                     ( HashSet )
import Data.List                                        ( nub, partition )
import Data.Maybe
import Data.String
import Data.Text.Prettyprint.Doc
import System.IO.Unsafe                                 ( unsafePerformIO )
import Prelude                                          hiding ( exp )
import qualified Data.HashSet                           as Set
import qualified Data.Sequence                          as Seq


-- Configuration options
-- ---------------------

cfgIncludeShape, cfgUnique :: Bool
cfgIncludeShape :: Bool
cfgIncludeShape = Bool
False             -- draw edges for uses of shape information
cfgUnique :: Bool
cfgUnique       = Bool
False             -- draw a single edge per data dependency


-- Environments
-- ------------

-- This is the standard environment typed by de Bruijn indices, where at each
-- index we need to record both the pretty printed label as well its 'NodeId',
-- which we use to track data dependencies.
--
data Aval env where
  Aempty ::                                Aval ()
  Apush  :: Aval env -> NodeId -> Label -> Aval (env, t)

-- Convert to the 'Val' used by the base pretty printing module by stripping out
-- the 'NodeId' part.
--
avalToVal :: Aval aenv -> Val aenv
avalToVal :: Aval aenv -> Val aenv
avalToVal Aval aenv
Aempty           = Val aenv
Val ()
Empty
avalToVal (Apush Aval env
aenv NodeId
_ Label
v) = Val env -> Adoc -> Val (env, t)
forall env t. Val env -> Adoc -> Val (env, t)
Push (Aval env -> Val env
forall aenv. Aval aenv -> Val aenv
avalToVal Aval env
aenv) (Label -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Label
v)

aprj :: Idx aenv t -> Aval aenv -> (NodeId, Label)        -- TLM: (Vertex, Label) ??
aprj :: Idx aenv t -> Aval aenv -> (NodeId, Label)
aprj Idx aenv t
ZeroIdx      (Apush Aval env
_    NodeId
n Label
v) = (NodeId
n,Label
v)
aprj (SuccIdx Idx env t
ix) (Apush Aval env
aenv NodeId
_ Label
_) = Idx env t -> Aval env -> (NodeId, Label)
forall aenv t. Idx aenv t -> Aval aenv -> (NodeId, Label)
aprj Idx env t
ix Aval env
Aval env
aenv


-- Graph construction
-- ------------------

mkNode :: PNode -> Maybe Label -> Dot NodeId
mkNode :: PNode -> Maybe Label -> Dot NodeId
mkNode (PNode NodeId
ident Tree (Maybe Label, Adoc)
tree [(Vertex, Maybe Label)]
deps) Maybe Label
label =
  let node :: Node
node  = Maybe Label -> NodeId -> Tree (Maybe Label, Adoc) -> Node
Node Maybe Label
label NodeId
ident Tree (Maybe Label, Adoc)
tree
      edges :: Seq Edge
edges = [Edge] -> Seq Edge
forall a. [a] -> Seq a
Seq.fromList
            ([Edge] -> Seq Edge) -> [Edge] -> Seq Edge
forall a b. (a -> b) -> a -> b
$ ((Vertex, Maybe Label) -> Edge)
-> [(Vertex, Maybe Label)] -> [Edge]
forall a b. (a -> b) -> [a] -> [b]
map (\(Vertex
from, Maybe Label
to) -> Vertex -> Vertex -> Edge
Edge Vertex
from (NodeId -> Maybe Label -> Vertex
Vertex NodeId
ident Maybe Label
to))
            ([(Vertex, Maybe Label)] -> [Edge])
-> [(Vertex, Maybe Label)] -> [Edge]
forall a b. (a -> b) -> a -> b
$ if Bool
cfgUnique then [(Vertex, Maybe Label)] -> [(Vertex, Maybe Label)]
forall a. Eq a => [a] -> [a]
nub [(Vertex, Maybe Label)]
deps else [(Vertex, Maybe Label)]
deps
  in
  (DotState -> (NodeId, DotState)) -> Dot NodeId
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((DotState -> (NodeId, DotState)) -> Dot NodeId)
-> (DotState -> (NodeId, DotState)) -> Dot NodeId
forall a b. (a -> b) -> a -> b
$ \DotState
s ->
    ( NodeId
ident
    , DotState
s { dotNodes :: Seq Node
dotNodes = Node
node  Node -> Seq Node -> Seq Node
forall a. a -> Seq a -> Seq a
Seq.<| DotState -> Seq Node
dotNodes DotState
s
        , dotEdges :: Seq Edge
dotEdges = Seq Edge
edges Seq Edge -> Seq Edge -> Seq Edge
forall a. Seq a -> Seq a -> Seq a
Seq.>< DotState -> Seq Edge
dotEdges DotState
s
        }
    )


-- Add [T|F] ports underneath the given tree.
--
mkTF :: Tree (Maybe Port, Adoc) -> Tree (Maybe Port, Adoc)
mkTF :: Tree (Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
mkTF Tree (Maybe Label, Adoc)
this =
  [Tree (Maybe Label, Adoc)] -> Tree (Maybe Label, Adoc)
forall a. [Tree a] -> Tree a
Forest [ Tree (Maybe Label, Adoc)
this
         , [Tree (Maybe Label, Adoc)] -> Tree (Maybe Label, Adoc)
forall a. [Tree a] -> Tree a
Forest [ (Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a. a -> Tree a
Leaf (Label -> Maybe Label
forall a. a -> Maybe a
Just Label
"T", Adoc
"T")
                  , (Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a. a -> Tree a
Leaf (Label -> Maybe Label
forall a. a -> Maybe a
Just Label
"F", Adoc
"F")
                  ]
         ]


-- Pretty Printing
-- ===============
--
-- The use of unsafePerformIO in the below is safe in the sense that we only
-- require IO to recover the stable names of terms. At worst, if we do not
-- recover the correct stable name for some reason, we will be left with
-- dandling edges in the graph.
--

class PrettyGraph g where
  ppGraph :: Detail -> g -> Graph

instance PrettyGraph (DelayedAcc a) where
  ppGraph :: Detail -> DelayedAcc a -> Graph
ppGraph = Detail -> DelayedAcc a -> Graph
forall a. HasCallStack => Detail -> DelayedAcc a -> Graph
graphDelayedAcc

instance PrettyGraph (DelayedAfun a) where
  ppGraph :: Detail -> DelayedAfun a -> Graph
ppGraph = Detail -> DelayedAfun a -> Graph
forall f. HasCallStack => Detail -> DelayedAfun f -> Graph
graphDelayedAfun

data Detail = Simple | Full

simple :: Detail -> Bool
simple :: Detail -> Bool
simple Detail
Simple = Bool
True
simple Detail
_      = Bool
False

-- | Generate a dependency graph for the given computation
--
{-# NOINLINE graphDelayedAcc #-}
graphDelayedAcc :: HasCallStack => Detail -> DelayedAcc a -> Graph
graphDelayedAcc :: Detail -> DelayedAcc a -> Graph
graphDelayedAcc Detail
detail DelayedAcc a
acc =
  IO Graph -> Graph
forall a. IO a -> a
unsafePerformIO (IO Graph -> Graph) -> IO Graph -> Graph
forall a b. (a -> b) -> a -> b
$! Dot Graph -> IO Graph
forall a. Dot a -> IO a
evalDot (Detail -> Aval () -> DelayedAcc a -> Dot Graph
forall aenv a.
HasCallStack =>
Detail -> Aval aenv -> DelayedOpenAcc aenv a -> Dot Graph
graphDelayedOpenAcc Detail
detail Aval ()
Aempty DelayedAcc a
acc)

-- | Generate a dependency graph for an array function
--
{-# NOINLINE graphDelayedAfun #-}
graphDelayedAfun :: HasCallStack => Detail -> DelayedAfun f -> Graph
graphDelayedAfun :: Detail -> DelayedAfun f -> Graph
graphDelayedAfun Detail
detail DelayedAfun f
afun = IO Graph -> Graph
forall a. IO a -> a
unsafePerformIO (IO Graph -> Graph)
-> (Dot Graph -> IO Graph) -> Dot Graph -> Graph
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Dot Graph -> IO Graph
forall a. Dot a -> IO a
evalDot (Dot Graph -> Graph) -> Dot Graph -> Graph
forall a b. (a -> b) -> a -> b
$! do
  Label
l <- Detail -> Aval () -> DelayedAfun f -> Dot Label
forall aenv afun.
HasCallStack =>
Detail -> Aval aenv -> DelayedOpenAfun aenv afun -> Dot Label
prettyDelayedAfun Detail
detail Aval ()
Aempty DelayedAfun f
afun
  (DotState -> (Graph, DotState)) -> Dot Graph
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((DotState -> (Graph, DotState)) -> Dot Graph)
-> (DotState -> (Graph, DotState)) -> Dot Graph
forall a b. (a -> b) -> a -> b
$ \DotState
s ->
    case Seq Graph -> ViewL Graph
forall a. Seq a -> ViewL a
Seq.viewl (DotState -> Seq Graph
dotGraph DotState
s) of
      g :: Graph
g@(Graph Label
l' [Statement]
_) Seq.:< Seq Graph
gs | Label
l Label -> Label -> Bool
forall a. Eq a => a -> a -> Bool
== Label
l' -> (Graph
g, DotState
s { dotGraph :: Seq Graph
dotGraph = Seq Graph
gs })
      ViewL Graph
_                                  -> String -> (Graph, DotState)
forall a. HasCallStack => String -> a
internalError String
"unexpected error"


-- Pretty-printing data-dependency graphs
-- --------------------------------------

-- Partially constructed graph nodes, consists of some body text and a list of
-- vertices which we will draw edges from (and later, the port we connect into).
--
data PDoc  = PDoc Adoc [Vertex]
data PNode = PNode NodeId (Tree (Maybe Port, Adoc)) [(Vertex, Maybe Port)]

graphDelayedOpenAcc
    :: HasCallStack
    => Detail
    -> Aval aenv
    -> DelayedOpenAcc aenv a
    -> Dot Graph
graphDelayedOpenAcc :: Detail -> Aval aenv -> DelayedOpenAcc aenv a -> Dot Graph
graphDelayedOpenAcc Detail
detail Aval aenv
aenv DelayedOpenAcc aenv a
acc = do
  PNode
r <- Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv a -> Dot PNode
forall aenv arrs.
HasCallStack =>
Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv arrs -> Dot PNode
prettyDelayedOpenAcc Detail
detail Context
context0 Aval aenv
aenv DelayedOpenAcc aenv a
acc
  NodeId
i <- PNode -> Dot NodeId
forall a. a -> Dot NodeId
mkNodeId PNode
r
  NodeId
v <- PNode -> Maybe Label -> Dot NodeId
mkNode PNode
r Maybe Label
forall a. Maybe a
Nothing
  NodeId
_ <- PNode -> Maybe Label -> Dot NodeId
mkNode (NodeId
-> Tree (Maybe Label, Adoc) -> [(Vertex, Maybe Label)] -> PNode
PNode NodeId
i ((Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a. a -> Tree a
Leaf (Maybe Label
forall a. Maybe a
Nothing,Adoc
"result")) [(NodeId -> Maybe Label -> Vertex
Vertex NodeId
v Maybe Label
forall a. Maybe a
Nothing, Maybe Label
forall a. Maybe a
Nothing)]) Maybe Label
forall a. Maybe a
Nothing
  Dot Graph
mkGraph

-- Generate a graph for the given term.
--
prettyDelayedOpenAcc
    :: forall aenv arrs. HasCallStack
    => Detail                               -- simplified output: only print operator name
    -> Context
    -> Aval aenv
    -> DelayedOpenAcc aenv arrs
    -> Dot PNode
prettyDelayedOpenAcc :: Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv arrs -> Dot PNode
prettyDelayedOpenAcc Detail
_      Context
_   Aval aenv
_    Delayed{}            = String -> Dot PNode
forall a. HasCallStack => String -> a
internalError String
"expected manifest array"
prettyDelayedOpenAcc Detail
detail Context
ctx Aval aenv
aenv atop :: DelayedOpenAcc aenv arrs
atop@(Manifest PreOpenAcc DelayedOpenAcc aenv arrs
pacc) =
  case PreOpenAcc DelayedOpenAcc aenv arrs
pacc of
    Avar ArrayVar aenv (Array sh e)
ix                 -> PDoc -> Dot PNode
pnode (ArrayVar aenv (Array sh e) -> PDoc
forall t. ArrayVar aenv t -> PDoc
avar ArrayVar aenv (Array sh e)
ix)
    Alet ALeftHandSide bndArrs aenv aenv'
lhs DelayedOpenAcc aenv bndArrs
bnd DelayedOpenAcc aenv' arrs
body       -> do
      bnd' :: PNode
bnd'@(PNode NodeId
ident Tree (Maybe Label, Adoc)
_ [(Vertex, Maybe Label)]
_) <- Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv bndArrs -> Dot PNode
forall aenv arrs.
HasCallStack =>
Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv arrs -> Dot PNode
prettyDelayedOpenAcc Detail
detail Context
context0 Aval aenv
aenv DelayedOpenAcc aenv bndArrs
bnd
      (Aval aenv'
aenv1, Label
a) <- NodeId
-> Aval aenv
-> ALeftHandSide bndArrs aenv aenv'
-> Dot (Aval aenv', Label)
forall repr aenv aenv'.
HasCallStack =>
NodeId
-> Aval aenv
-> ALeftHandSide repr aenv aenv'
-> Dot (Aval aenv', Label)
prettyLetALeftHandSide NodeId
ident Aval aenv
aenv ALeftHandSide bndArrs aenv aenv'
lhs
      NodeId
_ <- PNode -> Maybe Label -> Dot NodeId
mkNode PNode
bnd' (Label -> Maybe Label
forall a. a -> Maybe a
Just Label
a)
      PNode
body' <- Detail
-> Context -> Aval aenv' -> DelayedOpenAcc aenv' arrs -> Dot PNode
forall aenv arrs.
HasCallStack =>
Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv arrs -> Dot PNode
prettyDelayedOpenAcc Detail
detail Context
context0 Aval aenv'
aenv1 DelayedOpenAcc aenv' arrs
body
      PNode -> Dot PNode
forall (m :: * -> *) a. Monad m => a -> m a
return PNode
body'

    Acond Exp aenv PrimBool
p DelayedOpenAcc aenv arrs
t DelayedOpenAcc aenv arrs
e             -> do
      NodeId
ident <- DelayedOpenAcc aenv arrs -> Dot NodeId
forall a. a -> Dot NodeId
mkNodeId DelayedOpenAcc aenv arrs
atop
      Vertex
vt    <- DelayedOpenAcc aenv arrs -> Dot Vertex
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot Vertex
lift DelayedOpenAcc aenv arrs
t
      Vertex
ve    <- DelayedOpenAcc aenv arrs -> Dot Vertex
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot Vertex
lift DelayedOpenAcc aenv arrs
e
      PDoc Adoc
p' [Vertex]
vs <- Exp aenv PrimBool -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv PrimBool
p
      let port :: Maybe Label
port = Label -> Maybe Label
forall a. a -> Maybe a
Just Label
"P"
          doc :: Tree (Maybe Label, Adoc)
doc  = Tree (Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
mkTF (Tree (Maybe Label, Adoc) -> Tree (Maybe Label, Adoc))
-> Tree (Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a b. (a -> b) -> a -> b
$ (Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a. a -> Tree a
Leaf (Maybe Label
port, if Detail -> Bool
simple Detail
detail then Adoc
"?|" else Adoc
p')
          deps :: [(Vertex, Maybe Label)]
deps = (Vertex
vt, Label -> Maybe Label
forall a. a -> Maybe a
Just Label
"T") (Vertex, Maybe Label)
-> [(Vertex, Maybe Label)] -> [(Vertex, Maybe Label)]
forall a. a -> [a] -> [a]
: (Vertex
ve, Label -> Maybe Label
forall a. a -> Maybe a
Just Label
"F") (Vertex, Maybe Label)
-> [(Vertex, Maybe Label)] -> [(Vertex, Maybe Label)]
forall a. a -> [a] -> [a]
: (Vertex -> (Vertex, Maybe Label))
-> [Vertex] -> [(Vertex, Maybe Label)]
forall a b. (a -> b) -> [a] -> [b]
map (,Maybe Label
port) [Vertex]
vs
      PNode -> Dot PNode
forall (m :: * -> *) a. Monad m => a -> m a
return (PNode -> Dot PNode) -> PNode -> Dot PNode
forall a b. (a -> b) -> a -> b
$ NodeId
-> Tree (Maybe Label, Adoc) -> [(Vertex, Maybe Label)] -> PNode
PNode NodeId
ident Tree (Maybe Label, Adoc)
doc [(Vertex, Maybe Label)]
deps

    Apply ArraysR arrs
_ PreOpenAfun DelayedOpenAcc aenv (arrs1 -> arrs)
afun DelayedOpenAcc aenv arrs1
acc         -> Label -> PNode -> PNode
apply (Label -> PNode -> PNode)
-> Dot Label -> StateT DotState IO (PNode -> PNode)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Detail
-> Aval aenv
-> PreOpenAfun DelayedOpenAcc aenv (arrs1 -> arrs)
-> Dot Label
forall aenv afun.
HasCallStack =>
Detail -> Aval aenv -> DelayedOpenAfun aenv afun -> Dot Label
prettyDelayedAfun    Detail
detail     Aval aenv
aenv PreOpenAfun DelayedOpenAcc aenv (arrs1 -> arrs)
afun
                                      StateT DotState IO (PNode -> PNode) -> Dot PNode -> Dot PNode
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv arrs1 -> Dot PNode
forall aenv arrs.
HasCallStack =>
Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv arrs -> Dot PNode
prettyDelayedOpenAcc Detail
detail Context
ctx Aval aenv
aenv DelayedOpenAcc aenv arrs1
acc

    Awhile PreOpenAfun DelayedOpenAcc aenv (arrs -> Scalar PrimBool)
p PreOpenAfun DelayedOpenAcc aenv (arrs -> arrs)
f DelayedOpenAcc aenv arrs
x             -> do
      NodeId
ident <- DelayedOpenAcc aenv arrs -> Dot NodeId
forall a. a -> Dot NodeId
mkNodeId DelayedOpenAcc aenv arrs
atop
      PNode
x'    <- PNode -> Dot PNode
replant (PNode -> Dot PNode) -> Dot PNode -> Dot PNode
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv arrs -> Dot PNode
forall aenv arrs.
HasCallStack =>
Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv arrs -> Dot PNode
prettyDelayedOpenAcc Detail
detail Context
app Aval aenv
aenv DelayedOpenAcc aenv arrs
x
      Label
p'    <- Detail
-> Aval aenv
-> PreOpenAfun DelayedOpenAcc aenv (arrs -> Scalar PrimBool)
-> Dot Label
forall aenv afun.
HasCallStack =>
Detail -> Aval aenv -> DelayedOpenAfun aenv afun -> Dot Label
prettyDelayedAfun Detail
detail Aval aenv
aenv PreOpenAfun DelayedOpenAcc aenv (arrs -> Scalar PrimBool)
p
      Label
f'    <- Detail
-> Aval aenv
-> PreOpenAfun DelayedOpenAcc aenv (arrs -> arrs)
-> Dot Label
forall aenv afun.
HasCallStack =>
Detail -> Aval aenv -> DelayedOpenAfun aenv afun -> Dot Label
prettyDelayedAfun Detail
detail Aval aenv
aenv PreOpenAfun DelayedOpenAcc aenv (arrs -> arrs)
f
      --
      let PNode NodeId
_ (Leaf (Maybe Label
Nothing,Adoc
xb)) [(Vertex, Maybe Label)]
fvs = PNode
x'
          loop :: Adoc
loop                            = Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
nest Int
2 ([Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep [Adoc
"awhile", Label -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Label
p', Label -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Label
f', Adoc
xb ])
      PNode -> Dot PNode
forall (m :: * -> *) a. Monad m => a -> m a
return (PNode -> Dot PNode) -> PNode -> Dot PNode
forall a b. (a -> b) -> a -> b
$ NodeId
-> Tree (Maybe Label, Adoc) -> [(Vertex, Maybe Label)] -> PNode
PNode NodeId
ident ((Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a. a -> Tree a
Leaf (Maybe Label
forall a. Maybe a
Nothing,Adoc
loop)) [(Vertex, Maybe Label)]
fvs

    a :: PreOpenAcc DelayedOpenAcc aenv arrs
a@(Apair DelayedOpenAcc aenv as
a1 DelayedOpenAcc aenv bs
a2)          -> PreOpenAcc DelayedOpenAcc aenv arrs -> Dot NodeId
forall a. a -> Dot NodeId
mkNodeId PreOpenAcc DelayedOpenAcc aenv arrs
a Dot NodeId -> (NodeId -> Dot PNode) -> Dot PNode
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Detail
-> Aval aenv
-> DelayedOpenAcc aenv as
-> DelayedOpenAcc aenv bs
-> NodeId
-> Dot PNode
forall aenv a1 a2.
HasCallStack =>
Detail
-> Aval aenv
-> DelayedOpenAcc aenv a1
-> DelayedOpenAcc aenv a2
-> NodeId
-> Dot PNode
prettyDelayedApair Detail
detail Aval aenv
aenv DelayedOpenAcc aenv as
a1 DelayedOpenAcc aenv bs
a2

    PreOpenAcc DelayedOpenAcc aenv arrs
Anil                     -> Operator
"()"             Operator -> [Dot PDoc] -> Dot PNode
.$ []

    Use ArrayR (Array sh e)
repr Array sh e
arr             -> Operator
"use"            Operator -> [Dot PDoc] -> Dot PNode
.$ [ PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return (PDoc -> Dot PDoc) -> PDoc -> Dot PDoc
forall a b. (a -> b) -> a -> b
$ Adoc -> [Vertex] -> PDoc
PDoc (ArrayR (Array sh e) -> Array sh e -> Adoc
forall sh e. ArrayR (Array sh e) -> Array sh e -> Adoc
prettyArray ArrayR (Array sh e)
repr Array sh e
arr) [] ]
    Unit TypeR e
_ Exp aenv e
e                 -> Operator
"unit"           Operator -> [Dot PDoc] -> Dot PNode
.$ [ Exp aenv e -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv e
e ]
    Generate ArrayR (Array sh e)
_ Exp aenv sh
sh Fun aenv (sh -> e)
f          -> Operator
"generate"       Operator -> [Dot PDoc] -> Dot PNode
.$ [ Exp aenv sh -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv sh
sh, Fun aenv (sh -> e) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (sh -> e)
f ]
    Transform ArrayR (Array sh' b)
_ Exp aenv sh'
sh Fun aenv (sh' -> sh)
ix Fun aenv (a -> b)
f DelayedOpenAcc aenv (Array sh a)
xs   -> Operator
"transform"      Operator -> [Dot PDoc] -> Dot PNode
.$ [ Exp aenv sh' -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv sh'
sh, Fun aenv (sh' -> sh) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (sh' -> sh)
ix, Fun aenv (a -> b) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (a -> b)
f, DelayedOpenAcc aenv (Array sh a) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh a)
xs ]
    Reshape ShapeR sh
_ Exp aenv sh
sh DelayedOpenAcc aenv (Array sh' e)
xs          -> Operator
"reshape"        Operator -> [Dot PDoc] -> Dot PNode
.$ [ Exp aenv sh -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv sh
sh, DelayedOpenAcc aenv (Array sh' e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh' e)
xs ]
    Replicate SliceIndex slix sl co sh
_ty Exp aenv slix
ix DelayedOpenAcc aenv (Array sl e)
xs      -> Operator
"replicate"      Operator -> [Dot PDoc] -> Dot PNode
.$ [ Exp aenv slix -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv slix
ix, DelayedOpenAcc aenv (Array sl e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sl e)
xs ]
    Slice SliceIndex slix sl co sh
_ty DelayedOpenAcc aenv (Array sh e)
xs Exp aenv slix
ix          -> Operator
"slice"          Operator -> [Dot PDoc] -> Dot PNode
.$ [ DelayedOpenAcc aenv (Array sh e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh e)
xs, Exp aenv slix -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv slix
ix ]
    Map TypeR e'
_ Fun aenv (e -> e')
f DelayedOpenAcc aenv (Array sh e)
xs               -> Operator
"map"            Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (e -> e') -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (e -> e')
f, DelayedOpenAcc aenv (Array sh e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh e)
xs ]
    ZipWith TypeR e3
_ Fun aenv (e1 -> e2 -> e3)
f DelayedOpenAcc aenv (Array sh e1)
xs DelayedOpenAcc aenv (Array sh e2)
ys        -> Operator
"zipWith"        Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (e1 -> e2 -> e3) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (e1 -> e2 -> e3)
f, DelayedOpenAcc aenv (Array sh e1) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh e1)
xs, DelayedOpenAcc aenv (Array sh e2) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh e2)
ys ]
    Fold Fun aenv (e -> e -> e)
f (Just Exp aenv e
z) DelayedOpenAcc aenv (Array (sh, Int) e)
a        -> Operator
"fold"           Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (e -> e -> e) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (e -> e -> e)
f,  Exp aenv e -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv e
z, DelayedOpenAcc aenv (Array (sh, Int) e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array (sh, Int) e)
a ]
    Fold Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
Nothing  DelayedOpenAcc aenv (Array (sh, Int) e)
a        -> Operator
"fold1"          Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (e -> e -> e) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (e -> e -> e)
f,  DelayedOpenAcc aenv (Array (sh, Int) e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array (sh, Int) e)
a ]
    FoldSeg IntegralType i
_ Fun aenv (e -> e -> e)
f (Just Exp aenv e
z) DelayedOpenAcc aenv (Array (sh, Int) e)
a DelayedOpenAcc aenv (Segments i)
s -> Operator
"foldSeg"        Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (e -> e -> e) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (e -> e -> e)
f,  Exp aenv e -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv e
z, DelayedOpenAcc aenv (Array (sh, Int) e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array (sh, Int) e)
a, DelayedOpenAcc aenv (Segments i) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Segments i)
s ]
    FoldSeg IntegralType i
_ Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
Nothing  DelayedOpenAcc aenv (Array (sh, Int) e)
a DelayedOpenAcc aenv (Segments i)
s -> Operator
"fold1Seg"       Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (e -> e -> e) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (e -> e -> e)
f,  DelayedOpenAcc aenv (Array (sh, Int) e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array (sh, Int) e)
a, DelayedOpenAcc aenv (Segments i) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Segments i)
s ]
    Scan Direction
d Fun aenv (e -> e -> e)
f (Just Exp aenv e
z) DelayedOpenAcc aenv (Array (sh, Int) e)
a      -> String -> Direction -> String -> Operator
ppD String
"scan" Direction
d String
""  Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (e -> e -> e) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (e -> e -> e)
f,  Exp aenv e -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv e
z, DelayedOpenAcc aenv (Array (sh, Int) e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array (sh, Int) e)
a ]
    Scan Direction
d Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
Nothing  DelayedOpenAcc aenv (Array (sh, Int) e)
a      -> String -> Direction -> String -> Operator
ppD String
"scan" Direction
d String
"1" Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (e -> e -> e) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (e -> e -> e)
f,  DelayedOpenAcc aenv (Array (sh, Int) e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array (sh, Int) e)
a ]
    Scan' Direction
d Fun aenv (e -> e -> e)
f Exp aenv e
z DelayedOpenAcc aenv (Array (sh, Int) e)
a            -> String -> Direction -> String -> Operator
ppD String
"scan" Direction
d String
"'" Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (e -> e -> e) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (e -> e -> e)
f,  Exp aenv e -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv e
z, DelayedOpenAcc aenv (Array (sh, Int) e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array (sh, Int) e)
a ]
    Permute Fun aenv (e -> e -> e)
f DelayedOpenAcc aenv (Array sh' e)
dfts Fun aenv (sh -> PrimMaybe sh')
p DelayedOpenAcc aenv (Array sh e)
xs      -> Operator
"permute"        Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (e -> e -> e) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (e -> e -> e)
f, DelayedOpenAcc aenv (Array sh' e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh' e)
dfts, Fun aenv (sh -> PrimMaybe sh') -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (sh -> PrimMaybe sh')
p, DelayedOpenAcc aenv (Array sh e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh e)
xs ]
    Backpermute ShapeR sh'
_ Exp aenv sh'
sh Fun aenv (sh' -> sh)
p DelayedOpenAcc aenv (Array sh e)
xs    -> Operator
"backpermute"    Operator -> [Dot PDoc] -> Dot PNode
.$ [ Exp aenv sh' -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv sh'
sh, Fun aenv (sh' -> sh) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (sh' -> sh)
p, DelayedOpenAcc aenv (Array sh e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh e)
xs ]
    Stencil StencilR sh e stencil
s TypeR e'
_ Fun aenv (stencil -> e')
sten Boundary aenv (Array sh e)
bndy DelayedOpenAcc aenv (Array sh e)
xs -> Operator
"stencil"        Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (stencil -> e') -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (stencil -> e')
sten, TypeR e -> Boundary aenv (Array sh e) -> Dot PDoc
forall sh e.
HasCallStack =>
TypeR e -> Boundary aenv (Array sh e) -> Dot PDoc
ppB (StencilR sh e stencil -> TypeR e
forall sh e pat. StencilR sh e pat -> TypeR e
stencilEltR StencilR sh e stencil
s) Boundary aenv (Array sh e)
bndy, DelayedOpenAcc aenv (Array sh e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh e)
xs ]
    Stencil2 StencilR sh a stencil1
s1 StencilR sh b stencil2
s2 TypeR c
_ Fun aenv (stencil1 -> stencil2 -> c)
sten Boundary aenv (Array sh a)
bndy1 DelayedOpenAcc aenv (Array sh a)
acc1 Boundary aenv (Array sh b)
bndy2 DelayedOpenAcc aenv (Array sh b)
acc2
                            -> Operator
"stencil2"        Operator -> [Dot PDoc] -> Dot PNode
.$ [ Fun aenv (stencil1 -> stencil2 -> c) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (stencil1 -> stencil2 -> c)
sten, TypeR a -> Boundary aenv (Array sh a) -> Dot PDoc
forall sh e.
HasCallStack =>
TypeR e -> Boundary aenv (Array sh e) -> Dot PDoc
ppB (StencilR sh a stencil1 -> TypeR a
forall sh e pat. StencilR sh e pat -> TypeR e
stencilEltR StencilR sh a stencil1
s1) Boundary aenv (Array sh a)
bndy1, DelayedOpenAcc aenv (Array sh a) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh a)
acc1, TypeR b -> Boundary aenv (Array sh b) -> Dot PDoc
forall sh e.
HasCallStack =>
TypeR e -> Boundary aenv (Array sh e) -> Dot PDoc
ppB (StencilR sh b stencil2 -> TypeR b
forall sh e pat. StencilR sh e pat -> TypeR e
stencilEltR StencilR sh b stencil2
s2) Boundary aenv (Array sh b)
bndy2, DelayedOpenAcc aenv (Array sh b) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv (Array sh b)
acc2 ]
    Aforeign ArraysR arrs
_ asm (as -> arrs)
ff PreAfun DelayedOpenAcc (as -> arrs)
_afun DelayedOpenAcc aenv as
xs  -> Operator
"aforeign"        Operator -> [Dot PDoc] -> Dot PNode
.$ [ PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return (Adoc -> [Vertex] -> PDoc
PDoc (String -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty (asm (as -> arrs) -> String
forall (asm :: * -> *) args. Foreign asm => asm args -> String
strForeign asm (as -> arrs)
ff)) []), {- ppAf afun, -} DelayedOpenAcc aenv as -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA DelayedOpenAcc aenv as
xs ]
    -- Collect{}               -> error "Collect"

  where
    (.$) :: Operator -> [Dot PDoc] -> Dot PNode
    Operator
name .$ :: Operator -> [Dot PDoc] -> Dot PNode
.$ [Dot PDoc]
docs = PDoc -> Dot PNode
pnode (PDoc -> Dot PNode) -> Dot PDoc -> Dot PNode
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Operator -> [Dot PDoc] -> Dot PDoc
fmt Operator
name [Dot PDoc]
docs

    fmt :: Operator -> [Dot PDoc] -> Dot PDoc
    fmt :: Operator -> [Dot PDoc] -> Dot PDoc
fmt Operator
name [Dot PDoc]
docs = do
      [PDoc]
docs' <- [Dot PDoc] -> StateT DotState IO [PDoc]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [Dot PDoc]
docs
      let args :: [Adoc]
args = [ Adoc
x | PDoc Adoc
x [Vertex]
_ <- [PDoc]
docs' ]
          fvs :: [[Vertex]]
fvs  = [ [Vertex]
x | PDoc Adoc
_ [Vertex]
x <- [PDoc]
docs' ]
          doc :: Adoc
doc  = if Detail -> Bool
simple Detail
detail
                   then Operator -> Adoc
manifest Operator
name
                   else Bool -> Adoc -> Adoc
forall ann. Bool -> Doc ann -> Doc ann
parensIf (Context -> Operator -> Bool
needsParens Context
ctx Operator
name)
                      (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ Int -> Adoc -> Adoc
forall ann. Int -> Doc ann -> Doc ann
nest Int
shiftwidth
                      (Adoc -> Adoc) -> Adoc -> Adoc
forall a b. (a -> b) -> a -> b
$ [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
sep ( Operator -> Adoc
manifest Operator
name Adoc -> [Adoc] -> [Adoc]
forall a. a -> [a] -> [a]
: [Adoc]
args )
      PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return (PDoc -> Dot PDoc) -> PDoc -> Dot PDoc
forall a b. (a -> b) -> a -> b
$ Adoc -> [Vertex] -> PDoc
PDoc Adoc
doc ([[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Vertex]]
fvs)

    pnode :: PDoc -> Dot PNode
    pnode :: PDoc -> Dot PNode
pnode (PDoc Adoc
doc [Vertex]
vs) = do
      let port :: Maybe a
port = Maybe a
forall a. Maybe a
Nothing
      NodeId
ident <- DelayedOpenAcc aenv arrs -> Dot NodeId
forall a. a -> Dot NodeId
mkNodeId DelayedOpenAcc aenv arrs
atop
      PNode -> Dot PNode
forall (m :: * -> *) a. Monad m => a -> m a
return (PNode -> Dot PNode) -> PNode -> Dot PNode
forall a b. (a -> b) -> a -> b
$ NodeId
-> Tree (Maybe Label, Adoc) -> [(Vertex, Maybe Label)] -> PNode
PNode NodeId
ident ((Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a. a -> Tree a
Leaf (Maybe Label
forall a. Maybe a
port, Adoc
doc)) ((Vertex -> (Vertex, Maybe Label))
-> [Vertex] -> [(Vertex, Maybe Label)]
forall a b. (a -> b) -> [a] -> [b]
map (,Maybe Label
forall a. Maybe a
port) [Vertex]
vs)

    -- Free variables
    --
    fvF :: Fun aenv t -> [Vertex]
    fvF :: Fun aenv t -> [Vertex]
fvF = Val () -> Aval aenv -> Fun aenv t -> [Vertex]
forall env aenv fun.
Val env -> Aval aenv -> OpenFun env aenv fun -> [Vertex]
fvOpenFun Val ()
Empty Aval aenv
aenv

    fvE :: Exp aenv t -> [Vertex]
    fvE :: Exp aenv t -> [Vertex]
fvE = Val () -> Aval aenv -> Exp aenv t -> [Vertex]
forall env aenv exp.
Val env -> Aval aenv -> OpenExp env aenv exp -> [Vertex]
fvOpenExp Val ()
Empty Aval aenv
aenv

    -- Pretty-printing
    --
    avar :: ArrayVar aenv t -> PDoc
    avar :: ArrayVar aenv t -> PDoc
avar (Var ArrayR t
_ Idx aenv t
ix) = let (NodeId
ident, Label
v) = Idx aenv t -> Aval aenv -> (NodeId, Label)
forall aenv t. Idx aenv t -> Aval aenv -> (NodeId, Label)
aprj Idx aenv t
ix Aval aenv
aenv
                      in  Adoc -> [Vertex] -> PDoc
PDoc (Label -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Label
v) [NodeId -> Maybe Label -> Vertex
Vertex NodeId
ident Maybe Label
forall a. Maybe a
Nothing]

    aenv' :: Val aenv
    aenv' :: Val aenv
aenv' = Aval aenv -> Val aenv
forall aenv. Aval aenv -> Val aenv
avalToVal Aval aenv
aenv

    ppA :: HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
    ppA :: DelayedOpenAcc aenv a -> Dot PDoc
ppA (Manifest (Avar ArrayVar aenv (Array sh e)
ix)) = PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayVar aenv (Array sh e) -> PDoc
forall t. ArrayVar aenv t -> PDoc
avar ArrayVar aenv (Array sh e)
ix)
    ppA acc :: DelayedOpenAcc aenv a
acc@Manifest{}       = do
      -- Lift out and draw as a separate node. This can occur with the manifest
      -- array arguments to permute (defaults array) and stencil[2].
      PNode
acc'  <- Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv a -> Dot PNode
forall aenv arrs.
HasCallStack =>
Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv arrs -> Dot PNode
prettyDelayedOpenAcc Detail
detail Context
app Aval aenv
aenv DelayedOpenAcc aenv a
acc
      Label
v     <- Dot Label
mkLabel
      NodeId
ident <- PNode -> Maybe Label -> Dot NodeId
mkNode PNode
acc' (Label -> Maybe Label
forall a. a -> Maybe a
Just Label
v)
      PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return (PDoc -> Dot PDoc) -> PDoc -> Dot PDoc
forall a b. (a -> b) -> a -> b
$ Adoc -> [Vertex] -> PDoc
PDoc (Label -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Label
v) [NodeId -> Maybe Label -> Vertex
Vertex NodeId
ident Maybe Label
forall a. Maybe a
Nothing]
    ppA (Delayed ArrayR (Array sh e)
_ Exp aenv sh
sh Fun aenv (sh -> e)
f Fun aenv (Int -> e)
_)
      | Shape ArrayVar aenv (Array sh e)
a    <- Exp aenv sh
sh                   -- identical shape
      , Just ArrayVar aenv (Array sh e)
b     <- Fun aenv (sh -> e) -> Maybe (ArrayVar aenv (Array sh e))
forall env aenv a b.
OpenFun env aenv (a -> b) -> Maybe (ArrayVar aenv (Array a b))
isIdentityIndexing Fun aenv (sh -> e)
f -- function is `\ix -> b ! ix`
      , Just Array sh e :~: Array sh e
Refl  <- ArrayVar aenv (Array sh e)
-> ArrayVar aenv (Array sh e) -> Maybe (Array sh e :~: Array sh e)
forall (s :: * -> *) env t1 t2.
Var s env t1 -> Var s env t2 -> Maybe (t1 :~: t2)
matchVar ArrayVar aenv (Array sh e)
a ArrayVar aenv (Array sh e)
b         -- function thus is `\ix -> a ! ix`
      = DelayedOpenAcc aenv (Array sh e) -> Dot PDoc
forall a. HasCallStack => DelayedOpenAcc aenv a -> Dot PDoc
ppA (DelayedOpenAcc aenv (Array sh e) -> Dot PDoc)
-> DelayedOpenAcc aenv (Array sh e) -> Dot PDoc
forall a b. (a -> b) -> a -> b
$ PreOpenAcc DelayedOpenAcc aenv (Array sh e)
-> DelayedOpenAcc aenv (Array sh e)
forall aenv a.
PreOpenAcc DelayedOpenAcc aenv a -> DelayedOpenAcc aenv a
Manifest (PreOpenAcc DelayedOpenAcc aenv (Array sh e)
 -> DelayedOpenAcc aenv (Array sh e))
-> PreOpenAcc DelayedOpenAcc aenv (Array sh e)
-> DelayedOpenAcc aenv (Array sh e)
forall a b. (a -> b) -> a -> b
$ ArrayVar aenv (Array sh e)
-> PreOpenAcc DelayedOpenAcc aenv (Array sh e)
forall aenv sh e (acc :: * -> * -> *).
ArrayVar aenv (Array sh e) -> PreOpenAcc acc aenv (Array sh e)
Avar ArrayVar aenv (Array sh e)
a
    ppA (Delayed ArrayR (Array sh e)
_ Exp aenv sh
sh Fun aenv (sh -> e)
f Fun aenv (Int -> e)
_) = do
      PDoc Adoc
d [Vertex]
v <- Operator
"Delayed" Operator -> [Dot PDoc] -> Dot PDoc
`fmt` [ Exp aenv sh -> Dot PDoc
forall t. HasCallStack => Exp aenv t -> Dot PDoc
ppE Exp aenv sh
sh, Fun aenv (sh -> e) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (sh -> e)
f ]
      PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return    (PDoc -> Dot PDoc) -> PDoc -> Dot PDoc
forall a b. (a -> b) -> a -> b
$ Adoc -> [Vertex] -> PDoc
PDoc (Adoc -> Adoc
forall ann. Doc ann -> Doc ann
parens Adoc
d) [Vertex]
v

    ppB :: forall sh e. HasCallStack
        => TypeR e
        -> Boundary aenv (Array sh e)
        -> Dot PDoc
    ppB :: TypeR e -> Boundary aenv (Array sh e) -> Dot PDoc
ppB TypeR e
_  Boundary aenv (Array sh e)
Clamp        = PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return (Adoc -> [Vertex] -> PDoc
PDoc Adoc
"clamp"  [])
    ppB TypeR e
_  Boundary aenv (Array sh e)
Mirror       = PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return (Adoc -> [Vertex] -> PDoc
PDoc Adoc
"mirror" [])
    ppB TypeR e
_  Boundary aenv (Array sh e)
Wrap         = PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return (Adoc -> [Vertex] -> PDoc
PDoc Adoc
"wrap"   [])
    ppB TypeR e
tp (Constant e
e) = PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return (Adoc -> [Vertex] -> PDoc
PDoc (TypeR e -> e -> Adoc
forall e. TypeR e -> e -> Adoc
prettyConst TypeR e
tp e
e
e) [])
    ppB TypeR e
_  (Function Fun aenv (sh -> e)
f) = Fun aenv (sh -> e) -> Dot PDoc
forall t. HasCallStack => Fun aenv t -> Dot PDoc
ppF Fun aenv (sh -> e)
f

    ppF :: HasCallStack => Fun aenv t -> Dot PDoc
    ppF :: Fun aenv t -> Dot PDoc
ppF = PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return (PDoc -> Dot PDoc)
-> (Fun aenv t -> PDoc) -> Fun aenv t -> Dot PDoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Adoc -> [Vertex] -> PDoc) -> (Adoc, [Vertex]) -> PDoc
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Adoc -> [Vertex] -> PDoc
PDoc ((Adoc, [Vertex]) -> PDoc)
-> (Fun aenv t -> (Adoc, [Vertex])) -> Fun aenv t -> PDoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Adoc -> Adoc
forall ann. Doc ann -> Doc ann
parens (Adoc -> Adoc) -> (Fun aenv t -> Adoc) -> Fun aenv t -> Adoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Val aenv -> Fun aenv t -> Adoc
forall aenv f. Val aenv -> Fun aenv f -> Adoc
prettyFun Val aenv
aenv' (Fun aenv t -> Adoc)
-> (Fun aenv t -> [Vertex]) -> Fun aenv t -> (Adoc, [Vertex])
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& Fun aenv t -> [Vertex]
forall t. Fun aenv t -> [Vertex]
fvF)

    ppE :: HasCallStack => Exp aenv t -> Dot PDoc
    ppE :: Exp aenv t -> Dot PDoc
ppE = PDoc -> Dot PDoc
forall (m :: * -> *) a. Monad m => a -> m a
return (PDoc -> Dot PDoc)
-> (Exp aenv t -> PDoc) -> Exp aenv t -> Dot PDoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Adoc -> [Vertex] -> PDoc) -> (Adoc, [Vertex]) -> PDoc
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Adoc -> [Vertex] -> PDoc
PDoc ((Adoc, [Vertex]) -> PDoc)
-> (Exp aenv t -> (Adoc, [Vertex])) -> Exp aenv t -> PDoc
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Val aenv -> Exp aenv t -> Adoc
forall aenv t. Val aenv -> Exp aenv t -> Adoc
prettyExp Val aenv
aenv' (Exp aenv t -> Adoc)
-> (Exp aenv t -> [Vertex]) -> Exp aenv t -> (Adoc, [Vertex])
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& Exp aenv t -> [Vertex]
forall t. Exp aenv t -> [Vertex]
fvE)

    ppD :: String -> Direction -> String -> Operator
    ppD :: String -> Direction -> String -> Operator
ppD String
f Direction
LeftToRight String
k = String -> Operator
forall a. IsString a => String -> a
fromString (String
f String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"l" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
k)
    ppD String
f Direction
RightToLeft String
k = String -> Operator
forall a. IsString a => String -> a
fromString (String
f String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"r" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
k)

    lift :: HasCallStack => DelayedOpenAcc aenv a -> Dot Vertex
    lift :: DelayedOpenAcc aenv a -> Dot Vertex
lift Delayed{}                    = String -> Dot Vertex
forall a. HasCallStack => String -> a
internalError String
"expected manifest array"
    lift (Manifest (Avar (Var ArrayR (Array sh e)
_ Idx aenv (Array sh e)
ix))) = Vertex -> Dot Vertex
forall (m :: * -> *) a. Monad m => a -> m a
return (Vertex -> Dot Vertex) -> Vertex -> Dot Vertex
forall a b. (a -> b) -> a -> b
$ NodeId -> Maybe Label -> Vertex
Vertex ((NodeId, Label) -> NodeId
forall a b. (a, b) -> a
fst (Idx aenv (Array sh e) -> Aval aenv -> (NodeId, Label)
forall aenv t. Idx aenv t -> Aval aenv -> (NodeId, Label)
aprj Idx aenv (Array sh e)
ix Aval aenv
aenv)) Maybe Label
forall a. Maybe a
Nothing
    lift DelayedOpenAcc aenv a
acc                          = do
      PNode
acc'  <- Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv a -> Dot PNode
forall aenv arrs.
HasCallStack =>
Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv arrs -> Dot PNode
prettyDelayedOpenAcc Detail
detail Context
context0 Aval aenv
aenv DelayedOpenAcc aenv a
acc
      NodeId
ident <- PNode -> Maybe Label -> Dot NodeId
mkNode PNode
acc' Maybe Label
forall a. Maybe a
Nothing
      Vertex -> Dot Vertex
forall (m :: * -> *) a. Monad m => a -> m a
return (Vertex -> Dot Vertex) -> Vertex -> Dot Vertex
forall a b. (a -> b) -> a -> b
$ NodeId -> Maybe Label -> Vertex
Vertex NodeId
ident Maybe Label
forall a. Maybe a
Nothing

    apply :: Label -> PNode -> PNode
    apply :: Label -> PNode -> PNode
apply Label
f (PNode NodeId
ident Tree (Maybe Label, Adoc)
x [(Vertex, Maybe Label)]
vs) =
      let x' :: Tree (Maybe Label, Adoc)
x' = case Tree (Maybe Label, Adoc)
x of
                 Leaf (Maybe Label
p,Adoc
d) -> (Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a. a -> Tree a
Leaf (Maybe Label
p, Label -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Label
f Adoc -> Adoc -> Adoc
forall ann. Doc ann -> Doc ann -> Doc ann
<+> Adoc
d)
                 Forest [Tree (Maybe Label, Adoc)]
ts  -> [Tree (Maybe Label, Adoc)] -> Tree (Maybe Label, Adoc)
forall a. [Tree a] -> Tree a
Forest ((Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a. a -> Tree a
Leaf (Maybe Label
forall a. Maybe a
Nothing,Label -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Label
f) Tree (Maybe Label, Adoc)
-> [Tree (Maybe Label, Adoc)] -> [Tree (Maybe Label, Adoc)]
forall a. a -> [a] -> [a]
: [Tree (Maybe Label, Adoc)]
ts)
      in
      NodeId
-> Tree (Maybe Label, Adoc) -> [(Vertex, Maybe Label)] -> PNode
PNode NodeId
ident Tree (Maybe Label, Adoc)
x' [(Vertex, Maybe Label)]
vs


-- Pretty print array functions as separate sub-graphs, and return the name of
-- the sub-graph as if it can be called like a function. We will add additional
-- nodes at the top of the graph to represent the bound variables.
--
-- Note: [Edge placement]
--
-- If a node belongs to a particular graph, so too must all its edges (and
-- vertices). This means that if the subgraph references anything from the
-- enclosing environment, we must lift those edges out of this subgraph,
-- otherwise the referenced node will be drawn inside of the subgraph.
--
prettyDelayedAfun
    :: HasCallStack
    => Detail
    -> Aval aenv
    -> DelayedOpenAfun aenv afun
    -> Dot Label
prettyDelayedAfun :: Detail -> Aval aenv -> DelayedOpenAfun aenv afun -> Dot Label
prettyDelayedAfun Detail
detail Aval aenv
aenv DelayedOpenAfun aenv afun
afun = do
  Graph Label
_ [Statement]
ss  <- Dot Graph -> Dot Graph
mkSubgraph (Aval aenv -> DelayedOpenAfun aenv afun -> Dot Graph
forall aenv' a'.
Aval aenv' -> DelayedOpenAfun aenv' a' -> Dot Graph
go Aval aenv
aenv DelayedOpenAfun aenv afun
afun)
  Int
n           <- Seq Graph -> Int
forall a. Seq a -> Int
Seq.length (Seq Graph -> Int)
-> StateT DotState IO (Seq Graph) -> StateT DotState IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (DotState -> Seq Graph) -> StateT DotState IO (Seq Graph)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets DotState -> Seq Graph
dotGraph
  let label :: Label
label         = Label
"afun" Label -> Label -> Label
forall a. Semigroup a => a -> a -> a
<> String -> Label
forall a. IsString a => String -> a
fromString (Int -> String
forall a. Show a => a -> String
show (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))
      outer :: HashSet NodeId
outer         = Aval aenv -> HashSet NodeId
forall aenv'. Aval aenv' -> HashSet NodeId
collect Aval aenv
aenv
      ([Statement]
lifted,[Statement]
ss')  =
        ((Statement -> Bool) -> [Statement] -> ([Statement], [Statement]))
-> [Statement] -> (Statement -> Bool) -> ([Statement], [Statement])
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Statement -> Bool) -> [Statement] -> ([Statement], [Statement])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition [Statement]
ss ((Statement -> Bool) -> ([Statement], [Statement]))
-> (Statement -> Bool) -> ([Statement], [Statement])
forall a b. (a -> b) -> a -> b
$ \Statement
s ->
        case Statement
s of
          E (Edge (Vertex NodeId
ident Maybe Label
_) Vertex
_) -> NodeId -> HashSet NodeId -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
Set.member NodeId
ident HashSet NodeId
outer
          Statement
_                           -> Bool
False
  --
  (DotState -> DotState) -> StateT DotState IO ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((DotState -> DotState) -> StateT DotState IO ())
-> (DotState -> DotState) -> StateT DotState IO ()
forall a b. (a -> b) -> a -> b
$ \DotState
s -> DotState
s { dotGraph :: Seq Graph
dotGraph = DotState -> Seq Graph
dotGraph DotState
s                         Seq Graph -> Graph -> Seq Graph
forall a. Seq a -> a -> Seq a
Seq.|> Label -> [Statement] -> Graph
Graph Label
label [Statement]
ss'
                   , dotEdges :: Seq Edge
dotEdges = [Edge] -> Seq Edge
forall a. [a] -> Seq a
Seq.fromList [ Edge
e | E Edge
e <- [Statement]
lifted ] Seq Edge -> Seq Edge -> Seq Edge
forall a. Seq a -> Seq a -> Seq a
Seq.>< DotState -> Seq Edge
dotEdges DotState
s
                   }
  Label -> Dot Label
forall (m :: * -> *) a. Monad m => a -> m a
return Label
label
  where
    go :: Aval aenv' -> DelayedOpenAfun aenv' a' -> Dot Graph
    go :: Aval aenv' -> DelayedOpenAfun aenv' a' -> Dot Graph
go Aval aenv'
aenv' (Abody DelayedOpenAcc aenv' a'
b) = Detail -> Aval aenv' -> DelayedOpenAcc aenv' a' -> Dot Graph
forall aenv a.
HasCallStack =>
Detail -> Aval aenv -> DelayedOpenAcc aenv a -> Dot Graph
graphDelayedOpenAcc Detail
detail Aval aenv'
aenv' DelayedOpenAcc aenv' a'
b
    go Aval aenv'
aenv' (Alam ALeftHandSide a aenv' aenv'
lhs PreOpenAfun DelayedOpenAcc aenv' t
f) = do
      Aval aenv'
aenv'' <- Aval aenv' -> ALeftHandSide a aenv' aenv' -> Dot (Aval aenv')
forall repr aenv aenv'.
HasCallStack =>
Aval aenv -> ALeftHandSide repr aenv aenv' -> Dot (Aval aenv')
prettyLambdaALeftHandSide Aval aenv'
aenv' ALeftHandSide a aenv' aenv'
lhs
      Aval aenv' -> PreOpenAfun DelayedOpenAcc aenv' t -> Dot Graph
forall aenv' a'.
Aval aenv' -> DelayedOpenAfun aenv' a' -> Dot Graph
go Aval aenv'
aenv'' PreOpenAfun DelayedOpenAcc aenv' t
f

    collect :: Aval aenv' -> HashSet NodeId
    collect :: Aval aenv' -> HashSet NodeId
collect Aval aenv'
Aempty        = HashSet NodeId
forall a. HashSet a
Set.empty
    collect (Apush Aval env
a NodeId
i Label
_) = NodeId -> HashSet NodeId -> HashSet NodeId
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.insert NodeId
i (Aval env -> HashSet NodeId
forall aenv'. Aval aenv' -> HashSet NodeId
collect Aval env
a)

prettyLetALeftHandSide
    :: forall repr aenv aenv'. HasCallStack
    => NodeId
    -> Aval aenv
    -> ALeftHandSide repr aenv aenv'
    -> Dot (Aval aenv', Label)
prettyLetALeftHandSide :: NodeId
-> Aval aenv
-> ALeftHandSide repr aenv aenv'
-> Dot (Aval aenv', Label)
prettyLetALeftHandSide NodeId
_     Aval aenv
aenv (LeftHandSideWildcard TupR ArrayR repr
repr) = (Aval aenv, Label) -> StateT DotState IO (Aval aenv, Label)
forall (m :: * -> *) a. Monad m => a -> m a
return (Aval aenv
aenv, Label
doc)
  where
    doc :: Label
doc = case TupR ArrayR repr
repr of
      TupR ArrayR repr
TupRunit -> Label
"()"
      TupR ArrayR repr
_        -> Label
"_"
prettyLetALeftHandSide NodeId
ident Aval aenv
aenv (LeftHandSideSingle ArrayR repr
_) = do
  Label
a <- Dot Label
mkLabel
  (Aval (aenv, repr), Label)
-> StateT DotState IO (Aval (aenv, repr), Label)
forall (m :: * -> *) a. Monad m => a -> m a
return (Aval aenv -> NodeId -> Label -> Aval (aenv, repr)
forall env t. Aval env -> NodeId -> Label -> Aval (env, t)
Apush Aval aenv
aenv NodeId
ident Label
a, Label
a)
prettyLetALeftHandSide NodeId
ident Aval aenv
aenv (LeftHandSidePair LeftHandSide ArrayR v1 aenv env'
lhs1 LeftHandSide ArrayR v2 env' aenv'
lhs2) = do
  (Aval env'
aenv1, Label
d1) <- NodeId
-> Aval aenv
-> LeftHandSide ArrayR v1 aenv env'
-> Dot (Aval env', Label)
forall repr aenv aenv'.
HasCallStack =>
NodeId
-> Aval aenv
-> ALeftHandSide repr aenv aenv'
-> Dot (Aval aenv', Label)
prettyLetALeftHandSide NodeId
ident Aval aenv
aenv  LeftHandSide ArrayR v1 aenv env'
lhs1
  (Aval aenv'
aenv2, Label
d2) <- NodeId
-> Aval env'
-> LeftHandSide ArrayR v2 env' aenv'
-> Dot (Aval aenv', Label)
forall repr aenv aenv'.
HasCallStack =>
NodeId
-> Aval aenv
-> ALeftHandSide repr aenv aenv'
-> Dot (Aval aenv', Label)
prettyLetALeftHandSide NodeId
ident Aval env'
aenv1 LeftHandSide ArrayR v2 env' aenv'
lhs2
  (Aval aenv', Label) -> Dot (Aval aenv', Label)
forall (m :: * -> *) a. Monad m => a -> m a
return (Aval aenv'
aenv2, Label
"(" Label -> Label -> Label
forall a. Semigroup a => a -> a -> a
<> Label
d1 Label -> Label -> Label
forall a. Semigroup a => a -> a -> a
<> Label
", " Label -> Label -> Label
forall a. Semigroup a => a -> a -> a
<> Label
d2 Label -> Label -> Label
forall a. Semigroup a => a -> a -> a
<> Label
")")

prettyLambdaALeftHandSide
    :: forall repr aenv aenv'. HasCallStack
    => Aval aenv
    -> ALeftHandSide repr aenv aenv'
    -> Dot (Aval aenv')
prettyLambdaALeftHandSide :: Aval aenv -> ALeftHandSide repr aenv aenv' -> Dot (Aval aenv')
prettyLambdaALeftHandSide Aval aenv
aenv (LeftHandSideWildcard TupR ArrayR repr
_) = Aval aenv -> StateT DotState IO (Aval aenv)
forall (m :: * -> *) a. Monad m => a -> m a
return Aval aenv
aenv
prettyLambdaALeftHandSide Aval aenv
aenv lhs :: ALeftHandSide repr aenv aenv'
lhs@(LeftHandSideSingle ArrayR repr
_) = do
  Label
a     <- Dot Label
mkLabel
  NodeId
ident <- ALeftHandSide repr aenv aenv' -> Dot NodeId
forall a. a -> Dot NodeId
mkNodeId ALeftHandSide repr aenv aenv'
lhs
  NodeId
_     <- PNode -> Maybe Label -> Dot NodeId
mkNode (NodeId
-> Tree (Maybe Label, Adoc) -> [(Vertex, Maybe Label)] -> PNode
PNode NodeId
ident ((Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a. a -> Tree a
Leaf (Maybe Label
forall a. Maybe a
Nothing, Label -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Label
a)) []) Maybe Label
forall a. Maybe a
Nothing
  Aval (aenv, repr) -> StateT DotState IO (Aval (aenv, repr))
forall (m :: * -> *) a. Monad m => a -> m a
return (Aval (aenv, repr) -> StateT DotState IO (Aval (aenv, repr)))
-> Aval (aenv, repr) -> StateT DotState IO (Aval (aenv, repr))
forall a b. (a -> b) -> a -> b
$ Aval aenv -> NodeId -> Label -> Aval (aenv, repr)
forall env t. Aval env -> NodeId -> Label -> Aval (env, t)
Apush Aval aenv
aenv NodeId
ident Label
a
prettyLambdaALeftHandSide Aval aenv
aenv (LeftHandSidePair LeftHandSide ArrayR v1 aenv env'
lhs1 LeftHandSide ArrayR v2 env' aenv'
lhs2) = do
  Aval env'
aenv1 <- Aval aenv -> LeftHandSide ArrayR v1 aenv env' -> Dot (Aval env')
forall repr aenv aenv'.
HasCallStack =>
Aval aenv -> ALeftHandSide repr aenv aenv' -> Dot (Aval aenv')
prettyLambdaALeftHandSide Aval aenv
aenv LeftHandSide ArrayR v1 aenv env'
lhs1
  Aval env' -> LeftHandSide ArrayR v2 env' aenv' -> Dot (Aval aenv')
forall repr aenv aenv'.
HasCallStack =>
Aval aenv -> ALeftHandSide repr aenv aenv' -> Dot (Aval aenv')
prettyLambdaALeftHandSide Aval env'
aenv1 LeftHandSide ArrayR v2 env' aenv'
lhs2

-- Display array tuples. This is a little tricky...
--
prettyDelayedApair
    :: forall aenv a1 a2. HasCallStack
    => Detail
    -> Aval aenv
    -> DelayedOpenAcc aenv a1
    -> DelayedOpenAcc aenv a2
    -> NodeId
    -> Dot PNode
prettyDelayedApair :: Detail
-> Aval aenv
-> DelayedOpenAcc aenv a1
-> DelayedOpenAcc aenv a2
-> NodeId
-> Dot PNode
prettyDelayedApair Detail
detail Aval aenv
aenv DelayedOpenAcc aenv a1
a1 DelayedOpenAcc aenv a2
a2 NodeId
ident = do
  PNode NodeId
id1 Tree (Maybe Label, Adoc)
t1 [(Vertex, Maybe Label)]
v1 <- DelayedOpenAcc aenv a1 -> Dot PNode
forall a. DelayedOpenAcc aenv a -> Dot PNode
prettyElem DelayedOpenAcc aenv a1
a1
  PNode NodeId
id2 Tree (Maybe Label, Adoc)
t2 [(Vertex, Maybe Label)]
v2 <- DelayedOpenAcc aenv a2 -> Dot PNode
forall a. DelayedOpenAcc aenv a -> Dot PNode
prettyElem DelayedOpenAcc aenv a2
a2
  (DotState -> DotState) -> StateT DotState IO ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((DotState -> DotState) -> StateT DotState IO ())
-> (DotState -> DotState) -> StateT DotState IO ()
forall a b. (a -> b) -> a -> b
$ \DotState
s -> DotState
s { dotEdges :: Seq Edge
dotEdges = (Edge -> Edge) -> Seq Edge -> Seq Edge
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (NodeId -> [NodeId] -> Edge -> Edge
redirect NodeId
ident [NodeId
id1, NodeId
id2]) (DotState -> Seq Edge
dotEdges DotState
s) }
  PNode -> Dot PNode
forall (m :: * -> *) a. Monad m => a -> m a
return (PNode -> Dot PNode) -> PNode -> Dot PNode
forall a b. (a -> b) -> a -> b
$ NodeId
-> Tree (Maybe Label, Adoc) -> [(Vertex, Maybe Label)] -> PNode
PNode NodeId
ident ([Tree (Maybe Label, Adoc)] -> Tree (Maybe Label, Adoc)
forest [Tree (Maybe Label, Adoc)
t1, Tree (Maybe Label, Adoc)
t2]) ([(Vertex, Maybe Label)]
v1 [(Vertex, Maybe Label)]
-> [(Vertex, Maybe Label)] -> [(Vertex, Maybe Label)]
forall a. [a] -> [a] -> [a]
++ [(Vertex, Maybe Label)]
v2)
  where
    prettyElem :: DelayedOpenAcc aenv a -> Dot PNode
    prettyElem :: DelayedOpenAcc aenv a -> Dot PNode
prettyElem DelayedOpenAcc aenv a
a = PNode -> Dot PNode
replant (PNode -> Dot PNode) -> Dot PNode -> Dot PNode
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv a -> Dot PNode
forall aenv arrs.
HasCallStack =>
Detail
-> Context -> Aval aenv -> DelayedOpenAcc aenv arrs -> Dot PNode
prettyDelayedOpenAcc Detail
detail Context
context0 Aval aenv
aenv DelayedOpenAcc aenv a
a

    -- Redirect any edges that pointed into one of the nodes now part of this
    -- tuple, to instead point to the container node.
    --
    redirect :: NodeId -> [NodeId] -> Edge -> Edge
    redirect :: NodeId -> [NodeId] -> Edge -> Edge
redirect NodeId
new [NodeId]
subs edge :: Edge
edge@(Edge Vertex
from (Vertex NodeId
to Maybe Label
port))
      | NodeId
to NodeId -> [NodeId] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [NodeId]
subs = Vertex -> Vertex -> Edge
Edge Vertex
from (NodeId -> Maybe Label -> Vertex
Vertex NodeId
new Maybe Label
port)
      | Bool
otherwise      = Edge
edge

    -- Since we have lifted out any non-leaves into separate nodes, we can
    -- simply tuple-up all of the elements.
    --
    forest :: [Tree (Maybe Port, Adoc)] -> Tree (Maybe Port, Adoc)
    forest :: [Tree (Maybe Label, Adoc)] -> Tree (Maybe Label, Adoc)
forest [Tree (Maybe Label, Adoc)]
leaves = (Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a. a -> Tree a
Leaf (Maybe Label
forall a. Maybe a
Nothing, [Adoc] -> Adoc
forall ann. [Doc ann] -> Doc ann
tupled [ Adoc -> Adoc
forall ann. Doc ann -> Doc ann
align Adoc
d | Leaf (Maybe Label
Nothing,Adoc
d) <- [Tree (Maybe Label, Adoc)]
leaves ])


-- Lift out anything that isn't a Leaf node and output it to the graph
-- immediately as a new labelled node.
--
replant :: PNode -> Dot PNode
replant :: PNode -> Dot PNode
replant pnode :: PNode
pnode@(PNode NodeId
ident Tree (Maybe Label, Adoc)
tree [(Vertex, Maybe Label)]
_) =
  case Tree (Maybe Label, Adoc)
tree of
    Leaf (Maybe Label
Nothing, Adoc
_) -> PNode -> Dot PNode
forall (m :: * -> *) a. Monad m => a -> m a
return PNode
pnode
    Tree (Maybe Label, Adoc)
_                 -> do
      NodeId
vacuous <- PNode -> Dot NodeId
forall a. a -> Dot NodeId
mkNodeId PNode
pnode
      Label
a       <- Dot Label
mkLabel
      NodeId
_       <- PNode -> Maybe Label -> Dot NodeId
mkNode PNode
pnode (Label -> Maybe Label
forall a. a -> Maybe a
Just Label
a)
      PNode -> Dot PNode
forall (m :: * -> *) a. Monad m => a -> m a
return   (PNode -> Dot PNode) -> PNode -> Dot PNode
forall a b. (a -> b) -> a -> b
$ NodeId
-> Tree (Maybe Label, Adoc) -> [(Vertex, Maybe Label)] -> PNode
PNode NodeId
vacuous ((Maybe Label, Adoc) -> Tree (Maybe Label, Adoc)
forall a. a -> Tree a
Leaf (Maybe Label
forall a. Maybe a
Nothing, Label -> Adoc
forall a ann. Pretty a => a -> Doc ann
pretty Label
a)) [(NodeId -> Maybe Label -> Vertex
Vertex NodeId
ident Maybe Label
forall a. Maybe a
Nothing, Maybe Label
forall a. Maybe a
Nothing)]


-- Pretty printing scalar functions and expressions
-- ------------------------------------------------
--
-- This is done with the usual machinery. Note that we rely on knowing that all
-- array operations will be lifted out of scalar expressions. This means that we
-- don't really need to recurse into the scalar terms to uncover new graph
-- nodes.
--

-- Data dependencies
-- -----------------
--
-- Return the data-dependencies of the given term. This is just a tree traversal
-- to extract all of the free variables. We will draw an edge from each of those
-- nodes (vertices) into the current term.
--

fvAvar :: Aval aenv -> ArrayVar aenv a -> [Vertex]
fvAvar :: Aval aenv -> ArrayVar aenv a -> [Vertex]
fvAvar Aval aenv
env (Var ArrayR a
_ Idx aenv a
ix) = [ NodeId -> Maybe Label -> Vertex
Vertex ((NodeId, Label) -> NodeId
forall a b. (a, b) -> a
fst ((NodeId, Label) -> NodeId) -> (NodeId, Label) -> NodeId
forall a b. (a -> b) -> a -> b
$ Idx aenv a -> Aval aenv -> (NodeId, Label)
forall aenv t. Idx aenv t -> Aval aenv -> (NodeId, Label)
aprj Idx aenv a
ix Aval aenv
env) Maybe Label
forall a. Maybe a
Nothing ]

fvOpenFun
    :: forall env aenv fun.
       Val env
    -> Aval aenv
    -> OpenFun env aenv fun
    -> [Vertex]
fvOpenFun :: Val env -> Aval aenv -> OpenFun env aenv fun -> [Vertex]
fvOpenFun Val env
env Aval aenv
aenv (Body OpenExp env aenv fun
b)    = Val env -> Aval aenv -> OpenExp env aenv fun -> [Vertex]
forall env aenv exp.
Val env -> Aval aenv -> OpenExp env aenv exp -> [Vertex]
fvOpenExp Val env
env  Aval aenv
aenv OpenExp env aenv fun
b
fvOpenFun Val env
env Aval aenv
aenv (Lam ELeftHandSide a env env'
lhs OpenFun env' aenv t
f) = Val env' -> Aval aenv -> OpenFun env' aenv t -> [Vertex]
forall env aenv fun.
Val env -> Aval aenv -> OpenFun env aenv fun -> [Vertex]
fvOpenFun Val env'
env' Aval aenv
aenv OpenFun env' aenv t
f
      where
        (Val env'
env', Adoc
_) = Bool -> Val env -> ELeftHandSide a env env' -> (Val env', Adoc)
forall env (s :: * -> *) arrs env'.
Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc)
prettyELhs Bool
True Val env
env ELeftHandSide a env env'
lhs

fvOpenExp
    :: forall env aenv exp.
       Val env
    -> Aval aenv
    -> OpenExp env aenv exp
    -> [Vertex]
fvOpenExp :: Val env -> Aval aenv -> OpenExp env aenv exp -> [Vertex]
fvOpenExp Val env
env Aval aenv
aenv = OpenExp env aenv exp -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv
  where
    fvF :: OpenFun env aenv f -> [Vertex]
    fvF :: OpenFun env aenv f -> [Vertex]
fvF = Val env -> Aval aenv -> OpenFun env aenv f -> [Vertex]
forall env aenv fun.
Val env -> Aval aenv -> OpenFun env aenv fun -> [Vertex]
fvOpenFun Val env
env Aval aenv
aenv

    fv :: OpenExp env aenv e -> [Vertex]
    fv :: OpenExp env aenv e -> [Vertex]
fv (Shape ArrayVar aenv (Array e e)
acc)              = if Bool
cfgIncludeShape then Aval aenv -> ArrayVar aenv (Array e e) -> [Vertex]
forall aenv a. Aval aenv -> ArrayVar aenv a -> [Vertex]
fvAvar Aval aenv
aenv ArrayVar aenv (Array e e)
acc else []
    fv (Index ArrayVar aenv (Array dim e)
acc OpenExp env aenv dim
i)            = [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ Aval aenv -> ArrayVar aenv (Array dim e) -> [Vertex]
forall aenv a. Aval aenv -> ArrayVar aenv a -> [Vertex]
fvAvar Aval aenv
aenv ArrayVar aenv (Array dim e)
acc, OpenExp env aenv dim -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv dim
i ]
    fv (LinearIndex ArrayVar aenv (Array dim e)
acc OpenExp env aenv Int
i)      = [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ Aval aenv -> ArrayVar aenv (Array dim e) -> [Vertex]
forall aenv a. Aval aenv -> ArrayVar aenv a -> [Vertex]
fvAvar Aval aenv
aenv ArrayVar aenv (Array dim e)
acc, OpenExp env aenv Int -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv Int
i ]
    --
    fv (Let ELeftHandSide bnd_t env env'
lhs OpenExp env aenv bnd_t
e1 OpenExp env' aenv e
e2)          = [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ OpenExp env aenv bnd_t -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv bnd_t
e1, Val env' -> Aval aenv -> OpenExp env' aenv e -> [Vertex]
forall env aenv exp.
Val env -> Aval aenv -> OpenExp env aenv exp -> [Vertex]
fvOpenExp Val env'
env' Aval aenv
aenv OpenExp env' aenv e
e2 ]
      where
        (Val env'
env', Adoc
_) = Bool -> Val env -> ELeftHandSide bnd_t env env' -> (Val env', Adoc)
forall env (s :: * -> *) arrs env'.
Bool -> Val env -> LeftHandSide s arrs env env' -> (Val env', Adoc)
prettyELhs Bool
False Val env
env ELeftHandSide bnd_t env env'
lhs
    fv Evar{}                   = []
    fv Undef{}                  = []
    fv Const{}                  = []
    fv PrimConst{}              = []
    fv (PrimApp PrimFun (a -> e)
_ OpenExp env aenv a
x)            = OpenExp env aenv a -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv a
x
    fv (Pair OpenExp env aenv t1
e1 OpenExp env aenv t2
e2)             = [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ OpenExp env aenv t1 -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv t1
e1, OpenExp env aenv t2 -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv t2
e2]
    fv OpenExp env aenv e
Nil                      = []
    fv (VecPack   VecR n s tup
_ OpenExp env aenv tup
e)          = OpenExp env aenv tup -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv tup
e
    fv (VecUnpack VecR n s e
_ OpenExp env aenv (Vec n s)
e)          = OpenExp env aenv (Vec n s) -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv (Vec n s)
e
    fv (IndexSlice SliceIndex slix e co sh
_ OpenExp env aenv slix
slix OpenExp env aenv sh
sh)   = [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ OpenExp env aenv slix -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv slix
slix, OpenExp env aenv sh -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv sh
sh ]
    fv (IndexFull SliceIndex slix sl co e
_ OpenExp env aenv slix
slix OpenExp env aenv sl
sh)    = [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ OpenExp env aenv slix -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv slix
slix, OpenExp env aenv sl -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv sl
sh ]
    fv (ToIndex ShapeR sh
_ OpenExp env aenv sh
sh OpenExp env aenv sh
ix)        = [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ OpenExp env aenv sh -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv sh
sh, OpenExp env aenv sh -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv sh
ix ]
    fv (FromIndex ShapeR e
_ OpenExp env aenv e
sh OpenExp env aenv Int
ix)      = [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ OpenExp env aenv e -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv e
sh, OpenExp env aenv Int -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv Int
ix ]
    fv (ShapeSize ShapeR dim
_ OpenExp env aenv dim
sh)         = OpenExp env aenv dim -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv dim
sh
    fv Foreign{}                = []
    fv (Case OpenExp env aenv PrimBool
e [(PrimBool, OpenExp env aenv e)]
rhs Maybe (OpenExp env aenv e)
def)         = [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ OpenExp env aenv PrimBool -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv PrimBool
e, [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ OpenExp env aenv e -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv e
c | (PrimBool
_,OpenExp env aenv e
c) <- [(PrimBool, OpenExp env aenv e)]
rhs ], [Vertex]
-> (OpenExp env aenv e -> [Vertex])
-> Maybe (OpenExp env aenv e)
-> [Vertex]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] OpenExp env aenv e -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv Maybe (OpenExp env aenv e)
def ]
    fv (Cond OpenExp env aenv PrimBool
p OpenExp env aenv e
t OpenExp env aenv e
e)             = [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ OpenExp env aenv PrimBool -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv PrimBool
p, OpenExp env aenv e -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv e
t, OpenExp env aenv e -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv e
e ]
    fv (While OpenFun env aenv (e -> PrimBool)
p OpenFun env aenv (e -> e)
f OpenExp env aenv e
x)            = [[Vertex]] -> [Vertex]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ OpenFun env aenv (e -> PrimBool) -> [Vertex]
forall f. OpenFun env aenv f -> [Vertex]
fvF OpenFun env aenv (e -> PrimBool)
p, OpenFun env aenv (e -> e) -> [Vertex]
forall f. OpenFun env aenv f -> [Vertex]
fvF OpenFun env aenv (e -> e)
f, OpenExp env aenv e -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv e
x ]
    fv (Coerce ScalarType a
_ ScalarType e
_ OpenExp env aenv a
e)           = OpenExp env aenv a -> [Vertex]
forall e. OpenExp env aenv e -> [Vertex]
fv OpenExp env aenv a
e