{-# LANGUAGE ScopedTypeVariables, TupleSections #-}
module DataFlow(dataFlowAnalysis, Direction(..), module Lattice) where

import Data.Function
import Data.List
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Lattice


data Direction = Forwards | Backwards
  deriving (Show, Eq)

dataFlowAnalysis
        :: forall id state. (Ord id, Lattice state)
        -- | The direction of the dataflow analysis
        => Direction
        -- | The control flow graph of basic blocks
        -> Map.Map id [id]
        -- | The initial nodes in the analysis (i.e. either the initial
        -- node or the final nodes)
        -> [id]
        -- | The initial value for sources in the analysis
        -> state
        -- | The transfer function; note that the meaning of the states
        -- depends on the direction
        -> (state -> id -> state)
        -- | For each block that is reachable from 'entryNodes', the
        -- incoming state and the outgoing state (in CFG order, not
        -- analysis order)
        -> Map.Map id (state, state)
dataFlowAnalysis Backwards cfg entryNodes initValue transF =
    Map.map (\(a, b) -> (b, a)) $
        dataFlowAnalysis Forwards (reverseGraph cfg) entryNodes initValue transF
dataFlowAnalysis Forwards cfg entryNodes initValue transF =
    let initWorklist = [(from, to) | (from, tos) <- Map.assocs cfg, to <- tos]
        reachable = floodFill entryNodes cfg
        initStates = Map.fromList
            [(lab, if lab `elem` entryNodes then initValue else bottom)
            | lab <- reachable]
    in Map.mapWithKey (\k s -> (s, transF s k)) $ flowFixpoint initWorklist initStates
  where
    flowFixpoint :: [(id, id)] -> Map.Map id state -> Map.Map id state
    flowFixpoint [] stateMap = stateMap
    flowFixpoint ((from, to) : worklist) stateMap =
        let after = transF (stateMap Map.! from) from `join` (stateMap Map.! to)
        in if after /= stateMap Map.! to
               then flowFixpoint (map (to,) (cfg Map.! to) ++ worklist)
                                 (Map.insert to after stateMap)
               else flowFixpoint worklist stateMap


reverseGraph :: Ord a => Map.Map a [a] -> Map.Map a [a]
reverseGraph graph =
    Map.fromList
    . map ((,) <$> (fst . head) <*> map snd)
    . groupBy ((==) `on` fst)
    . sortBy (compare `on` fst)
    $ [(to, from)
      | (from, tos) <- Map.assocs graph
      , to <- tos]

floodFill :: Ord a => [a] -> Map.Map a [a] -> [a]
floodFill startNodes graph = Set.toList $ go $ Set.fromList startNodes
  where
    go seen =
        let front = Set.fromList (concatMap (graph Map.!) seen) Set.\\ seen
        in if Set.size front > 0
               then go (Set.union seen front)
               else seen