{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
module GHC.Cmm.Dataflow.Graph
    ( Body
    , Graph
    , Graph'(..)
    , NonLocal(..)
    , addBlock
    , bodyList
    , bodyToBlockList
    , emptyBody
    , labelsDefined
    , mapGraph
    , mapGraphBlocks
    , revPostorderFrom
    ) where


import GHC.Prelude
import GHC.Utils.Misc

import GHC.Cmm.Dataflow.Label
import GHC.Cmm.Dataflow.Block
import GHC.Cmm.Dataflow.Collections

import Data.Kind

-- | A (possibly empty) collection of closed/closed blocks
type Body n = LabelMap (Block n C C)

-- | @Body@ abstracted over @block@
type Body' block (n :: Extensibility -> Extensibility -> Type) = LabelMap (block n C C)

-------------------------------
-- | Gives access to the anchor points for
-- nonlocal edges as well as the edges themselves
class NonLocal thing where
  entryLabel :: thing C x -> Label   -- ^ The label of a first node or block
  successors :: thing e C -> [Label] -- ^ Gives control-flow successors

instance NonLocal n => NonLocal (Block n) where
  entryLabel :: forall (x :: Extensibility). Block n C x -> Label
entryLabel (BlockCO n C 'Open
f Block n 'Open 'Open
_)   = forall (thing :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
NonLocal thing =>
thing C x -> Label
entryLabel n C 'Open
f
  entryLabel (BlockCC n C 'Open
f Block n 'Open 'Open
_ n 'Open C
_) = forall (thing :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
NonLocal thing =>
thing C x -> Label
entryLabel n C 'Open
f

  successors :: forall (e :: Extensibility). Block n e C -> [Label]
successors (BlockOC   Block n 'Open 'Open
_ n 'Open C
n) = forall (thing :: Extensibility -> Extensibility -> *)
       (e :: Extensibility).
NonLocal thing =>
thing e C -> [Label]
successors n 'Open C
n
  successors (BlockCC n C 'Open
_ Block n 'Open 'Open
_ n 'Open C
n) = forall (thing :: Extensibility -> Extensibility -> *)
       (e :: Extensibility).
NonLocal thing =>
thing e C -> [Label]
successors n 'Open C
n


emptyBody :: Body' block n
emptyBody :: forall (block :: (Extensibility -> Extensibility -> *)
                 -> Extensibility -> Extensibility -> *)
       (n :: Extensibility -> Extensibility -> *).
Body' block n
emptyBody = forall (map :: * -> *) a. IsMap map => map a
mapEmpty

bodyList :: Body' block n -> [(Label,block n C C)]
bodyList :: forall (block :: (Extensibility -> Extensibility -> *)
                 -> Extensibility -> Extensibility -> *)
       (n :: Extensibility -> Extensibility -> *).
Body' block n -> [(Label, block n C C)]
bodyList Body' block n
body = forall (map :: * -> *) a. IsMap map => map a -> [(KeyOf map, a)]
mapToList Body' block n
body

bodyToBlockList :: Body n -> [Block n C C]
bodyToBlockList :: forall (n :: Extensibility -> Extensibility -> *).
Body n -> [Block n C C]
bodyToBlockList Body n
body = forall (map :: * -> *) a. IsMap map => map a -> [a]
mapElems Body n
body

addBlock
    :: (NonLocal block, HasDebugCallStack)
    => block C C -> LabelMap (block C C) -> LabelMap (block C C)
addBlock :: forall (block :: Extensibility -> Extensibility -> *).
(NonLocal block, HasDebugCallStack) =>
block C C -> LabelMap (block C C) -> LabelMap (block C C)
addBlock block C C
block LabelMap (block C C)
body = forall (map :: * -> *) a.
IsMap map =>
(Maybe a -> Maybe a) -> KeyOf map -> map a -> map a
mapAlter Maybe (block C C) -> Maybe (block C C)
add Label
lbl LabelMap (block C C)
body
  where
    lbl :: Label
lbl = forall (thing :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
NonLocal thing =>
thing C x -> Label
entryLabel block C C
block
    add :: Maybe (block C C) -> Maybe (block C C)
add Maybe (block C C)
Nothing = forall a. a -> Maybe a
Just block C C
block
    add Maybe (block C C)
_ = forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"duplicate label " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Label
lbl forall a. [a] -> [a] -> [a]
++ [Char]
" in graph"


-- ---------------------------------------------------------------------------
-- Graph

-- | A control-flow graph, which may take any of four shapes (O/O,
-- O/C, C/O, C/C).  A graph open at the entry has a single,
-- distinguished, anonymous entry point; if a graph is closed at the
-- entry, its entry point(s) are supplied by a context.
type Graph = Graph' Block

-- | @Graph'@ is abstracted over the block type, so that we can build
-- graphs of annotated blocks for example (Compiler.Hoopl.Dataflow
-- needs this).
data Graph' block (n :: Extensibility -> Extensibility -> Type) e x where
  GNil  :: Graph' block n O O
  GUnit :: block n O O -> Graph' block n O O
  GMany :: MaybeO e (block n O C)
        -> Body' block n
        -> MaybeO x (block n C O)
        -> Graph' block n e x


-- -----------------------------------------------------------------------------
-- Mapping over graphs

-- | Maps over all nodes in a graph.
mapGraph :: (forall e x. n e x -> n' e x) -> Graph n e x -> Graph n' e x
mapGraph :: forall (n :: Extensibility -> Extensibility -> *)
       (n' :: Extensibility -> Extensibility -> *) (e :: Extensibility)
       (x :: Extensibility).
(forall (e :: Extensibility) (x :: Extensibility). n e x -> n' e x)
-> Graph n e x -> Graph n' e x
mapGraph forall (e :: Extensibility) (x :: Extensibility). n e x -> n' e x
f = forall (block :: (Extensibility -> Extensibility -> *)
                 -> Extensibility -> Extensibility -> *)
       (n :: Extensibility -> Extensibility -> *)
       (block' :: (Extensibility -> Extensibility -> *)
                  -> Extensibility -> Extensibility -> *)
       (n' :: Extensibility -> Extensibility -> *) (e :: Extensibility)
       (x :: Extensibility).
(forall (e :: Extensibility) (x :: Extensibility).
 block n e x -> block' n' e x)
-> Graph' block n e x -> Graph' block' n' e x
mapGraphBlocks (forall (n :: Extensibility -> Extensibility -> *)
       (n' :: Extensibility -> Extensibility -> *) (e :: Extensibility)
       (x :: Extensibility).
(forall (e1 :: Extensibility) (x1 :: Extensibility).
 n e1 x1 -> n' e1 x1)
-> Block n e x -> Block n' e x
mapBlock forall (e :: Extensibility) (x :: Extensibility). n e x -> n' e x
f)

-- | Function 'mapGraphBlocks' enables a change of representation of blocks,
-- nodes, or both.  It lifts a polymorphic block transform into a polymorphic
-- graph transform.  When the block representation stabilizes, a similar
-- function should be provided for blocks.
mapGraphBlocks :: forall block n block' n' e x .
                  (forall e x . block n e x -> block' n' e x)
               -> (Graph' block n e x -> Graph' block' n' e x)

mapGraphBlocks :: forall (block :: (Extensibility -> Extensibility -> *)
                 -> Extensibility -> Extensibility -> *)
       (n :: Extensibility -> Extensibility -> *)
       (block' :: (Extensibility -> Extensibility -> *)
                  -> Extensibility -> Extensibility -> *)
       (n' :: Extensibility -> Extensibility -> *) (e :: Extensibility)
       (x :: Extensibility).
(forall (e :: Extensibility) (x :: Extensibility).
 block n e x -> block' n' e x)
-> Graph' block n e x -> Graph' block' n' e x
mapGraphBlocks forall (e :: Extensibility) (x :: Extensibility).
block n e x -> block' n' e x
f = Graph' block n e x -> Graph' block' n' e x
map
  where map :: Graph' block n e x -> Graph' block' n' e x
        map :: Graph' block n e x -> Graph' block' n' e x
map Graph' block n e x
GNil = forall (block :: (Extensibility -> Extensibility -> *)
                 -> Extensibility -> Extensibility -> *)
       (n :: Extensibility -> Extensibility -> *).
Graph' block n 'Open 'Open
GNil
        map (GUnit block n 'Open 'Open
b) = forall (block :: (Extensibility -> Extensibility -> *)
                 -> Extensibility -> Extensibility -> *)
       (n :: Extensibility -> Extensibility -> *).
block n 'Open 'Open -> Graph' block n 'Open 'Open
GUnit (forall (e :: Extensibility) (x :: Extensibility).
block n e x -> block' n' e x
f block n 'Open 'Open
b)
        map (GMany MaybeO e (block n 'Open C)
e Body' block n
b MaybeO x (block n C 'Open)
x) = forall (e :: Extensibility)
       (block :: (Extensibility -> Extensibility -> *)
                 -> Extensibility -> Extensibility -> *)
       (n :: Extensibility -> Extensibility -> *) (x :: Extensibility).
MaybeO e (block n 'Open C)
-> Body' block n
-> MaybeO x (block n C 'Open)
-> Graph' block n e x
GMany (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (e :: Extensibility) (x :: Extensibility).
block n e x -> block' n' e x
f MaybeO e (block n 'Open C)
e) (forall (map :: * -> *) a b. IsMap map => (a -> b) -> map a -> map b
mapMap forall (e :: Extensibility) (x :: Extensibility).
block n e x -> block' n' e x
f Body' block n
b) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (e :: Extensibility) (x :: Extensibility).
block n e x -> block' n' e x
f MaybeO x (block n C 'Open)
x)

-- -----------------------------------------------------------------------------
-- Extracting Labels from graphs

labelsDefined :: forall block n e x . NonLocal (block n) => Graph' block n e x
              -> LabelSet
labelsDefined :: forall (block :: (Extensibility -> Extensibility -> *)
                 -> Extensibility -> Extensibility -> *)
       (n :: Extensibility -> Extensibility -> *) (e :: Extensibility)
       (x :: Extensibility).
NonLocal (block n) =>
Graph' block n e x -> LabelSet
labelsDefined Graph' block n e x
GNil      = forall set. IsSet set => set
setEmpty
labelsDefined (GUnit{}) = forall set. IsSet set => set
setEmpty
labelsDefined (GMany MaybeO e (block n 'Open C)
_ Body' block n
body MaybeO x (block n C 'Open)
x) = forall (map :: * -> *) b a.
IsMap map =>
(b -> KeyOf map -> a -> b) -> b -> map a -> b
mapFoldlWithKey forall a. LabelSet -> ElemOf LabelSet -> a -> LabelSet
addEntry (MaybeO x (block n C 'Open) -> LabelSet
exitLabel MaybeO x (block n C 'Open)
x) Body' block n
body
  where addEntry :: forall a. LabelSet -> ElemOf LabelSet -> a -> LabelSet
        addEntry :: forall a. LabelSet -> ElemOf LabelSet -> a -> LabelSet
addEntry LabelSet
labels ElemOf LabelSet
label a
_ = forall set. IsSet set => ElemOf set -> set -> set
setInsert ElemOf LabelSet
label LabelSet
labels
        exitLabel :: MaybeO x (block n C O) -> LabelSet
        exitLabel :: MaybeO x (block n C 'Open) -> LabelSet
exitLabel MaybeO x (block n C 'Open)
NothingO  = forall set. IsSet set => set
setEmpty
        exitLabel (JustO block n C 'Open
b) = forall set. IsSet set => ElemOf set -> set
setSingleton (forall (thing :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
NonLocal thing =>
thing C x -> Label
entryLabel block n C 'Open
b)


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

-- | Returns a list of blocks reachable from the provided Labels in the reverse
-- postorder.
--
-- This is the most important traversal over this data structure.  It drops
-- unreachable code and puts blocks in an order that is good for solving forward
-- dataflow problems quickly.  The reverse order is good for solving backward
-- dataflow problems quickly.  The forward order is also reasonably good for
-- emitting instructions, except that it will not usually exploit Forrest
-- Baskett's trick of eliminating the unconditional branch from a loop.  For
-- that you would need a more serious analysis, probably based on dominators, to
-- identify loop headers.
--
-- For forward analyses we want reverse postorder visitation, consider:
-- @
--      A -> [B,C]
--      B -> D
--      C -> D
-- @
-- Postorder: [D, C, B, A] (or [D, B, C, A])
-- Reverse postorder: [A, B, C, D] (or [A, C, B, D])
-- This matters for, e.g., forward analysis, because we want to analyze *both*
-- B and C before we analyze D.
revPostorderFrom
  :: forall block.  (NonLocal block)
  => LabelMap (block C C) -> Label -> [block C C]
revPostorderFrom :: forall (block :: Extensibility -> Extensibility -> *).
NonLocal block =>
LabelMap (block C C) -> Label -> [block C C]
revPostorderFrom LabelMap (block C C)
graph Label
start = DfsStack (block C C) -> LabelSet -> [block C C] -> [block C C]
go DfsStack (block C C)
start_worklist forall set. IsSet set => set
setEmpty []
  where
    start_worklist :: DfsStack (block C C)
start_worklist = Label -> DfsStack (block C C) -> DfsStack (block C C)
lookup_for_descend Label
start forall a. DfsStack a
Nil

    -- To compute the postorder we need to "visit" a block (mark as done) *after*
    -- visiting all its successors. So we need to know whether we already
    -- processed all successors of each block (and @NonLocal@ allows arbitrary
    -- many successors). So we use an explicit stack with an extra bit
    -- of information:
    -- - @ConsTodo@ means to explore the block if it wasn't visited before
    -- - @ConsMark@ means that all successors were already done and we can add
    --   the block to the result.
    --
    -- NOTE: We add blocks to the result list in postorder, but we *prepend*
    -- them (i.e., we use @(:)@), which means that the final list is in reverse
    -- postorder.
    go :: DfsStack (block C C) -> LabelSet -> [block C C] -> [block C C]
    go :: DfsStack (block C C) -> LabelSet -> [block C C] -> [block C C]
go DfsStack (block C C)
Nil                      !LabelSet
_           ![block C C]
result = [block C C]
result
    go (ConsMark block C C
block DfsStack (block C C)
rest)    !LabelSet
wip_or_done ![block C C]
result =
        DfsStack (block C C) -> LabelSet -> [block C C] -> [block C C]
go DfsStack (block C C)
rest LabelSet
wip_or_done (block C C
block forall a. a -> [a] -> [a]
: [block C C]
result)
    go (ConsTodo block C C
block DfsStack (block C C)
rest)    !LabelSet
wip_or_done ![block C C]
result
        | forall (thing :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
NonLocal thing =>
thing C x -> Label
entryLabel block C C
block forall set. IsSet set => ElemOf set -> set -> Bool
`setMember` LabelSet
wip_or_done = DfsStack (block C C) -> LabelSet -> [block C C] -> [block C C]
go DfsStack (block C C)
rest LabelSet
wip_or_done [block C C]
result
        | Bool
otherwise =
            let new_worklist :: DfsStack (block C C)
new_worklist =
                    forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Label -> DfsStack (block C C) -> DfsStack (block C C)
lookup_for_descend
                          (forall a. a -> DfsStack a -> DfsStack a
ConsMark block C C
block DfsStack (block C C)
rest)
                          (forall (thing :: Extensibility -> Extensibility -> *)
       (e :: Extensibility).
NonLocal thing =>
thing e C -> [Label]
successors block C C
block)
            in DfsStack (block C C) -> LabelSet -> [block C C] -> [block C C]
go DfsStack (block C C)
new_worklist (forall set. IsSet set => ElemOf set -> set -> set
setInsert (forall (thing :: Extensibility -> Extensibility -> *)
       (x :: Extensibility).
NonLocal thing =>
thing C x -> Label
entryLabel block C C
block) LabelSet
wip_or_done) [block C C]
result

    lookup_for_descend :: Label -> DfsStack (block C C) -> DfsStack (block C C)
    lookup_for_descend :: Label -> DfsStack (block C C) -> DfsStack (block C C)
lookup_for_descend Label
label DfsStack (block C C)
wl
      | Just block C C
b <- forall (map :: * -> *) a.
IsMap map =>
KeyOf map -> map a -> Maybe a
mapLookup Label
label LabelMap (block C C)
graph = forall a. a -> DfsStack a -> DfsStack a
ConsTodo block C C
b DfsStack (block C C)
wl
      | Bool
otherwise =
           forall a. HasCallStack => [Char] -> a
error forall a b. (a -> b) -> a -> b
$ [Char]
"Label that doesn't have a block?! " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Label
label

data DfsStack a = ConsTodo a (DfsStack a) | ConsMark a (DfsStack a) | Nil