{-# LANGUAGE ScopedTypeVariables #-}

module LifetimeAnalysis(fullLifetimeAnalysis, lifetimeAnalysis, Access(..), unAccess, BB) where

import Data.List
import Data.Maybe
import qualified Data.Map.Strict as Map
import Debug.Trace

import Utils


data Access a = Write a | Read a
  deriving (Show, Eq)

unAccess :: Access a -> a
unAccess (Write x) = x
unAccess (Read x) = x

type BB a = ([[Access a]], [Int])

data DUIO a = DUIO {defs :: [a], uses :: [a], ins :: [a], outs :: [a]}
  deriving Eq

lifetimeAnalysis :: (Eq a, Ord a) => a -> [BB a] -> [[Bool]]
lifetimeAnalysis target bbs = map (map (target `elem`)) $ fullLifetimeAnalysis bbs

fullLifetimeAnalysis :: (Eq a, Ord a) => [BB a] -> [[[a]]]
fullLifetimeAnalysis bbs =
    let duios = map fst $ eqFixpoint analysisIterator $ flip map bbs $
                    \bb@(_, nexts) -> let (d,u) = defUseAnalysis bb
                                      in (DUIO d u [] [], nexts)
    in map markLive $ zip bbs duios

markLive :: forall a. (Eq a, Ord a) => (BB a, DUIO a) -> [[a]]
markLive ((inaccs, _), duio) = init $ go fullaccs 0 (ins duio)
  where
    fullaccs = inaccs ++ [map Read (outs duio)]

    hasReadBeforeWrite _ [] = False
    hasReadBeforeWrite t (acl : rest)
        | Read t `elem` acl = True
        | Write t `elem` acl = False
        | otherwise = hasReadBeforeWrite t rest

    go :: (Eq a, Ord a) => [[Access a]] -> Int -> [a] -> [[a]]
    go [] _ _ = []
    go (acl : rest) i lives =
        let (ws, rs) = partitionAccess acl
            newlives = union rs $ flip filter (union ws lives) $ \v -> hasReadBeforeWrite v rest
        in newlives : go rest (i+1) newlives

analysisIterator :: (Eq a, Ord a) => [(DUIO a, [Int])] -> [(DUIO a, [Int])]
analysisIterator toplist = map updateIns $ map updateOuts (zip toplist [0..])
  where
    updateIns (duio, nexts) = (duio {ins = sort $ union (uses duio) (outs duio \\ defs duio)}, nexts)
    updateOuts ((duio, nexts), idx) = (duio {outs = sort $ foldl union [] (insAfter idx)}, nexts)

    insAfter = map (ins . fst . (toplist !!)) . snd . (toplist !!)

defUseAnalysis :: Eq a => BB a -> ([a], [a])
defUseAnalysis (inss, _) = foldl foldfunc ([], []) inss
  where
    foldfunc (d, u) accs =
        let (ws, rs) = partitionAccess accs
            newds = filter (not . (`elem` u)) ws
            newus = filter (not . (`elem` d)) rs
        in (union d newds, union u newus)

partitionAccess :: [Access a] -> ([a], [a])
partitionAccess [] = ([], [])
partitionAccess (Write x : rest) = let (w, r) = partitionAccess rest in (x : w, r)
partitionAccess (Read x : rest) = let (w, r) = partitionAccess rest in (w, x : r)

eqFixpoint :: Eq a => (a -> a) -> a -> a
eqFixpoint f x = let y = f x in if y == x then x else eqFixpoint f y