{-# LANGUAGE TupleSections, ScopedTypeVariables #-}
module Liveness (
    livenessAnalysis
) where

import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Data.Set (Set)


class    IntoSet s   where intoSet :: Ord a => s a -> Set a
instance IntoSet []  where intoSet = Set.fromList
instance IntoSet Set where intoSet = id


livenessAnalysis
    :: forall bbid bb ins var set.
       (Ord bbid, Ord var, IntoSet set)
    => [bb]  -- basic blocks
    -> (bb -> bbid)  -- ID of the basic block; should probably be Int
    -> (bb -> [ins])  -- instructions in the basic block
    -> (bb -> [bbid])  -- control flow graph: subsequents of a basic block
    -> (ins -> set var)  -- read set (variables required to be live when entering this instruction)
    -> (ins -> set var)  -- write set (variables made available after this instruction)
    -> [[(Set var, Set var)]]  -- variables live (before, after) that instruction
livenessAnalysis bblocks bidOf bbInss bbNexts fread fwrite =
    let bbEndLive = computeFlow (Map.fromList (map (,Set.empty) bids))
    in zipWith (\fs endlive -> let lives = scanr id endlive fs
                               in zip lives (tail lives))
               (mapToList insTransFs) (mapToList bbEndLive)
  where
    mapToList m = map (m Map.!) bids
    bids = map bidOf bblocks
    bidInss = Map.fromList (zip bids (map bbInss bblocks))
    bidNexts = Map.fromList (zip bids (map bbNexts bblocks))
    rwSets = Map.map (map (\ins -> (intoSet (fread ins), intoSet (fwrite ins)))) bidInss
    insTransFs = Map.map (map (\(rs, ws) live -> (live Set.\\ ws) <> rs)) rwSets
    bbTransFs = Map.map (foldr (.) id) insTransFs
    computeFlow state = let l = iterate flowStep state in fst . head . dropWhile (uncurry (/=)) $ zip l (tail l)
    flowStep state = foldl updateFlow state bids
    updateFlow state bid = Map.insert bid (Set.unions (map (\n -> (bbTransFs Map.! n) (state Map.! n)) (bidNexts Map.! bid))) state