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


import GhcPrelude
import Util

import Hoopl.Label
import Hoopl.Block
import Hoopl.Collections

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

-- | @Body@ abstracted over @block@
type Body' block (n :: * -> * -> *) = 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 :: Block n C x -> Label
entryLabel (BlockCO f :: n C O
f _)   = n C O -> Label
forall (thing :: * -> * -> *) x.
NonLocal thing =>
thing C x -> Label
entryLabel n C O
f
  entryLabel (BlockCC f :: n C O
f _ _) = n C O -> Label
forall (thing :: * -> * -> *) x.
NonLocal thing =>
thing C x -> Label
entryLabel n C O
f

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


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

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

addBlock
    :: (NonLocal block, HasDebugCallStack)
    => block C C -> LabelMap (block C C) -> LabelMap (block C C)
addBlock :: block C C -> LabelMap (block C C) -> LabelMap (block C C)
addBlock block :: block C C
block body :: LabelMap (block C C)
body = (Maybe (block C C) -> Maybe (block C C))
-> KeyOf LabelMap -> LabelMap (block C C) -> LabelMap (block C C)
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 KeyOf LabelMap
Label
lbl LabelMap (block C C)
body
  where
    lbl :: Label
lbl = block C C -> Label
forall (thing :: * -> * -> *) x.
NonLocal thing =>
thing C x -> Label
entryLabel block C C
block
    add :: Maybe (block C C) -> Maybe (block C C)
add Nothing = block C C -> Maybe (block C C)
forall a. a -> Maybe a
Just block C C
block
    add _ = [Char] -> Maybe (block C C)
