module LivenessAnalysis(livenessAnalysis) where

import Data.List
import qualified Data.Map.Strict as Map
import DataFlow


livenessAnalysis :: (Ord id, Ord var)
                 -- | The basic blocks, containing a list of statements and
                 -- a list of possible subsequent blocks
                 => Map.Map id ([stmt], [id])
                 -- | The blocks that can exit the program
                 -> [id]
                 -- | The gen set for a statement: the variables required
                 -- to be live right before the statement
                 -> (stmt -> [var])
                 -- | The kill set for a statement: the variables that are
                 -- assigned a new value in the statement; note that
                 -- a variable may be in both the gen and kill sets, in
                 -- which case the kill entry is ignored
                 -> (stmt -> [var])
                 -- | The list of live variables before and after each
                 -- statement; if a block contains n statements, then its
                 -- entry in the returned map contains (n+1) lists of
                 -- variables. These lists are to be placed in between the
                 -- n statements.
                 -> Map.Map id [[var]]
livenessAnalysis cfg exitNodes genF killF =
    let cfg' = Map.map snd cfg
        stmtTransF lives stmt = (lives \\ nub (killF stmt)) `union` nub (genF stmt)
        gkMap = Map.fromList
            [(i, (g, k))
            | (i, (stmts, _)) <- Map.assocs cfg
            , let (g, k) = foldl' (\(g1, k1) stmt ->
                                        let (g2, k2) = (genF stmt, killF stmt)
                                        in (k1 `union` k2, (g1 \\ k1) `union` g2))
                                  ([], [])
                                  (reverse stmts)]
        genF' = fst . (gkMap Map.!)
        killF' = snd . (gkMap Map.!)
        blockEnds = liveness cfg' exitNodes genF' killF'
    in Map.mapWithKey (\i endLives ->
                            scanl stmtTransF
                                  endLives
                                  (fst $ cfg Map.! i))
                      blockEnds

data LivenessLattice var = Lives { getLives :: [var] }

instance Ord var => Eq (LivenessLattice var) where
    Lives v1 == Lives v2 = sort v1 == sort v2

instance Ord var => Lattice (LivenessLattice var) where
    bottom = Lives []
    Lives v1 `join` Lives v2 = Lives (union v1 v2)

liveness :: (Ord id, Ord var)
         -- | The control flow graph of basic blocks
         => Map.Map id [id]
         -- | The blocks that can exit the program
         -> [id]
         -- | The gen set for a basic block
         -> (id -> [var])
         -- | The kill set for a basic block
         -> (id -> [var])
         -- | The list of live variables at the end of the blocks
         -> Map.Map id [var]
liveness cfg exitNodes genF killF =
    -- Nothing is considered to be live at the end of the program, so
    -- exitValue is bottom
    let exitValue = bottom
        genMap = Map.mapWithKey (\k _ -> nub (genF k)) cfg
        killMap = Map.mapWithKey (\k _ -> nub (killF k)) cfg
        transF (Lives vars) i =
            Lives $ (vars \\ (killMap Map.! i)) `union` (genMap Map.! i)
    in Map.map (getLives . snd) $
        dataFlowAnalysis Backwards cfg exitNodes exitValue transF