forall a. HasCallStack => [Char] -> a
error ([Char] -> Maybe (block C C)) -> [Char] -> Maybe (block C C)
forall a b. (a -> b) -> a -> b
$ "duplicate label " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Label -> [Char]
forall a. Show a => a -> [Char]
show Label
lbl [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ " 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 :: * -> * -> *) 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 e x. n e x -> n' e x) -> Graph n e x -> Graph n' e x
mapGraph f :: forall e x. n e x -> n' e x
f = (forall e x. Block n e x -> Block n' e x)
-> Graph n e x -> Graph n' e x
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 e x. n e x -> n' e x) -> Block n e x -> Block n' e x
forall (n :: * -> * -> *) (n' :: * -> * -> *) e x.
(forall e1 x1. n e1 x1 -> n' e1 x1) -> Block n e x -> Block n' e x
mapBlock forall e x. 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 e x. block n e x -> block' n' e x)
-> Graph' block n e x -> Graph' block' n' e x
mapGraphBlocks f :: forall e x. 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 GNil = Graph' block' n' e x
forall (block :: (* -> * -> *) -> * -> * -> *) (n :: * -> * -> *).
Graph' block n O O
GNil
        map (GUnit b :: block n O O
b) = block' n' O O -> Graph' block' n' O O
forall (block :: (* -> * -> *) -> * -> * -> *) (n :: * -> * -> *).
block n O O -> Graph' block n O O
GUnit (block n O O -> block' n' O O
forall e x. block n e x -> block' n' e x
f block n O O
b)
        map (GMany e :: MaybeO e (block n O C)
e b :: Body' block n
b x :: MaybeO x (block n C O)
x) = MaybeO e (block' n' O C)
-> Body' block' n'
-> MaybeO x (block' n' C O)
-> Graph' block' n' e x
forall e (block :: (* -> * -> *) -> * -> * -> *) (n :: * -> * -> *)
       x.
MaybeO e (block n O C)
-> Body' block n -> MaybeO x (block n C O) -> Graph' block n e x
GMany ((block n O C -> block' n' O C)
-> MaybeO e (block n O C) -> MaybeO e (block' n' O C)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap block n O C -> block' n' O C
forall e x. block n e x -> block' n' e x
f MaybeO e (block n O C)
e) ((block n C C -> block' n' C C) -> Body' block n -> Body' block' n'
forall (map :: * -> *) a b. IsMap map => (a -> b) -> map a -> map b
mapMap block n C C -> block' n' C C
forall e x. block n e x -> block' n' e x
f Body' block n
b) ((block n C O -> block' n' C O)
-> MaybeO x (block n C O) -> MaybeO x (block' n' C O)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap block n C O -> block' n' C O
forall e x. block n e x -> block' n' e x
f MaybeO x (block n C O)
x)

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

labelsDefined :: forall block n e x . NonLocal (block n) => Graph' block n e x
              -> LabelSet
labelsDefined :: Graph' block n e x -> LabelSet
labelsDefined GNil      = LabelSet
forall set. IsSet set => set
setEmpty
labelsDefined (GUnit{}) = LabelSet
forall set. IsSet set => set
setEmpty
labelsDefined (GMany _ body :: Body' block n
body x :: MaybeO x (block n C O)
x) = (LabelSet -> KeyOf LabelMap -> block n C C -> LabelSet)
-> LabelSet -> Body' block n -> LabelSet
forall (map :: * -> *) b a.
IsMap map =>
(b -> KeyOf map -> a -> b) -> b -> map a -> b
mapFoldlWithKey LabelSet -> KeyOf LabelMap -> block n C C -> LabelSet
forall a. LabelSet -> ElemOf LabelSet -> a -> LabelSet
addEntry (MaybeO x (block n C O) -> LabelSet
exitLabel MaybeO x (block n C O)
x) Body' block n
body
  where addEntry :: forall a. LabelSet -> ElemOf LabelSet -> a -> LabelSet
        addEntry :: LabelSet -> ElemOf LabelSet -> a -> LabelSet
addEntry labels :: LabelSet
labels label :: ElemOf LabelSet
label _ = ElemOf LabelSet -> LabelSet -> LabelSet
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 O) -> LabelSet
exitLabel NothingO  = LabelSet
forall set. IsSet set => set
setEmpty
        exitLabel (JustO b :: block n C O
b) = ElemOf LabelSet -> LabelSet
forall set. IsSet set => ElemOf set -> set
setSingleton (block n C O -> Label
forall (thing :: * -> * -> *) x.
NonLocal thing =>
thing C x -> Label
entryLabel block n C O
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 :: LabelMap (block C C) -> Label -> [block C C]
revPostorderFrom graph :: LabelMap (block C C)
graph start :: Label
start = DfsStack (block C C) -> LabelSet -> [block C C] -> [block C C]
go DfsStack (block C C)
start_worklist LabelSet
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 DfsStack (block C C)
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 Nil                      !LabelSet
_           ![block C C]
result = [block C C]
result
    go (ConsMark block :: block C C
block rest :: 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 block C C -> [block C C] -> [block C C]
forall a. a -> [a] -> [a]
: [block C C]
result)
    go (ConsTodo block :: block C C
block rest :: DfsStack (block C C)
rest)    !LabelSet
wip_or_done ![block C C]
result
        | block C C -> Label
forall (thing :: * -> * -> *) x.
NonLocal thing =>
thing C x -> Label
entryLabel block C C
block ElemOf LabelSet -> LabelSet -> Bool
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 =
                    (Label -> DfsStack (block C C) -> DfsStack (block C C))
-> DfsStack (block C C) -> [Label] -> DfsStack (block C C)
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
                          (block C C -> DfsStack (block C C) -> DfsStack (block C C)
forall a. a -> DfsStack a -> DfsStack a
ConsMark block C C
block DfsStack (block C C)
rest)
                          (block C C -> [Label]
forall (thing :: * -> * -> *) e.
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 (ElemOf LabelSet -> LabelSet -> LabelSet
forall set. IsSet set => ElemOf set -> set -> set
setInsert (block C C -> Label
forall (thing :: * -> * -> *) x.
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
label wl :: DfsStack (block C C)
wl
      | Just b :: block C C
b <- KeyOf LabelMap -> LabelMap (block C C) -> Maybe (block C C)
forall (map :: * -> *) a.
IsMap map =>
KeyOf map -> map a -> Maybe a
mapLookup KeyOf LabelMap
Label
label LabelMap (block C C)
graph = block C C -> DfsStack (block C C) -> DfsStack (block C C)
forall a. a -> DfsStack a -> DfsStack a
ConsTodo block C C
b DfsStack (block C C)
wl
      | Bool
otherwise =
           [Char] -> DfsStack (block C C)
forall a. HasCallStack => [Char] -> a
error ([Char] -> DfsStack (block C C)) -> [Char] -> DfsStack (block C C)
forall a b. (a -> b) -> a -> b
$ "Label that doesn't have a block?! " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Label -> [Char]
forall a. Show a => a -> [Char]
show Label
label

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